[mlir] Added ctlz and cttz to math dialect and LLVM dialect

[rsuderman]

Summary

Count leading/trailing zeros are an existing LLVM intrinsic. Added LLVM
support for the intrinsics with lowerings from the math dialect to LLVM
dialect.

[ftynse]

mlir/include/mlir/Dialect/LLVMIR/LLVMOpBase.td

97 ↗

(On Diff #392233)

The comment does not match the code.

98–104 ↗

(On Diff #392233)

Why do we need this? There is no "LLVM integer type", the dialect uses built-in integers. Why can't this match just use I1?

mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td

1407

Please give the name to the second argument.

mlir/include/mlir/Dialect/Math/IR/MathOps.td

304

Bikeshed: can we consider more pronounceable names at this level?

[rsuderman]

Updated names and addressed code review comments.

[rsuderman]

Removed extra line.

[rsuderman]

mlir/include/mlir/Dialect/LLVMIR/LLVMOpBase.td

97 ↗

(On Diff #392233)

Removed.

98–104 ↗

(On Diff #392233)

I wanted to match the existing type naming behavior. I changed to directly using I<1> after.

mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td

 }
 
 class LLVM_TernarySameArgsIntrinsicOp<string func, list<OpTrait> traits = []> :
     LLVM_OneResultIntrOp<func, [], [0],
            !listconcat([NoSideEffect, SameOperandsAndResultType], traits)> {
   let arguments = (ins LLVM_Type:$a, LLVM_Type:$b, LLVM_Type:$c);
 }
 
+class LLVM_CountZerosIntrinsicOp<string func, list<OpTrait> traits = []> :
+    LLVM_OneResultIntrOp<func, [], [0],
+           !listconcat([NoSideEffect], traits)> {
+  let arguments = (ins LLVM_Type:$in, I<1>:$zero_undefined);
+}
+
 def LLVM_CopySignOp : LLVM_BinarySameArgsIntrinsicOp<"copysign">;
 def LLVM_CosOp : LLVM_UnaryIntrinsicOp<"cos">;
 def LLVM_ExpOp : LLVM_UnaryIntrinsicOp<"exp">;
 def LLVM_Exp2Op : LLVM_UnaryIntrinsicOp<"exp2">;
 def LLVM_FAbsOp : LLVM_UnaryIntrinsicOp<"fabs">;
 def LLVM_FCeilOp : LLVM_UnaryIntrinsicOp<"ceil">;
 def LLVM_FFloorOp : LLVM_UnaryIntrinsicOp<"floor">;
 def LLVM_FMAOp : LLVM_TernarySameArgsIntrinsicOp<"fma">;
 def LLVM_FMulAddOp : LLVM_TernarySameArgsIntrinsicOp<"fmuladd">;
 def LLVM_Log10Op : LLVM_UnaryIntrinsicOp<"log10">;
 def LLVM_Log2Op : LLVM_UnaryIntrinsicOp<"log2">;
 def LLVM_LogOp : LLVM_UnaryIntrinsicOp<"log">;
 def LLVM_Prefetch : LLVM_ZeroResultIntrOp<"prefetch", [0]> {
   let arguments = (ins LLVM_Type:$addr, LLVM_Type:$rw, LLVM_Type:$hint,
                    LLVM_Type:$cache);
 }
 def LLVM_SinOp : LLVM_UnaryIntrinsicOp<"sin">;
 def LLVM_SqrtOp : LLVM_UnaryIntrinsicOp<"sqrt">;
 def LLVM_PowOp : LLVM_BinarySameArgsIntrinsicOp<"pow">;
 def LLVM_BitReverseOp : LLVM_UnaryIntrinsicOp<"bitreverse">;
+def LLVM_CountLeadingZerosOp : LLVM_CountZerosIntrinsicOp<"ctlz">;
+def LLVM_CountTrailingZerosOp : LLVM_CountZerosIntrinsicOp<"cttz">;
 def LLVM_CtPopOp : LLVM_UnaryIntrinsicOp<"ctpop">;
 def LLVM_MaxNumOp : LLVM_BinarySameArgsIntrinsicOp<"maxnum">;
 def LLVM_MinNumOp : LLVM_BinarySameArgsIntrinsicOp<"minnum">;
 def LLVM_MaximumOp : LLVM_BinarySameArgsIntrinsicOp<"maximum">;
 def LLVM_MinimumOp : LLVM_BinarySameArgsIntrinsicOp<"minimum">;
 def LLVM_SMaxOp : LLVM_BinarySameArgsIntrinsicOp<"smax">;
 def LLVM_SMinOp : LLVM_BinarySameArgsIntrinsicOp<"smin">;
 

mlir/include/mlir/Dialect/Math/IR/MathOps.td

let description = [{
 
     // Tensor element-wise sine value.
     %x = math.sin %y : tensor<4x?xf8>
     ```
   }];
 }
 
 //===----------------------------------------------------------------------===//
+// CountLeadingZerosOp
+//===----------------------------------------------------------------------===//
+
+def Math_CountLeadingZerosOp : Math_IntegerUnaryOp<"ctlz"> {
+  let summary = "counts the leading zeros an integer value";
+  let description = [{
+    The `ctlz` operation computes the number of leading zeros of an integer value.
+
+    Example:
+
+    ```mlir
+    // Scalar ctlz function value.
+    %a = math.ctlz %b : i32
+
+    // SIMD vector element-wise ctlz function value.
+    %f = math.ctlz %g : vector<4xi16>
+
+    // Tensor element-wise ctlz function value.
+    %x = math.ctlz %y : tensor<4x?xi8>
+    ```
+  }];
+}
+
+//===----------------------------------------------------------------------===//
+// CountTrailingZerosOp
+//===----------------------------------------------------------------------===//
+
+def Math_CountTrailingZerosOp : Math_IntegerUnaryOp<"cttz"> {
+  let summary = "counts the trailing zeros an integer value";
+  let description = [{
+    The `cttz` operation computes the number of trailing zeros of an integer value.
+
+    Example:
+
+    ```mlir
+    // Scalar cttz function value.
+    %a = math.cttz %b : i32
+
+    // SIMD vector element-wise cttz function value.
+    %f = math.cttz %g : vector<4xi16>
+
+    // Tensor element-wise cttz function value.
+    %x = math.cttz %y : tensor<4x?xi8>
+    ```
+  }];
+}
+
+//===----------------------------------------------------------------------===//
 // CtPopOp
 //===----------------------------------------------------------------------===//
 
 def Math_CtPopOp : Math_IntegerUnaryOp<"ctpop"> {
   let summary = "counts the number of set bits of an integer value";
   let description = [{
     The `ctpop` operation computes the number of set bits of an integer value.
 

mlir/lib/Conversion/MathToLLVM/MathToLLVM.cpp

 using Log10OpLowering =
     VectorConvertToLLVMPattern<math::Log10Op, LLVM::Log10Op>;
 using Log2OpLowering = VectorConvertToLLVMPattern<math::Log2Op, LLVM::Log2Op>;
 using LogOpLowering = VectorConvertToLLVMPattern<math::LogOp, LLVM::LogOp>;
 using PowFOpLowering = VectorConvertToLLVMPattern<math::PowFOp, LLVM::PowOp>;
 using SinOpLowering = VectorConvertToLLVMPattern<math::SinOp, LLVM::SinOp>;
 using SqrtOpLowering = VectorConvertToLLVMPattern<math::SqrtOp, LLVM::SqrtOp>;
 
+// A `CtLz/CtTz(a)` is converted into `CtLz/CtTz(a, false)`.
+template <typename MathOp, typename LLVMOp>
+struct CountOpLowering : public ConvertOpToLLVMPattern<MathOp> {
+  using ConvertOpToLLVMPattern<MathOp>::ConvertOpToLLVMPattern;
+  using Super = CountOpLowering<MathOp, LLVMOp>;
+
+  LogicalResult
+  matchAndRewrite(MathOp op, typename MathOp::Adaptor adaptor,
+                  ConversionPatternRewriter &rewriter) const override {
+    auto operandType = adaptor.getOperand().getType();
+
+    if (!operandType || !LLVM::isCompatibleType(operandType))
+      return failure();
+
+    auto loc = op.getLoc();
+    auto resultType = op.getResult().getType();
+    auto boolType = rewriter.getIntegerType(1);
+    auto boolZero = rewriter.getIntegerAttr(boolType, 0);
+
+    if (!operandType.template isa<LLVM::LLVMArrayType>()) {
+      LLVM::ConstantOp zero =
+          rewriter.create<LLVM::ConstantOp>(loc, boolType, boolZero);
+      rewriter.replaceOpWithNewOp<LLVMOp>(op, resultType, adaptor.getOperand(),
+                                          zero);
+      return success();
+    }
+
+    auto vectorType = resultType.template dyn_cast<VectorType>();
+    if (!vectorType)
+      return failure();
+
+    return LLVM::detail::handleMultidimensionalVectors(
+        op.getOperation(), adaptor.getOperands(), *this->getTypeConverter(),
+        [&](Type llvm1DVectorTy, ValueRange operands) {
+          LLVM::ConstantOp zero =
+              rewriter.create<LLVM::ConstantOp>(loc, boolType, boolZero);
+          return rewriter.replaceOpWithNewOp<LLVMOp>(op, llvm1DVectorTy,
+                                                     operands[0], zero);
+        },
+        rewriter);
+  }
+};
+
+using CountLeadingZerosOpLowering =
+    CountOpLowering<math::CountLeadingZerosOp, LLVM::CountLeadingZerosOp>;
+using CountTrailingZerosOpLowering =
+    CountOpLowering<math::CountTrailingZerosOp, LLVM::CountTrailingZerosOp>;
+
 // A `expm1` is converted into `exp - 1`.
 struct ExpM1OpLowering : public ConvertOpToLLVMPattern<math::ExpM1Op> {
   using ConvertOpToLLVMPattern<math::ExpM1Op>::ConvertOpToLLVMPattern;
 
   LogicalResult
   matchAndRewrite(math::ExpM1Op op, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
     auto operandType = adaptor.getOperand().getType();

 void mlir::populateMathToLLVMConversionPatterns(LLVMTypeConverter &converter,
                                                 RewritePatternSet &patterns) {
   // clang-format off
   patterns.add<
     AbsOpLowering,
     CeilOpLowering,
     CopySignOpLowering,
     CosOpLowering,
+    CountLeadingZerosOpLowering,
+    CountTrailingZerosOpLowering,
     CtPopFOpLowering,
     ExpOpLowering,
     Exp2OpLowering,
     ExpM1OpLowering,
     FloorOpLowering,
     FmaOpLowering,
     Log10OpLowering,
     Log1pOpLowering,

mlir/test/Conversion/MathToLLVM/math-to-llvm.mlir

 func @sine(%arg0 : f32) {
   // CHECK: "llvm.intr.sin"(%arg0) : (f32) -> f32
   %0 = math.sin %arg0 : f32
   std.return
 }
 
 // -----
 
+// CHECK-LABEL: func @ctlz(
+// CHECK-SAME: i32
+func @ctlz(%arg0 : i32) {
+  // CHECK: %[[ZERO:.+]] = llvm.mlir.constant(false) : i1
+  // CHECK: "llvm.intr.ctlz"(%arg0, %[[ZERO]]) : (i32, i1) -> i32
+  %0 = math.ctlz %arg0 : i32
+  std.return
+}
+
+// -----
+
+// CHECK-LABEL: func @cttz(
+// CHECK-SAME: i32
+func @cttz(%arg0 : i32) {
+  // CHECK: %[[ZERO:.+]] = llvm.mlir.constant(false) : i1
+  // CHECK: "llvm.intr.cttz"(%arg0, %[[ZERO]]) : (i32, i1) -> i32
+  %0 = math.cttz %arg0 : i32
+  std.return
+}
+
+// -----
+
+// CHECK-LABEL: func @cttz_vec(
+// CHECK-SAME: i32
+func @cttz_vec(%arg0 : vector<4xi32>) {
+  // CHECK: %[[ZERO:.+]] = llvm.mlir.constant(false) : i1
+  // CHECK: "llvm.intr.cttz"(%arg0, %[[ZERO]]) : (vector<4xi32>, i1) -> vector<4xi32>
+  %0 = math.cttz %arg0 : vector<4xi32>
+  std.return
+}
+
+// -----
+
 // CHECK-LABEL: func @ctpop(
 // CHECK-SAME: i32
 func @ctpop(%arg0 : i32) {
   // CHECK: "llvm.intr.ctpop"(%arg0) : (i32) -> i32
   %0 = math.ctpop %arg0 : i32
   std.return
 }
 

mlir/test/Target/LLVMIR/llvmir-intrinsics.mlir

 llvm.func @bitreverse_test(%arg0: i32, %arg1: vector<8xi32>) {
   // CHECK: call i32 @llvm.bitreverse.i32
   "llvm.intr.bitreverse"(%arg0) : (i32) -> i32
   // CHECK: call <8 x i32> @llvm.bitreverse.v8i32
   "llvm.intr.bitreverse"(%arg1) : (vector<8xi32>) -> vector<8xi32>
   llvm.return
 }
 
+// CHECK-LABEL: @ctlz_test
+llvm.func @ctlz_test(%arg0: i32, %arg1: vector<8xi32>) {
+  %i1 = llvm.mlir.constant(false) : i1
+  // CHECK: call i32 @llvm.ctlz.i32
+  "llvm.intr.ctlz"(%arg0, %i1) : (i32, i1) -> i32
+  // CHECK: call <8 x i32> @llvm.ctlz.v8i32
+  "llvm.intr.ctlz"(%arg1, %i1) : (vector<8xi32>, i1) -> vector<8xi32>
+  llvm.return
+}
+
+// CHECK-LABEL: @cttz_test
+llvm.func @cttz_test(%arg0: i32, %arg1: vector<8xi32>) {
+  %i1 = llvm.mlir.constant(false) : i1
+  // CHECK: call i32 @llvm.cttz.i32
+  "llvm.intr.cttz"(%arg0, %i1) : (i32, i1) -> i32
+  // CHECK: call <8 x i32> @llvm.cttz.v8i32
+  "llvm.intr.cttz"(%arg1, %i1) : (vector<8xi32>, i1) -> vector<8xi32>
+  llvm.return
+}
+
 // CHECK-LABEL: @ctpop_test
 llvm.func @ctpop_test(%arg0: i32, %arg1: vector<8xi32>) {
   // CHECK: call i32 @llvm.ctpop.i32
   "llvm.intr.ctpop"(%arg0) : (i32) -> i32
   // CHECK: call <8 x i32> @llvm.ctpop.v8i32
   "llvm.intr.ctpop"(%arg1) : (vector<8xi32>) -> vector<8xi32>
   llvm.return
 }