[mlir][math] Add math.ctlz expansion to control flow + arith operations

[rsuderman]

Summary

Ctlz is an intrinsic in LLVM but does not have equivalent operations in SPIR-V.
Including a decomposition gives an alternative path for these platforms.

[rsuderman]

Added missing newline

[NatashaKnk]

mlir/lib/Dialect/Math/Transforms/ExpandPatterns.cpp

87

Can the input not be negative? The variable is called largerThanZero but the comparison is about equality -- would it be better to use a gt predicates?

[rsuderman]

Changed var name to be more correct.

[rsuderman]

Woops, slight typo

[rsuderman]

Update for natashaknk@'s comments.

mlir/include/mlir/Dialect/Math/Transforms/Passes.h

 //===- Passes.h - Pass Entrypoints ------------------------------*- C++ -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_DIALECT_MATH_TRANSFORMS_PASSES_H_
 #define MLIR_DIALECT_MATH_TRANSFORMS_PASSES_H_
 
 namespace mlir {
 
 class RewritePatternSet;
 
+void populateExpandCtlzPattern(RewritePatternSet &patterns);
 void populateExpandTanhPattern(RewritePatternSet &patterns);
 
 void populateMathAlgebraicSimplificationPatterns(RewritePatternSet &patterns);
 
 struct MathPolynomialApproximationOptions {
   // Enables the use of AVX2 intrinsics in some of the approximations.
   bool enableAvx2 = false;
 };
 
 void populateMathPolynomialApproximationPatterns(
     RewritePatternSet &patterns,
     const MathPolynomialApproximationOptions &options = {});
 
 } // namespace mlir
 
 #endif // MLIR_DIALECT_MATH_TRANSFORMS_PASSES_H_

mlir/lib/Dialect/Math/Transforms/CMakeLists.txt

 add_mlir_dialect_library(MLIRMathTransforms
   AlgebraicSimplification.cpp
-  ExpandTanh.cpp
+  ExpandPatterns.cpp
   PolynomialApproximation.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/Math/Transforms
 
   LINK_LIBS PUBLIC
   MLIRArithmetic
   MLIRDialectUtils  
   MLIRIR
   MLIRMath
   MLIRPass
   MLIRTransforms
   MLIRX86Vector
   MLIRVector
   MLIRVectorUtils
   )

mlir/lib/Dialect/Math/Transforms/ExpandPatterns.cpp

 //===- ExpandTanh.cpp - Code to perform expanding tanh op -----------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // This file implements expansion of tanh op.
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/Arithmetic/IR/Arithmetic.h"
 #include "mlir/Dialect/Math/IR/Math.h"
 #include "mlir/Dialect/Math/Transforms/Passes.h"
+#include "mlir/Dialect/SCF/SCF.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/Transforms/DialectConversion.h"
 
 using namespace mlir;
 
 /// Expands tanh op into
 ///   1) 1-exp^{-2x} / 1+exp^{-2x}, if x => 0
 ///   2) exp^{2x}-1 / exp^{2x}+1  , if x < 0
static LogicalResult convertTanhOp(math::TanhOp op, PatternRewriter &rewriter) {
   Value zero = rewriter.create<arith::ConstantOp>(loc, floatZero);
   Value cmpRes = rewriter.create<arith::CmpFOp>(loc, arith::CmpFPredicate::OGE,
                                                 op.getOperand(), zero);
   rewriter.replaceOpWithNewOp<arith::SelectOp>(op, cmpRes, positiveRes,
                                                negativeRes);
   return success();
 }
 
