[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");
----------------
amd-eochoalo wrote:

Again, my intention was for indexBitwidth to be the actual and not an upper bound. But with the assumption that this pass should work on estimates of the actual index's bitwidth, I think we would also need a lower bound.

actual index bits is 32 but the assumed upper bound is 64
i32 -> index (64) --> exact

is not sound if the lower bound is i16.

With the original semantics:

i64 -> index (64) --> exact

I'm still thinking though.


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