[Mlir-commits] [mlir] [mlir][MathToFuncs] `MathToFuncs` only support integer type (PR #113693)

[Longsheng Mou]

https://github.com/CoTinker created https://github.com/llvm/llvm-project/pull/113693

This PR fixes a bug in `MathToFuncs` where it incorrectly converts index type for `math.ctlz` and `math.ipowi`, leading to a crash. Fixes #108150.

>From 05ef46070d5bb13499a05b6d50132d66263c1902 Mon Sep 17 00:00:00 2001
From: Longsheng Mou <longshengmou at gmail.com>
Date: Fri, 25 Oct 2024 18:23:04 +0800
Subject: [PATCH] [mlir][MathToFuncs] `MathToFuncs` only support integer type

This PR fixes a bug in `MathToFuncs` where it incorrectly converts
index type for `math.ctlz` and `math.ipowi`, leading to a crash.
---
 mlir/lib/Conversion/MathToFuncs/MathToFuncs.cpp | 17 +++++++++++++----
 mlir/test/Conversion/MathToFuncs/ctlz.mlir      | 10 ++++++++++
 mlir/test/Conversion/MathToFuncs/ipowi.mlir     | 11 +++++++++++
 3 files changed, 34 insertions(+), 4 deletions(-)

diff --git a/mlir/lib/Conversion/MathToFuncs/MathToFuncs.cpp b/mlir/lib/Conversion/MathToFuncs/MathToFuncs.cpp
index 3a567643ffdb8f..f8af53355ea0fc 100644
--- a/mlir/lib/Conversion/MathToFuncs/MathToFuncs.cpp
+++ b/mlir/lib/Conversion/MathToFuncs/MathToFuncs.cpp
@@ -781,6 +781,9 @@ struct ConvertMathToFuncsPass
   // or equal to minWidthOfFPowIExponent option value.
   bool isFPowIConvertible(math::FPowIOp op);
 
+  // Reture true, if operation is integer type.
+  bool isConvertible(Operation *op);
+
   // Generate outlined implementations for power operations
   // and store them in funcImpls map.
   void generateOpImplementations();
@@ -804,7 +807,7 @@ void ConvertMathToFuncsPass::generateOpImplementations() {
   module.walk([&](Operation *op) {
     TypeSwitch<Operation *>(op)
         .Case<math::CountLeadingZerosOp>([&](math::CountLeadingZerosOp op) {
-          if (!convertCtlz)
+          if (!convertCtlz || !isConvertible(op))
             return;
           Type resultType = getElementTypeOrSelf(op.getResult().getType());
 
@@ -816,6 +819,9 @@ void ConvertMathToFuncsPass::generateOpImplementations() {
             entry.first->second = createCtlzFunc(&module, resultType);
         })
         .Case<math::IPowIOp>([&](math::IPowIOp op) {
+          if (!isConvertible(op))
+            return;
+
           Type resultType = getElementTypeOrSelf(op.getResult().getType());
 
           // Generate the software implementation of this operation,
@@ -873,9 +879,12 @@ void ConvertMathToFuncsPass::runOnOperation() {
                          func::FuncDialect, scf::SCFDialect,
                          vector::VectorDialect>();
 
-  target.addIllegalOp<math::IPowIOp>();
-  if (convertCtlz)
-    target.addIllegalOp<math::CountLeadingZerosOp>();
+  target.addDynamicallyLegalOp<math::IPowIOp>(
+      [this](math::IPowIOp op) { return !isConvertible(op); });
+  if (convertCtlz) {
+    target.addDynamicallyLegalOp<math::CountLeadingZerosOp>(
+        [this](math::CountLeadingZerosOp op) { return !isConvertible(op); });
+  }
   target.addDynamicallyLegalOp<math::FPowIOp>(
       [this](math::FPowIOp op) { return !isFPowIConvertible(op); });
   if (failed(applyPartialConversion(module, target, std::move(patterns))))
diff --git a/mlir/test/Conversion/MathToFuncs/ctlz.mlir b/mlir/test/Conversion/MathToFuncs/ctlz.mlir
index 4e262417d6a959..b7ef0a8928912d 100644
--- a/mlir/test/Conversion/MathToFuncs/ctlz.mlir
+++ b/mlir/test/Conversion/MathToFuncs/ctlz.mlir
@@ -91,3 +91,13 @@ func.func @main(%arg0: i8) {
   func.return
 }
 
+// -----
+
+// Check that index is not converted
+
+// CHECK-LABEL: func.func @ctlz_index
+// CHECK:         math.ctlz
+func.func @ctlz_index(%arg0: index) {
+  %0 = math.ctlz %arg0 : index
+  func.return
+}
diff --git a/mlir/test/Conversion/MathToFuncs/ipowi.mlir b/mlir/test/Conversion/MathToFuncs/ipowi.mlir
index e464e9ca9564fc..2702a1e22e621d 100644
--- a/mlir/test/Conversion/MathToFuncs/ipowi.mlir
+++ b/mlir/test/Conversion/MathToFuncs/ipowi.mlir
@@ -170,3 +170,14 @@ func.func @ipowi_vec(%arg0: vector<2x3xi64>, %arg1: vector<2x3xi64>) {
   %0 = math.ipowi %arg0, %arg1 : vector<2x3xi64>
   func.return
 }
+
+// -----
+
+// Check that index is not converted
+
+// CHECK-LABEL: func.func @ipowi_index
+// CHECK:         math.ipowi
+func.func @ipowi_index(%arg0: index, %arg1: index) {
+  %0 = math.ipowi %arg0, %arg1 : index
+  func.return
+}