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

[llvmlistbot at llvm.org]

Author: Longsheng Mou
Date: 2024-10-28T09:54:51+08:00
New Revision: 7ad63c0e44ef277591497a176991e7723165611e

URL: https://github.com/llvm/llvm-project/commit/7ad63c0e44ef277591497a176991e7723165611e
DIFF: https://github.com/llvm/llvm-project/commit/7ad63c0e44ef277591497a176991e7723165611e.diff

LOG: [mlir][MathToFuncs] `MathToFuncs` only support integer type (#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.

Added: 
    

Modified: 
    mlir/lib/Conversion/MathToFuncs/MathToFuncs.cpp
    mlir/test/Conversion/MathToFuncs/ctlz.mlir
    mlir/test/Conversion/MathToFuncs/ipowi.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Conversion/MathToFuncs/MathToFuncs.cpp b/mlir/lib/Conversion/MathToFuncs/MathToFuncs.cpp
index 3a567643ffdb8f..df5396ac628cf6 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();
@@ -798,13 +801,17 @@ bool ConvertMathToFuncsPass::isFPowIConvertible(math::FPowIOp op) {
   return (expTy && expTy.getWidth() >= minWidthOfFPowIExponent);
 }
 
+bool ConvertMathToFuncsPass::isConvertible(Operation *op) {
+  return isa<IntegerType>(getElementTypeOrSelf(op->getResult(0).getType()));
+}
+
 void ConvertMathToFuncsPass::generateOpImplementations() {
   ModuleOp module = getOperation();
 
   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 +823,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 +883,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
+}