[mlir][spirv] Convert math.ctlz to spv.GLSL.FindUMsb

[antiagainst]

Add some negative tests

[ThomasRaoux]

mlir/lib/Conversion/MathToSPIRV/MathToSPIRV.cpp

41

nit: remove braces for single line if

mlir/include/mlir/Dialect/SPIRV/IR/SPIRVGLSLOps.td

def SPV_GLSLFMixOp :
 
   let assemblyFormat = [{
     attr-dict $x `:` type($x) `,` $y `:` type($y) `,` $a `:` type($a) `->` type($result)
   }];
 
   let hasVerifier = 0;
 }
 
+def SPV_GLSLFindUMsbOp : SPV_GLSLUnaryArithmeticOp<"FindUMsb", 75, SPV_Int32> {
+  let summary = "Unsigned-integer most-significant bit";
+
+  let description = [{
+    Results in the bit number of the most-significant 1-bit in the binary
+    representation of Value. If Value is 0, the result is -1.
+
+    Result Type and the type of Value must both be integer scalar or
+    integer vector types. Result Type and operand types must have the
+    same number of components with the same component width. Results are
+    computed per component.
+
+    This instruction is currently limited to 32-bit width components.
+  }];
+}
+
 #endif // MLIR_DIALECT_SPIRV_IR_GLSL_OPS

mlir/lib/Conversion/MathToSPIRV/MathToSPIRV.cpp

 //===----------------------------------------------------------------------===//
 
 #include "../SPIRVCommon/Pattern.h"
 #include "mlir/Dialect/Math/IR/Math.h"
 #include "mlir/Dialect/SPIRV/IR/SPIRVDialect.h"
 #include "mlir/Dialect/SPIRV/IR/SPIRVOps.h"
 #include "mlir/Dialect/SPIRV/Transforms/SPIRVConversion.h"
 #include "mlir/IR/BuiltinTypes.h"
+#include "mlir/Transforms/DialectConversion.h"
 #include "llvm/Support/Debug.h"
 
 #define DEBUG_TYPE "math-to-spirv-pattern"
 
 using namespace mlir;
 
 //===----------------------------------------------------------------------===//
+// Utility functions
+//===----------------------------------------------------------------------===//
+
+/// Creates a 32-bit scalar/vector integer constant. Returns nullptr if the
+/// given type is not a 32-bit scalar/vector type.
+static Value getScalarOrVectorI32Constant(Type type, int value,
+                                          OpBuilder &builder, Location loc) {
+  if (auto vectorType = type.dyn_cast<VectorType>()) {
+    if (!vectorType.getElementType().isInteger(32))
+      return nullptr;
+    SmallVector<int> values(vectorType.getNumElements(), value);
+    return builder.create<spirv::ConstantOp>(loc, type,
+                                             builder.getI32VectorAttr(values));
+  }
+  if (type.isInteger(32)) {
+    return builder.create<spirv::ConstantOp>(loc, type,
+                                             builder.getI32IntegerAttr(value));
+  }
+  return nullptr;
+}
+
+//===----------------------------------------------------------------------===//
 // Operation conversion
 //===----------------------------------------------------------------------===//
 
 // Note that DRR cannot be used for the patterns in this file: we may need to
 // convert type along the way, which requires ConversionPattern. DRR generates
 // normal RewritePattern.
 
 namespace {
matchAndRewrite(math::CopySignOp copySignOp, OpAdaptor adaptor,
 
     Value result = rewriter.create<spirv::BitwiseOrOp>(loc, intType,
                                                        ValueRange{value, sign});
     rewriter.replaceOpWithNewOp<spirv::BitcastOp>(copySignOp, type, result);
     return success();
   }
 };
 
