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

[Darshan Bhat]

https://github.com/darshan-opensource updated https://github.com/llvm/llvm-project/pull/181539

>From a2e58c34452d475bc5ce3283f946518777a6e5a6 Mon Sep 17 00:00:00 2001
From: Darshan Bhat <darshanbhatsirsi at gmail.com>
Date: Sun, 15 Feb 2026 10:53:23 +0100
Subject: [PATCH] [mlir][Math] Fix math-expand-ops crash on math.ctlz with
 index type

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).
---
 mlir/lib/Dialect/Math/Transforms/ExpandOps.cpp |  5 +++++
 mlir/test/Dialect/Math/expand-math.mlir        | 11 +++++++++++
 2 files changed, 16 insertions(+)

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 {