[Mlir-commits] [mlir] [mlir][SPIR-V] Convert math.cttz to spirv.GLFindILsb (PR #200455)

[llvmlistbot at llvm.org]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir-spirv

Author: Arseniy Obolenskiy (aobolensk)

<details>
<summary>Changes</summary>



---
Full diff: https://github.com/llvm/llvm-project/pull/200455.diff


2 Files Affected:

- (modified) mlir/lib/Conversion/MathToSPIRV/MathToSPIRV.cpp (+41-1) 
- (modified) mlir/test/Conversion/MathToSPIRV/math-to-gl-spirv.mlir (+22) 


``````````diff
diff --git a/mlir/lib/Conversion/MathToSPIRV/MathToSPIRV.cpp b/mlir/lib/Conversion/MathToSPIRV/MathToSPIRV.cpp
index cf5ea9716b9f7..6eeb8f842a87c 100644
--- a/mlir/lib/Conversion/MathToSPIRV/MathToSPIRV.cpp
+++ b/mlir/lib/Conversion/MathToSPIRV/MathToSPIRV.cpp
@@ -242,6 +242,45 @@ struct CountLeadingZerosPattern final
   }
 };
 
+/// Converts math.cttz to GL FindILsb. GL FindILsb returns -1 for a zero
+/// input while math.cttz must return the bitwidth, so the zero case is
+/// patched up with a select.
+struct CountTrailingZerosPattern final
+    : public OpConversionPattern<math::CountTrailingZerosOp> {
+  using Base::Base;
+
+  LogicalResult
+  matchAndRewrite(math::CountTrailingZerosOp countOp, OpAdaptor adaptor,
+                  ConversionPatternRewriter &rewriter) const override {
+    if (LogicalResult res = checkSourceOpTypes(rewriter, countOp); failed(res))
+      return res;
+
+    Type type = getTypeConverter()->convertType(countOp.getType());
+    if (!type)
+      return failure();
+
+    unsigned bitwidth = 0;
+    if (isa<IntegerType>(type))
+      bitwidth = type.getIntOrFloatBitWidth();
+    if (auto vectorType = dyn_cast<VectorType>(type))
+      bitwidth = vectorType.getElementTypeBitWidth();
+    if (bitwidth != 32)
+      return failure();
+
+    Location loc = countOp.getLoc();
+    Value input = adaptor.getOperand();
+    Value val0 = getScalarOrVectorI32Constant(type, 0, rewriter, loc);
+    Value valBitwidth =
+        getScalarOrVectorI32Constant(type, bitwidth, rewriter, loc);
+
+    Value lsb = spirv::GLFindILsbOp::create(rewriter, loc, input);
+    Value isZero = spirv::IEqualOp::create(rewriter, loc, input, val0);
+    rewriter.replaceOpWithNewOp<spirv::SelectOp>(countOp, isZero, valBitwidth,
+                                                 lsb);
+    return success();
+  }
+};
+
 /// Converts math.expm1 to SPIR-V ops.
 ///
 /// SPIR-V does not have a direct operations for exp(x)-1. Explicitly lower to
@@ -530,7 +569,8 @@ void populateMathToSPIRVPatterns(const SPIRVTypeConverter &typeConverter,
 
   // GLSL patterns
   patterns.add<
-      CountLeadingZerosPattern, Log1pOpPattern<spirv::GLLogOp>, Log10OpPattern,
+      CountLeadingZerosPattern, CountTrailingZerosPattern,
+      Log1pOpPattern<spirv::GLLogOp>, Log10OpPattern,
       ExpM1OpPattern<spirv::GLExpOp>, PowFOpPattern, RoundOpPattern,
       CheckedElementwiseOpPattern<math::AbsFOp, spirv::GLFAbsOp>,
       CheckedElementwiseOpPattern<math::AbsIOp, spirv::GLSAbsOp>,
diff --git a/mlir/test/Conversion/MathToSPIRV/math-to-gl-spirv.mlir b/mlir/test/Conversion/MathToSPIRV/math-to-gl-spirv.mlir
index 08d7822d04cc1..94c285ddb0ccc 100644
--- a/mlir/test/Conversion/MathToSPIRV/math-to-gl-spirv.mlir
+++ b/mlir/test/Conversion/MathToSPIRV/math-to-gl-spirv.mlir
@@ -183,6 +183,28 @@ func.func @ctlz_vector2(%val: vector<2xi32>) -> vector<2xi32> {
   return %0 : vector<2xi32>
 }
 
+// CHECK-LABEL: @cttz_scalar
+//  CHECK-SAME: (%[[VAL:.+]]: i32)
+func.func @cttz_scalar(%val: i32) -> i32 {
+  // CHECK-DAG: %[[V0:.+]] = spirv.Constant 0 : i32
+  // CHECK-DAG: %[[V32:.+]] = spirv.Constant 32 : i32
+  // CHECK: %[[LSB:.+]] = spirv.GL.FindILsb %[[VAL]] : i32
+  // CHECK: %[[CMP:.+]] = spirv.IEqual %[[VAL]], %[[V0]] : i32
+  // CHECK: %[[R:.+]] = spirv.Select %[[CMP]], %[[V32]], %[[LSB]] : i1, i32
+  // CHECK: return %[[R]]
+  %0 = math.cttz %val : i32
+  return %0 : i32
+}
+
+// CHECK-LABEL: @cttz_vector
+func.func @cttz_vector(%val: vector<2xi32>) -> vector<2xi32> {
+  // CHECK: spirv.GL.FindILsb
+  // CHECK: spirv.IEqual
+  // CHECK: spirv.Select
+  %0 = math.cttz %val : vector<2xi32>
+  return %0 : vector<2xi32>
+}
+
 // Dynamic exponent: exp(y * log(x)); yields NaN for x<0.
 // CHECK-LABEL: @powf_scalar
 //  CHECK-SAME: (%[[LHS:.+]]: f32, %[[RHS:.+]]: f32)

``````````

</details>


https://github.com/llvm/llvm-project/pull/200455