[Mlir-commits] [mlir] [MLIR][Arith] Fix BitcastOp fold crashing on unhandled constant attributes (PR #212483)

[llvmlistbot at llvm.org]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir-arith

Author: Berke Ates (Berke-Ates)

<details>
<summary>Changes</summary>

`BitcastOp::fold` assumed any non-poison scalar operand attribute is a `FloatAttr` or `IntegerAttr` and hard-casted it. 
Constant attributes from other dialects, e.g. the `LLVM::UndefAttr` produced by `llvm.mlir.undef`'s fold, hit the cast assertion.
This crashed SCCP on IR where `llvm.mlir.undef` feeds `arith.bitcast`.
Bail out on attributes and result types the fold does not handle.

---
Full diff: https://github.com/llvm/llvm-project/pull/212483.diff


2 Files Affected:

- (modified) mlir/lib/Dialect/Arith/IR/ArithOps.cpp (+11-4) 
- (modified) mlir/test/Transforms/sccp.mlir (+16) 


``````````diff
diff --git a/mlir/lib/Dialect/Arith/IR/ArithOps.cpp b/mlir/lib/Dialect/Arith/IR/ArithOps.cpp
index b7fdb97aba335..14cc8718d51a5 100644
--- a/mlir/lib/Dialect/Arith/IR/ArithOps.cpp
+++ b/mlir/lib/Dialect/Arith/IR/ArithOps.cpp
@@ -2270,16 +2270,23 @@ OpFoldResult arith::BitcastOp::fold(FoldAdaptor adaptor) {
     return ub::PoisonAttr::get(getContext());
 
   /// Bitcast integer or float to integer or float.
-  APInt bits = llvm::isa<FloatAttr>(operand)
-                   ? llvm::cast<FloatAttr>(operand).getValue().bitcastToAPInt()
-                   : llvm::cast<IntegerAttr>(operand).getValue();
+  APInt bits;
+  if (auto floatAttr = dyn_cast<FloatAttr>(operand))
+    bits = floatAttr.getValue().bitcastToAPInt();
+  else if (auto intAttr = dyn_cast<IntegerAttr>(operand))
+    bits = intAttr.getValue();
+  else
+    return {};
+
   assert(resType.getIntOrFloatBitWidth() == bits.getBitWidth() &&
          "trying to fold on broken IR: operands have incompatible types");
 
   if (auto resFloatType = dyn_cast<FloatType>(resType))
     return FloatAttr::get(resType,
                           APFloat(resFloatType.getFloatSemantics(), bits));
-  return IntegerAttr::get(resType, bits);
+  if (auto resIntType = dyn_cast<IntegerType>(resType))
+    return IntegerAttr::get(resType, bits);
+  return {};
 }
 
 void arith::BitcastOp::getCanonicalizationPatterns(RewritePatternSet &patterns,
diff --git a/mlir/test/Transforms/sccp.mlir b/mlir/test/Transforms/sccp.mlir
index 251a74dc20647..b7ce60320d465 100644
--- a/mlir/test/Transforms/sccp.mlir
+++ b/mlir/test/Transforms/sccp.mlir
@@ -335,3 +335,19 @@ func.func @fold_to_non_operand_value(%x: i64, %cond: i1) -> i64 {
   %cast2 = builtin.unrealized_conversion_cast %cast1 : index to i64
   return %cast2 : i64
 }
+
+// -----
+
+// SCCP propagates the constant attribute produced by llvm.mlir.undef's fold
+// (an LLVM::UndefAttr) into arith.bitcast's fold, which must gracefully bail
+// instead of asserting.
+
+// CHECK-LABEL: func @bitcast_of_foreign_constant_attr
+func.func @bitcast_of_foreign_constant_attr() -> f64 {
+  // CHECK: %[[UNDEF:.*]] = llvm.mlir.undef : i64
+  // CHECK: %[[CAST:.*]] = arith.bitcast %[[UNDEF]] : i64 to f64
+  // CHECK: return %[[CAST]] : f64
+  %0 = llvm.mlir.undef : i64
+  %1 = arith.bitcast %0 : i64 to f64
+  return %1 : f64
+}

``````````

</details>


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