[mlir][math] Properly disable ctlz conversion in MathToFuncs.

[vzakhari]

Summary

This fixes issues caused by D146261.

[mehdi_amini]

Thanks!

mlir/test/Conversion/MathToFuncs/ctlz.mlir

2

We could have added the test to the other file that already run without the option enabled.

[vzakhari]

mlir/test/Conversion/MathToFuncs/ctlz.mlir

2

Thanks for the review. I do not think there is another test for math.ctlz that enables convert-math-to-funcs converter explicitly. @j2kun, can you please follow up on this comment?

I will merge it as-is shortly.

[PeteSteinfeld]

Thanks, Slava! This passes all of my tests.

[j2kun]

What exactly is the followup here? I see the extra test added that ensures ctlz is not lowered when the flag is disabled. Is another test helpful?

mlir/lib/Conversion/MathToFuncs/MathToFuncs.cpp

 }
 
 void ConvertMathToFuncsPass::generateOpImplementations() {
   ModuleOp module = getOperation();
 
   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,
           // if it has not been generated yet.
           auto key = std::pair(op->getName(), resultType);
           auto entry = funcImpls.try_emplace(key, func::FuncOp{});
           if (entry.second)
             entry.first->second = createCtlzFunc(&module, resultType);
if (convertCtlz)
     patterns.add<CtlzOpLowering>(patterns.getContext(), getFuncOpByType);
 
   ConversionTarget target(getContext());
   target.addLegalDialect<arith::ArithDialect, cf::ControlFlowDialect,
                          func::FuncDialect, scf::SCFDialect,
                          vector::VectorDialect>();
 
   target.addIllegalOp<math::IPowIOp>();
+  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))))
     signalPassFailure();
 }

mlir/test/Conversion/MathToFuncs/ctlz.mlir

 // 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
 
 // CHECK-LABEL:   func.func @main(
 // CHECK-SAME:                       %[[VAL_0:.*]]: i32
 // CHECK-SAME:                       ) {
 // CHECK:           %[[VAL_1:.*]] = call @__mlir_math_ctlz_i32(%[[VAL_0]]) : (i32) -> i32
 // CHECK:           return

 // CHECK:                 scf.yield %[[ARG_NEXT]], %[[N_NEXT]] : i32, i32
 // CHECK:               }
 // CHECK:               scf.yield %[[IF_RET]]#0, %[[IF_RET]]#1 : i32, i32
 // CHECK:             }
 // CHECK:             scf.yield %[[FOR_RET]]#1 : i32
 // CHECK:           }
 // CHECK:           return %[[OUT]] : i32
 // CHECK:         }
+// NOCVT-NOT: __mlir_math_ctlz_i32
 func.func @main(%arg0: i32) {
   %0 = math.ctlz %arg0 : i32
   func.return
 }
 
 // -----
 
 // Check that i8 input is preserved

 // CHECK:                 scf.yield %[[ARG_NEXT]], %[[N_NEXT]] : i8, i8
 // CHECK:               }
 // CHECK:               scf.yield %[[IF_RET]]#0, %[[IF_RET]]#1 : i8, i8
 // CHECK:             }
 // CHECK:             scf.yield %[[FOR_RET]]#1 : i8
 // CHECK:           }
 // CHECK:           return %[[OUT]] : i8
 // CHECK:         }
+// NOCVT-NOT: __mlir_math_ctlz_i32
 func.func @main(%arg0: i8) {
   %0 = math.ctlz %arg0 : i8
   func.return
 }