From ce47090d00cb0091902634ce67c90ff8ef3999f3 Mon Sep 17 00:00:00 2001
From: Slava Zakharin <szakharin@nvidia.com>
Date: Sun, 16 Apr 2023 13:07:55 -0700
Subject: [PATCH] [mlir][math] Properly disable ctlz conversion in MathToFuncs.

This fixes issues caused by D146261.
Differential Revision: https://reviews.llvm.org/D148477
---
 mlir/lib/Conversion/MathToFuncs/MathToFuncs.cpp | 5 ++++-
 mlir/test/Conversion/MathToFuncs/ctlz.mlir      | 3 +++
 2 files changed, 7 insertions(+), 1 deletion(-)

diff --git a/mlir/lib/Conversion/MathToFuncs/MathToFuncs.cpp b/mlir/lib/Conversion/MathToFuncs/MathToFuncs.cpp
index 6eac3b8..10832a1 100644
--- a/mlir/lib/Conversion/MathToFuncs/MathToFuncs.cpp
+++ b/mlir/lib/Conversion/MathToFuncs/MathToFuncs.cpp
@@ -804,6 +804,8 @@ void ConvertMathToFuncsPass::generateOpImplementations() {
   module.walk([&](Operation *op) {
     TypeSwitch<Operation *>(op)
         .Case<math::CountLeadingZerosOp>([&](math::CountLeadingZerosOp op) {
+          if (!convertCtlz)
+            return;
           Type resultType = getElementTypeOrSelf(op.getResult().getType());
 
           // Generate the software implementation of this operation,
@@ -872,7 +874,8 @@ void ConvertMathToFuncsPass::runOnOperation() {
                          vector::VectorDialect>();
 
   target.addIllegalOp<math::IPowIOp>();
-  target.addIllegalOp<math::CountLeadingZerosOp>();
+  if (convertCtlz)
+    target.addIllegalOp<math::CountLeadingZerosOp>();
   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 8678c22..4e26241 100644
--- a/mlir/test/Conversion/MathToFuncs/ctlz.mlir
+++ b/mlir/test/Conversion/MathToFuncs/ctlz.mlir
@@ -1,4 +1,5 @@
 // RUN: mlir-opt %s -split-input-file -pass-pipeline="builtin.module(convert-math-to-funcs{convert-ctlz})" | FileCheck %s
+// RUN: mlir-opt %s -split-input-file -pass-pipeline="builtin.module(convert-math-to-funcs{convert-ctlz=false})" | FileCheck --check-prefix=NOCVT %s
 
 // Check a golden-path i32 conversion
 
@@ -38,6 +39,7 @@
 // CHECK:           }
 // CHECK:           return %[[OUT]] : i32
 // CHECK:         }
+// NOCVT-NOT: __mlir_math_ctlz_i32
 func.func @main(%arg0: i32) {
   %0 = math.ctlz %arg0 : i32
   func.return
@@ -83,6 +85,7 @@ func.func @main(%arg0: i32) {
 // CHECK:           }
 // CHECK:           return %[[OUT]] : i8
 // CHECK:         }
+// NOCVT-NOT: __mlir_math_ctlz_i32
 func.func @main(%arg0: i8) {
   %0 = math.ctlz %arg0 : i8
   func.return
-- 
2.7.4