+static LogicalResult convertCtlzOp(math::CountLeadingZerosOp op,
+                                   PatternRewriter &rewriter) {
+  auto operand = op.getOperand();
+  auto elementTy = operand.getType();
+  auto resultTy = op.getType();
+  Location loc = op.getLoc();
+
+  int bitWidth = elementTy.getIntOrFloatBitWidth();
+  auto zero =
+      rewriter.create<arith::ConstantOp>(loc, IntegerAttr::get(elementTy, 0));
+  auto leadingZeros = rewriter.create<arith::ConstantOp>(
+      loc, IntegerAttr::get(elementTy, bitWidth));
+
+  SmallVector<Value> operands = {operand, leadingZeros, zero};
+  SmallVector<Type> types = {elementTy, elementTy, elementTy};
+  SmallVector<Location> locations = {loc, loc, loc};
+
+  auto whileOp = rewriter.create<scf::WhileOp>(loc, types, operands);
+  Block *before =
+      rewriter.createBlock(&whileOp.getBefore(), {}, types, locations);
+  Block *after =
+      rewriter.createBlock(&whileOp.getAfter(), {}, types, locations);
+
+  // The conditional block of the while loop.
+  {
+    rewriter.setInsertionPointToStart(&whileOp.getBefore().front());
+    Value input = before->getArgument(0);
+    Value zero = before->getArgument(2);
+
+    Value inputNotZero = rewriter.create<arith::CmpIOp>(
+        loc, arith::CmpIPredicate::ne, input, zero);
+    rewriter.create<scf::ConditionOp>(loc, inputNotZero,
+                                      before->getArguments());
+  }
+
+  // The body of the while loop: shift right until reaching a value of 0.
+  {
+    rewriter.setInsertionPointToStart(&whileOp.getAfter().front());
+    Value input = after->getArgument(0);
+    Value leadingZeros = after->getArgument(1);
+
+    auto one =
+        rewriter.create<arith::ConstantOp>(loc, IntegerAttr::get(elementTy, 1));
+    auto shifted = rewriter.create<arith::ShRUIOp>(loc, resultTy, input, one);
+    auto leadingZerosMinusOne =
+        rewriter.create<arith::SubIOp>(loc, resultTy, leadingZeros, one);
+
+    rewriter.create<scf::YieldOp>(
+        loc,
+        ValueRange({shifted, leadingZerosMinusOne, after->getArgument(2)}));
+  }
+
+  rewriter.setInsertionPointAfter(whileOp);
+  rewriter.replaceOp(op, whileOp->getResult(1));
+  return success();
+}
+
+void mlir::populateExpandCtlzPattern(RewritePatternSet &patterns) {
+  patterns.add(convertCtlzOp);
+}
+
 void mlir::populateExpandTanhPattern(RewritePatternSet &patterns) {
   patterns.add(convertTanhOp);
 }

mlir/test/Dialect/Math/expand-math.mlir

-// RUN: mlir-opt %s -test-expand-tanh | FileCheck %s
+// RUN: mlir-opt %s --split-input-file -test-expand-math | FileCheck %s
 
 // CHECK-LABEL: func @tanh
 func.func @tanh(%arg: f32) -> f32 {
   %res = math.tanh %arg : f32
   return %res : f32
 }
 // CHECK-DAG: %[[ZERO:.+]] = arith.constant 0.000000e+00 : f32
 // CHECK-DAG: %[[ONE:.+]] = arith.constant 1.000000e+00 : f32
 // CHECK-DAG: %[[TWO:.+]] = arith.constant 2.000000e+00 : f32
 // CHECK: %[[DOUBLEDX:.+]] = arith.mulf %arg0, %[[TWO]] : f32
 // CHECK: %[[NEGDOUBLEDX:.+]] = arith.negf %[[DOUBLEDX]] : f32
 // CHECK: %[[EXP1:.+]] = math.exp %[[NEGDOUBLEDX]] : f32
 // CHECK: %[[DIVIDEND1:.+]] = arith.subf %[[ONE]], %[[EXP1]] : f32
 // CHECK: %[[DIVISOR1:.+]] = arith.addf %[[EXP1]], %[[ONE]] : f32
 // CHECK: %[[RES1:.+]] = arith.divf %[[DIVIDEND1]], %[[DIVISOR1]] : f32
 // CHECK: %[[EXP2:.+]] = math.exp %[[DOUBLEDX]] : f32
 // CHECK: %[[DIVIDEND2:.+]] = arith.subf %[[EXP2]], %[[ONE]] : f32
 // CHECK: %[[DIVISOR2:.+]] = arith.addf %[[EXP2]], %[[ONE]] : f32
 // CHECK: %[[RES2:.+]] = arith.divf %[[DIVIDEND2]], %[[DIVISOR2]] : f32
 // CHECK: %[[COND:.+]] = arith.cmpf oge, %arg0, %[[ZERO]] : f32
 // CHECK: %[[RESULT:.+]] = arith.select %[[COND]], %[[RES1]], %[[RES2]] : f32
 // CHECK: return %[[RESULT]]
+
+// ----
+
+// CHECK-LABEL: func @ctlz
+func.func @ctlz(%arg: i32) -> i32 {
+  // CHECK: %[[C0:.+]] = arith.constant 0 : i32
+  // CHECK: %[[C32:.+]] = arith.constant 32 : i32
+  // CHECK: %[[C1:.+]] = arith.constant 1 : i32
+  // CHECK: %[[WHILE:.+]]:3 = scf.while (%[[A1:.+]] = %arg0, %[[A2:.+]] = %[[C32]], %[[A3:.+]] = %[[C0]])
+  // CHECK:   %[[CMP:.+]] = arith.cmpi ne, %[[A1]], %[[A3]]
+  // CHECK:   scf.condition(%[[CMP]]) %[[A1]], %[[A2]], %[[A3]]
+  // CHECK:   %[[SHR:.+]] = arith.shrui %[[A1]], %[[C1]]
+  // CHECK:   %[[SUB:.+]] = arith.subi %[[A2]], %[[C1]]
+  // CHECK:   scf.yield %[[SHR]], %[[SUB]], %[[A3]]
+  %res = math.ctlz %arg : i32
+
+  // CHECK: return %[[WHILE]]#1
+  return %res : i32
+}

