[Mlir-commits] [mlir] [mlir][arith] Add InferExactOnIndexCast to cast folders (PR #184631)

[Erick Ochoa Lopez]

================
@@ -1844,9 +1844,42 @@ bool arith::IndexCastOp::areCastCompatible(TypeRange inputs,
   return areIndexCastCompatible(inputs, outputs);
 }
 
+static unsigned getBitwidth(Type type, unsigned indexBitwidth) {
+  Type elemType = getElementTypeOrSelf(type);
+  if (isa<IndexType>(elemType))
+    return indexBitwidth;
+  return elemType.getIntOrFloatBitWidth();
+}
+
+template <typename CastOp>
+struct InferExactOnIndexCast final : OpRewritePattern<CastOp> {
+  InferExactOnIndexCast(MLIRContext *context)
+      : OpRewritePattern<CastOp>(context) {}
+
+  LogicalResult matchAndRewrite(CastOp op,
+                                PatternRewriter &rewriter) const override {
+    if (op.getExact())
+      return failure();
+
+    DataLayout layout = DataLayout::closest(op);
+    unsigned indexBitwidth =
+        layout.getTypeSizeInBits(IndexType::get(op.getContext()));
+    unsigned srcBW = getBitwidth(op.getIn().getType(), indexBitwidth);
+    unsigned dstBW = getBitwidth(op.getType(), indexBitwidth);
+    if (srcBW > dstBW)
+      return rewriter.notifyMatchFailure(op, "source is wider than dest");
+
+    rewriter.modifyOpInPlace(op, [&] { op.setExact(true); });
+    return success();
+  }
+};
+
 OpFoldResult arith::IndexCastOp::fold(FoldAdaptor adaptor) {
   // index_cast(constant) -> constant
-  unsigned resultBitwidth = 64; // Default for index integer attributes.
+  DataLayout layout = DataLayout::closest(*this);
----------------
amd-eochoalo wrote:

>  I think what we are relying on here that cast is exact when it's a noop or an extension, and we assume that the concrete index type will be <= the default

I think this folder (ignoring the InferExactOnIndexCast pattern) should be changed to instead look for `exact` and if exact is set then it can fold correctly. (Which brings us back to InferExactOnIndexCast ...)

https://github.com/llvm/llvm-project/pull/184631