[Mlir-commits] [mlir] [mlir][Math] Fix math-expand-ops crash on math.ctlz with index type (PR #181539)

[llvmlistbot at llvm.org]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir-math

Author: Darshan Bhat (darshan-opensource)

<details>
<summary>Changes</summary>

math.ctlz expansion called getIntOrFloatBitWidth() on the operand type without checking. Index type has no fixed bitwidth and is not int/float, so the assertion in Type::getIntOrFloatBitWidth() could fire.

- In convertCtlzOp, bail out with notifyMatchFailure when the element type is not integer or float, so expansion is only applied to types with a defined bitwidth.
- Add a test in expand-math.mlir that math.ctlz on index is left unchanged by the pass (no crash, op preserved).

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


2 Files Affected:

- (modified) mlir/lib/Dialect/Math/Transforms/ExpandOps.cpp (+5) 
- (modified) mlir/test/Dialect/Math/expand-math.mlir (+11) 


``````````diff
diff --git a/mlir/lib/Dialect/Math/Transforms/ExpandOps.cpp b/mlir/lib/Dialect/Math/Transforms/ExpandOps.cpp
index 4fc435533a1c5..f76ddfae2a67a 100644
--- a/mlir/lib/Dialect/Math/Transforms/ExpandOps.cpp
+++ b/mlir/lib/Dialect/Math/Transforms/ExpandOps.cpp
@@ -523,6 +523,11 @@ static LogicalResult convertCtlzOp(math::CountLeadingZerosOp op,
   auto eTy = getElementTypeOrSelf(operandTy);
   Location loc = op.getLoc();
 
+  // Only expand for integer or float element types (index has no fixed bitwidth).
+  if (!eTy.isIntOrFloat()) {
+    return rewriter.notifyMatchFailure(op, "ctlz expansion only supports int or float types");
+  }
+
   int32_t bitwidth = eTy.getIntOrFloatBitWidth();
   if (bitwidth > 64)
     return failure();
diff --git a/mlir/test/Dialect/Math/expand-math.mlir b/mlir/test/Dialect/Math/expand-math.mlir
index 5f2843045b885..126270ca40130 100644
--- a/mlir/test/Dialect/Math/expand-math.mlir
+++ b/mlir/test/Dialect/Math/expand-math.mlir
@@ -128,6 +128,17 @@ func.func @ctlz_vector(%arg: vector<4xi32>) -> vector<4xi32> {
 
 // -----
 
+// Index type has no fixed bitwidth; expand pass must not expand ctlz on index
+// (would assert in getIntOrFloatBitWidth). Op is left unchanged.
+// CHECK-LABEL: func @ctlz_index
+func.func @ctlz_index(%arg: index) -> index {
+  // CHECK: math.ctlz %{{.*}} : index
+  %res = math.ctlz %arg : index
+  return %res : index
+}
+
+// -----
+
 // CHECK-LABEL:    func @fmaf_func
 // CHECK-SAME:     ([[ARG0:%.+]]: f64, [[ARG1:%.+]]: f64, [[ARG2:%.+]]: f64) -> f64
 func.func @fmaf_func(%a: f64, %b: f64, %c: f64) -> f64 {

``````````

</details>


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