[mlir] Add CtPop to MathOps with lowering to LLVM

math.ctpop maths to the llvm.ctpop intrinsic.

Reviewed By: ftynse

Differential Revision: https://reviews.llvm.org/D114998

diff --git a/mlir/include/mlir/Dialect/Math/IR/MathOps.td b/mlir/include/mlir/Dialect/Math/IR/MathOps.td
index 2cc977f..42a5b57 100644 (file)
--- a/mlir/include/mlir/Dialect/Math/IR/MathOps.td
+++ b/mlir/include/mlir/Dialect/Math/IR/MathOps.td
@@ -21,6 +21,17 @@ class Math_Op<string mnemonic, list<OpTrait> traits = []> :
     DeclareOpInterfaceMethods<VectorUnrollOpInterface>] #
     ElementwiseMappable.traits>;
 
+// Base class for unary math operations on integer types. Require a operand and
+// result of the same type. This type can be an integer type, or vector or tensor
+// thereof.
+class Math_IntegerUnaryOp<string mnemonic, list<OpTrait> traits = []> :
+    Math_Op<mnemonic, traits # [SameOperandsAndResultType]> {
+  let arguments = (ins SignlessIntegerLike:$operand);
+  let results = (outs SignlessIntegerLike:$result);
+
+  let assemblyFormat = "$operand attr-dict `:` type($result)";
+}
+
 // Base class for unary math operations on floating point types. Require a
 // operand and result of the same type. This type can be a floating point type,
 // or vector or tensor thereof.
@@ -287,6 +298,30 @@ def Math_SinOp : Math_FloatUnaryOp<"sin"> {
 }
 
 //===----------------------------------------------------------------------===//
+// 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.
+
+    Example:
+
+    ```mlir
+    // Scalar ctpop function value.
+    %a = math.ctpop %b : i32
+
+    // SIMD vector element-wise ctpop function value.
+    %f = math.ctpop %g : vector<4xi16>
+
+    // Tensor element-wise ctpop function value.
+    %x = math.ctpop %y : tensor<4x?xi8>
+    ```
+  }];
+}
+
+//===----------------------------------------------------------------------===//
 // ErfOp
 //===----------------------------------------------------------------------===//
 

diff --git a/mlir/lib/Conversion/MathToLLVM/MathToLLVM.cpp b/mlir/lib/Conversion/MathToLLVM/MathToLLVM.cpp
index 29c6ad3..551cefe 100644 (file)
--- a/mlir/lib/Conversion/MathToLLVM/MathToLLVM.cpp
+++ b/mlir/lib/Conversion/MathToLLVM/MathToLLVM.cpp
@@ -23,6 +23,8 @@ using CeilOpLowering = VectorConvertToLLVMPattern<math::CeilOp, LLVM::FCeilOp>;
 using CopySignOpLowering =
     VectorConvertToLLVMPattern<math::CopySignOp, LLVM::CopySignOp>;
 using CosOpLowering = VectorConvertToLLVMPattern<math::CosOp, LLVM::CosOp>;
+using CtPopFOpLowering =
+    VectorConvertToLLVMPattern<math::CtPopOp, LLVM::CtPopOp>;
 using ExpOpLowering = VectorConvertToLLVMPattern<math::ExpOp, LLVM::ExpOp>;
 using Exp2OpLowering = VectorConvertToLLVMPattern<math::Exp2Op, LLVM::Exp2Op>;
 using FloorOpLowering =
@@ -220,6 +222,7 @@ void mlir::populateMathToLLVMConversionPatterns(LLVMTypeConverter &converter,
     CeilOpLowering,
     CopySignOpLowering,
     CosOpLowering,
+    CtPopFOpLowering,
     ExpOpLowering,
     Exp2OpLowering,
     ExpM1OpLowering,

diff --git a/mlir/test/Conversion/MathToLLVM/math-to-llvm.mlir b/mlir/test/Conversion/MathToLLVM/math-to-llvm.mlir
index 3eeaea2..2b22ef9 100644 (file)
--- a/mlir/test/Conversion/MathToLLVM/math-to-llvm.mlir
+++ b/mlir/test/Conversion/MathToLLVM/math-to-llvm.mlir
@@ -74,6 +74,26 @@ func @sine(%arg0 : f32) {
 
 // -----
 
+// 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
+}
+
+// -----
+
+// CHECK-LABEL: func @ctpop_vector(
+// CHECK-SAME: vector<3xi32>
+func @ctpop_vector(%arg0 : vector<3xi32>) {
+  // CHECK: "llvm.intr.ctpop"(%arg0) : (vector<3xi32>) -> vector<3xi32>
+  %0 = math.ctpop %arg0 : vector<3xi32>
+  std.return
+}
+
+// -----
+
 // CHECK-LABEL: func @rsqrt_double(
 // CHECK-SAME: f64
 func @rsqrt_double(%arg0 : f64) {