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

[ThomasRaoux]

Summary

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

[mravishankar]

mlir/test/Dialect/Arith/canonicalize.mlir

333

This is super confusing to read! Not an issue with this patch though. I guess its an issue with printer.... I would rather hope that printer prints in unsigned form by default.

[ThomasRaoux]

mlir/test/Dialect/Arith/canonicalize.mlir

333

Yes this is quite confusing. This is the same in llvm if I recall. Not sure there is an easy solution since it really depends on the context.

mlir/lib/Dialect/Arith/IR/ArithOps.cpp

bool arith::IndexCastUIOp::areCastCompatible(TypeRange inputs,
   return areIndexCastCompatible(inputs, outputs);
 }
 
 OpFoldResult arith::IndexCastUIOp::fold(ArrayRef<Attribute> operands) {
   // index_castui(constant) -> constant
   // 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 {};
 }
 
 void arith::IndexCastUIOp::getCanonicalizationPatterns(
     RewritePatternSet &patterns, MLIRContext *context) {
   patterns.add<IndexCastUIOfIndexCastUI, IndexCastUIOfExtUI>(context);
 }

mlir/test/Dialect/Arith/canonicalize.mlir

 //       CHECK:   %[[res:.+]] = arith.index_castui %arg0 : i8 to index
 //       CHECK:   return %[[res]]
 func.func @indexCastUIOfUnsignedExtend(%arg0: i8) -> index {
   %ext = arith.extui %arg0 : i8 to i16
   %idx = arith.index_castui %ext : i16 to 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]]
 func.func @signExtendConstant() -> i16 {
   %c-2 = arith.constant -2 : i8
   %ext = arith.extsi %c-2 : i8 to i16
   return %ext : i16
 }