[Mlir-commits] [mlir] b9a0eb6 - [mlir][arithmetic] Add tests for IndexCast folding ops and fix assert

[Thomas Raoux]

Author: Thomas Raoux
Date: 2022-10-03T20:28:09Z
New Revision: b9a0eb610642ed67a82a4f50f07c2f7c82c7c32b

URL: https://github.com/llvm/llvm-project/commit/b9a0eb610642ed67a82a4f50f07c2f7c82c7c32b
DIFF: https://github.com/llvm/llvm-project/commit/b9a0eb610642ed67a82a4f50f07c2f7c82c7c32b.diff

LOG: [mlir][arithmetic] Add tests for IndexCast folding ops and fix assert

Fix assert in IndexCastUI folding and add tests for both IndexCastOp and
IndexCastUIOp folding

Differential Revision: https://reviews.llvm.org/D135098

Added: 
    

Modified: 
    mlir/lib/Dialect/Arith/IR/ArithOps.cpp
    mlir/test/Dialect/Arith/canonicalize.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Dialect/Arith/IR/ArithOps.cpp b/mlir/lib/Dialect/Arith/IR/ArithOps.cpp
index 190a1ef72d5bd..0378f5f9a1a74 100644
--- a/mlir/lib/Dialect/Arith/IR/ArithOps.cpp
+++ b/mlir/lib/Dialect/Arith/IR/ArithOps.cpp
@@ -1308,7 +1308,7 @@ OpFoldResult arith::IndexCastUIOp::fold(ArrayRef<Attribute> operands) {
   // A little hack because we go through int. Otherwise, the size of the
   // constant might need to change.
   if (auto value = operands[0].dyn_cast_or_null<IntegerAttr>())
-    return IntegerAttr::get(getType(), value.getUInt());
+    return IntegerAttr::get(getType(), value.getValue().getZExtValue());
 
   return {};
 }

diff  --git a/mlir/test/Dialect/Arith/canonicalize.mlir b/mlir/test/Dialect/Arith/canonicalize.mlir
index be680acea733c..337eec00f3bf9 100644
--- a/mlir/test/Dialect/Arith/canonicalize.mlir
+++ b/mlir/test/Dialect/Arith/canonicalize.mlir
@@ -317,6 +317,24 @@ func.func @indexCastUIOfUnsignedExtend(%arg0: i8) -> index {
   return %idx : index
 }
 
+// CHECK-LABEL: @indexCastFold
+//       CHECK:   %[[res:.*]] = arith.constant -2 : index
+//       CHECK:   return %[[res]]
+func.func @indexCastFold(%arg0: i8) -> index {
+  %c-2 = arith.constant -2 : i8
+  %idx = arith.index_cast %c-2 : i8 to index
+  return %idx : index
+}
+
+// CHECK-LABEL: @indexCastUIFold
+//       CHECK:   %[[res:.*]] = arith.constant 254 : index
+//       CHECK:   return %[[res]]
+func.func @indexCastUIFold(%arg0: i8) -> index {
+  %c-2 = arith.constant -2 : i8
+  %idx = arith.index_castui %c-2 : i8 to index
+  return %idx : index
+}
+
 // CHECK-LABEL: @signExtendConstant
 //       CHECK:   %[[cres:.+]] = arith.constant -2 : i16
 //       CHECK:   return %[[cres]]