From 733fc6d0329363d7074e5ef60471fafc91c2fad8 Mon Sep 17 00:00:00 2001 From: Smit Hinsu Date: Wed, 2 Dec 2020 07:22:51 -0800 Subject: [PATCH] Fix handling of negative seeds in random number generator op kernels for XLA Casting negative s32 number to u64 directly will have leading 1s in the representation which is not what we want to get a single u64 out of two s32 seeds. Fixed this by first getting unsigned number of the same bit-width. PiperOrigin-RevId: 345227848 --- lib/Dialect/mhlo/IR/hlo_ops.cc | 12 ++++++++++-- 1 file changed, 10 insertions(+), 2 deletions(-) diff --git a/lib/Dialect/mhlo/IR/hlo_ops.cc b/lib/Dialect/mhlo/IR/hlo_ops.cc index 64b42c9..4fc3e7e 100644 --- a/lib/Dialect/mhlo/IR/hlo_ops.cc +++ b/lib/Dialect/mhlo/IR/hlo_ops.cc @@ -518,11 +518,19 @@ void ConvertOp::build(OpBuilder& builder, OperationState& result, Value operand, } OpFoldResult ConvertOp::fold(ArrayRef operands) { - if (getOperand().getType() == getResult().getType()) return getOperand(); + auto operand_ty = getOperand().getType().cast(); + auto result_ty = getResult().getType().cast(); + if (operand_ty == result_ty) return getOperand(); // If the result has non-static shape, a convert op is necessary to go from // static shape to non-static shape. - if (!getResult().getType().cast().hasStaticShape()) return {}; + if (!result_ty.hasStaticShape()) return {}; + + // TODO(hinsu): Handle unsigned types. + if (operand_ty.getElementType().isUnsignedInteger() || + result_ty.getElementType().isUnsignedInteger()) { + return {}; + } // If the operand is constant, we can do the conversion now. if (auto elementsAttr = operands.front().dyn_cast_or_null()) {