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

[llvmlistbot at llvm.org]

Author: Darshan Bhat
Date: 2026-02-15T23:56:44Z
New Revision: 901e543c15bd5e46a33dc19354a35da0278aaf7b

URL: https://github.com/llvm/llvm-project/commit/901e543c15bd5e46a33dc19354a35da0278aaf7b
DIFF: https://github.com/llvm/llvm-project/commit/901e543c15bd5e46a33dc19354a35da0278aaf7b.diff

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

Fixes #179847
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).

Added: 
    

Modified: 
    mlir/lib/Dialect/Math/Transforms/ExpandOps.cpp
    mlir/test/Dialect/Math/expand-math.mlir

Removed: 
    


################################################################################
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 {