mlir/test/lib/Dialect/Math/CMakeLists.txt

 # Exclude tests from libMLIR.so
 add_mlir_library(MLIRMathTestPasses
   TestAlgebraicSimplification.cpp
-  TestExpandTanh.cpp
+  TestExpandMath.cpp
   TestPolynomialApproximation.cpp
 
   EXCLUDE_FROM_LIBMLIR
 
   LINK_LIBS PUBLIC
   MLIRMathTransforms
   MLIRPass
   MLIRTransformUtils
   MLIRVector
   MLIRX86Vector
   )

mlir/test/lib/Dialect/Math/TestExpandMath.cpp

-//===- TestExpandTanh.cpp - Test expand tanh op into exp form -------------===//
+//===- TestExpandMath.cpp - Test expand math op into exp form -------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
-// This file contains test passes for expanding tanh.
+// This file contains test passes for expanding math operations.
 //
 //===----------------------------------------------------------------------===//
 
+#include "mlir/Dialect/Arithmetic/IR/Arithmetic.h"
 #include "mlir/Dialect/Math/Transforms/Passes.h"
+#include "mlir/Dialect/SCF/SCF.h"
 #include "mlir/Pass/Pass.h"
 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
 
 using namespace mlir;
 
 namespace {
-struct TestExpandTanhPass
-    : public PassWrapper<TestExpandTanhPass, OperationPass<>> {
-  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestExpandTanhPass)
+struct TestExpandMathPass
+    : public PassWrapper<TestExpandMathPass, OperationPass<>> {
+  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestExpandMathPass)
 
   void runOnOperation() override;
-  StringRef getArgument() const final { return "test-expand-tanh"; }
-  StringRef getDescription() const final { return "Test expanding tanh"; }
+  StringRef getArgument() const final { return "test-expand-math"; }
+  void getDependentDialects(DialectRegistry &registry) const override {
+    registry.insert<arith::ArithmeticDialect, scf::SCFDialect>();
+  }
+  StringRef getDescription() const final { return "Test expanding math"; }
 };
 } // namespace
 
-void TestExpandTanhPass::runOnOperation() {
+void TestExpandMathPass::runOnOperation() {
   RewritePatternSet patterns(&getContext());
+  populateExpandCtlzPattern(patterns);
   populateExpandTanhPattern(patterns);
   (void)applyPatternsAndFoldGreedily(getOperation(), std::move(patterns));
 }
 
 namespace mlir {
 namespace test {
-void registerTestExpandTanhPass() { PassRegistration<TestExpandTanhPass>(); }
+void registerTestExpandMathPass() { PassRegistration<TestExpandMathPass>(); }
 } // namespace test
 } // namespace mlir

mlir/tools/mlir-opt/mlir-opt.cpp

 void registerTestGpuSerializeToCubinPass();
 void registerTestGpuSerializeToHsacoPass();
 void registerTestDataFlowPass();
 void registerTestDataLayoutQuery();
 void registerTestDecomposeCallGraphTypes();
 void registerTestDiagnosticsPass();
 void registerTestDominancePass();
 void registerTestDynamicPipelinePass();
-void registerTestExpandTanhPass();
+void registerTestExpandMathPass();
 void registerTestComposeSubView();
 void registerTestMultiBuffering();
 void registerTestIRVisitorsPass();
 void registerTestGenericIRVisitorsPass();
 void registerTestGenericIRVisitorsInterruptPass();
 void registerTestInterfaces();
 void registerTestLinalgCodegenStrategy();
 void registerTestLinalgElementwiseFusion();

 #if MLIR_ROCM_CONVERSIONS_ENABLED
   mlir::test::registerTestGpuSerializeToHsacoPass();
 #endif
   mlir::test::registerTestDecomposeCallGraphTypes();
   mlir::test::registerTestDataFlowPass();
   mlir::test::registerTestDataLayoutQuery();
   mlir::test::registerTestDominancePass();
   mlir::test::registerTestDynamicPipelinePass();
-  mlir::test::registerTestExpandTanhPass();
+  mlir::test::registerTestExpandMathPass();
   mlir::test::registerTestComposeSubView();
   mlir::test::registerTestMultiBuffering();
   mlir::test::registerTestIRVisitorsPass();
   mlir::test::registerTestGenericIRVisitorsPass();
   mlir::test::registerTestInterfaces();
   mlir::test::registerTestLinalgCodegenStrategy();
   mlir::test::registerTestLinalgElementwiseFusion();
   mlir::test::registerTestLinalgFusionTransforms();