+/// Converts math.ctlz to SPIR-V ops.
+///
+/// SPIR-V does not have a direct operations for counting leading zeros. If
+/// Shader capability is supported, we can leverage GLSL FindUMsb to calculate
+/// it.
+class CountLeadingZerosPattern final
+    : public OpConversionPattern<math::CountLeadingZerosOp> {
+  using OpConversionPattern::OpConversionPattern;
+
+  LogicalResult
+  matchAndRewrite(math::CountLeadingZerosOp countOp, OpAdaptor adaptor,
+                  ConversionPatternRewriter &rewriter) const override {
+    auto type = getTypeConverter()->convertType(countOp.getType());
+    if (!type)
+      return failure();
+
+    // We can only support 32-bit integer types for now.
+    unsigned bitwidth = 0;
+    if (type.isa<IntegerType>())
+      bitwidth = type.getIntOrFloatBitWidth();
+    if (auto vectorType = type.dyn_cast<VectorType>())
+      bitwidth = vectorType.getElementTypeBitWidth();
+    if (bitwidth != 32)
+      return failure();
+
+    Location loc = countOp.getLoc();
+    Value val31 = getScalarOrVectorI32Constant(type, 31, rewriter, loc);
+    Value msb =
+        rewriter.create<spirv::GLSLFindUMsbOp>(loc, adaptor.getOperand());
+    // We need to subtract from 31 given that the index is from the least
+    // significant bit.
+    rewriter.replaceOpWithNewOp<spirv::ISubOp>(countOp, val31, msb);
+    return success();
+  }
+};
+
 /// Converts math.expm1 to SPIR-V ops.
 ///
 /// SPIR-V does not have a direct operations for exp(x)-1. Explicitly lower to
 /// these operations.
 template <typename ExpOp>
 struct ExpM1OpPattern final : public OpConversionPattern<math::ExpM1Op> {
   using OpConversionPattern::OpConversionPattern;
 

 namespace mlir {
 void populateMathToSPIRVPatterns(SPIRVTypeConverter &typeConverter,
                                  RewritePatternSet &patterns) {
   // Core patterns
   patterns.add<CopySignPattern>(typeConverter, patterns.getContext());
 
   // GLSL patterns
   patterns
-      .add<Log1pOpPattern<spirv::GLSLLogOp>, ExpM1OpPattern<spirv::GLSLExpOp>,
+      .add<CountLeadingZerosPattern, Log1pOpPattern<spirv::GLSLLogOp>,
+           ExpM1OpPattern<spirv::GLSLExpOp>,
            spirv::ElementwiseOpPattern<math::AbsOp, spirv::GLSLFAbsOp>,
            spirv::ElementwiseOpPattern<math::CeilOp, spirv::GLSLCeilOp>,
            spirv::ElementwiseOpPattern<math::CosOp, spirv::GLSLCosOp>,
            spirv::ElementwiseOpPattern<math::ExpOp, spirv::GLSLExpOp>,
            spirv::ElementwiseOpPattern<math::FloorOp, spirv::GLSLFloorOp>,
            spirv::ElementwiseOpPattern<math::FmaOp, spirv::GLSLFmaOp>,
            spirv::ElementwiseOpPattern<math::LogOp, spirv::GLSLLogOp>,
            spirv::ElementwiseOpPattern<math::PowFOp, spirv::GLSLPowOp>,

mlir/lib/Conversion/MathToSPIRV/MathToSPIRVPass.cpp

void ConvertMathToSPIRVPass::runOnOperation() {
   ModuleOp module = getOperation();
 
   auto targetAttr = spirv::lookupTargetEnvOrDefault(module);
   std::unique_ptr<ConversionTarget> target =
       SPIRVConversionTarget::get(targetAttr);
 
   SPIRVTypeConverter typeConverter(targetAttr);
 
+  // Use UnrealizedConversionCast as the bridge so that we don't need to pull
+  // in patterns for other dialects.
+  auto addUnrealizedCast = [](OpBuilder &builder, Type type, ValueRange inputs,
+                              Location loc) {
+    auto cast = builder.create<UnrealizedConversionCastOp>(loc, type, inputs);
+    return Optional<Value>(cast.getResult(0));
+  };
+  typeConverter.addSourceMaterialization(addUnrealizedCast);
+  typeConverter.addTargetMaterialization(addUnrealizedCast);
+  target->addLegalOp<UnrealizedConversionCastOp>();
+
   RewritePatternSet patterns(context);
   populateMathToSPIRVPatterns(typeConverter, patterns);
 
   if (failed(applyPartialConversion(module, *target, std::move(patterns))))
     return signalPassFailure();
 }
 
 std::unique_ptr<OperationPass<ModuleOp>> mlir::createConvertMathToSPIRVPass() {
   return std::make_unique<ConvertMathToSPIRVPass>();
 }

mlir/test/Conversion/MathToSPIRV/math-to-glsl-spirv.mlir

 // CHECK-LABEL: @float32_ternary_vector
 func.func @float32_ternary_vector(%a: vector<4xf32>, %b: vector<4xf32>,
                             %c: vector<4xf32>) {
   // CHECK: spv.GLSL.Fma %{{.*}}: vector<4xf32>
   %0 = math.fma %a, %b, %c : vector<4xf32>
   return
 }
 
+// CHECK-LABEL: @ctlz_scalar
+//  CHECK-SAME: (%[[VAL:.+]]: i32)
+func.func @ctlz_scalar(%val: i32) -> i32 {
+  // CHECK: %[[V31:.+]] = spv.Constant 31 : i32
+  // CHECK: %[[MSB:.+]] = spv.GLSL.FindUMsb %[[VAL]] : i32
+  // CHECK: %[[SUB:.+]] = spv.ISub %[[V31]], %[[MSB]] : i32
+  // CHECK: return %[[SUB]]
+  %0 = math.ctlz %val : i32
+  return %0 : i32
+}
+
+// CHECK-LABEL: @ctlz_vector1
+func.func @ctlz_vector1(%val: vector<1xi32>) -> vector<1xi32> {
+  // CHECK: spv.GLSL.FindUMsb
+  // CHECK: spv.ISub
+  %0 = math.ctlz %val : vector<1xi32>
+  return %0 : vector<1xi32>
+}
+
+// CHECK-LABEL: @ctlz_vector2
+//  CHECK-SAME: (%[[VAL:.+]]: vector<2xi32>)
+func.func @ctlz_vector2(%val: vector<2xi32>) -> vector<2xi32> {
+  // CHECK-DAG: %[[V31:.+]] = spv.Constant dense<31> : vector<2xi32>
+  // CHECK: %[[MSB:.+]] = spv.GLSL.FindUMsb %[[VAL]] : vector<2xi32>
+  // CHECK: %[[SUB:.+]] = spv.ISub %[[V31]], %[[MSB]] : vector<2xi32>
+  // CHECK: return %[[SUB]]
+  %0 = math.ctlz %val : vector<2xi32>
+  return %0 : vector<2xi32>
+}
+
 } // end module

mlir/test/Dialect/SPIRV/IR/glsl-ops.mlir

 //===----------------------------------------------------------------------===//
 
 func.func @fmix(%arg0 : f32, %arg1 : f32, %arg2 : f32) -> () {
   // CHECK: {{%.*}} = spv.GLSL.FMix {{%.*}} : f32, {{%.*}} : f32, {{%.*}} : f32 -> f32
   %0 = spv.GLSL.FMix %arg0 : f32, %arg1 : f32, %arg2 : f32 -> f32
   return
 }
 
-// -----
 func.func @fmix_vector(%arg0 : vector<3xf32>, %arg1 : vector<3xf32>, %arg2 : vector<3xf32>) -> () {
   // CHECK: {{%.*}} = spv.GLSL.FMix {{%.*}} : vector<3xf32>, {{%.*}} : vector<3xf32>, {{%.*}} : vector<3xf32> -> vector<3xf32>
   %0 = spv.GLSL.FMix %arg0 : vector<3xf32>, %arg1 : vector<3xf32>, %arg2 : vector<3xf32> -> vector<3xf32>
   return
 }
 
+// -----
+
+//===----------------------------------------------------------------------===//
+// spv.GLSL.Exp
+//===----------------------------------------------------------------------===//
+
+func.func @findumsb(%arg0 : i32) -> () {
+  // CHECK: spv.GLSL.FindUMsb {{%.*}} : i32
+  %2 = spv.GLSL.FindUMsb %arg0 : i32
+  return
+}
+
+func.func @findumsb_vector(%arg0 : vector<3xi32>) -> () {
+  // CHECK: spv.GLSL.FindUMsb {{%.*}} : vector<3xi32>
+  %2 = spv.GLSL.FindUMsb %arg0 : vector<3xi32>
+  return
+}
+
+// -----
+
+func.func @findumsb(%arg0 : i64) -> () {
+  // expected-error @+1 {{operand #0 must be Int32 or vector of Int32}}
+  %2 = spv.GLSL.FindUMsb %arg0 : i64
+  return
+}

mlir/test/Target/SPIRV/glsl-ops.mlir

spv.func @sclamp(%arg0 : si32, %arg1 : si32, %arg2 : si32) "None" {
     spv.Return
   }
 
   spv.func @fma(%arg0 : f32, %arg1 : f32, %arg2 : f32) "None" {
     // CHECK: spv.GLSL.Fma {{%[^,]*}}, {{%[^,]*}}, {{%[^,]*}} : f32
     %13 = spv.GLSL.Fma %arg0, %arg1, %arg2 : f32
     spv.Return
   }
+
+  spv.func @findumsb(%arg0 : i32) "None" {
+    // CHECK: spv.GLSL.FindUMsb {{%.*}} : i32
+    %2 = spv.GLSL.FindUMsb %arg0 : i32
+    spv.Return
+  }
 }