[MLIR][SPIRVToLLVM] Added spv.constant conversion pattern for scalar and vector types

[georgemitenkov]

Summary

This patch introduces conversion pattern for spv.constant with scalar and vector types. There is a special case when the constant value is a signed/unsigned integer (vector of integers). Since LLVM dialect does not have signedness semantics, the types had to be converted to signless ints.

[ftynse]

mlir/lib/Conversion/SPIRVToLLVM/ConvertSPIRVToLLVM.cpp

157

Is this-> necessary here?

194

Getting a sign-extended value for originally-unsigned type sounds dangerous

[rriddle]

mlir/lib/Conversion/SPIRVToLLVM/ConvertSPIRVToLLVM.cpp

146

Please format this.

172

nit: Drop the newline.

173

dyn_cast can return null. Please either check for null here or use cast instead.

175

nit: Please avoid using the attribute values whenever you possibly can. They have to be reconstructed, meaning locking the context per element. Use the APInt range instead if you know it is an integer, getValues<APInt>().

179

getIntegerAttr can't fail, so there is no need to check anything here.

184

Seems like what you really want here is to do something like:

dstElementsAttr = dstElementsAttr.mapValues(signlessType, [](APInt value) { return value; });

That would remove the need for the loop above.

194

Please just use the APInt methods instead.

[georgemitenkov]

• [MLIR][SPIRVToLLVM] Addressed comments on constant conversion
• [MLIR][SPIRVToLLVM] Refactored the code to follow clang style

Changed code style:

• Ckang reformatted, removed extra newlines
• Used mapValues() with a lambda instead of the loop
• Removed sign extension for APInt values in favour of APInt method

Also used cast() instead of dyn_cast().

mlir/lib/Conversion/SPIRVToLLVM/ConvertSPIRVToLLVM.cpp

 #define DEBUG_TYPE "spirv-to-llvm-pattern"
 
 using namespace mlir;
 
 //===----------------------------------------------------------------------===//
 // Utility functions
 //===----------------------------------------------------------------------===//
 
+/// Returns true if the given type is a signed integer or vector type.
+static bool isSignedIntegerOrVector(Type type) {
+  if (type.isSignedInteger())
+    return true;
+  if (auto vecType = type.dyn_cast<VectorType>())
+    return vecType.getElementType().isSignedInteger();
+  return false;
+}
+
 /// Returns true if the given type is an unsigned integer or vector type
 static bool isUnsignedIntegerOrVector(Type type) {
   if (type.isUnsignedInteger())
     return true;
   if (auto vecType = type.dyn_cast<VectorType>())
     return vecType.getElementType().isUnsignedInteger();
   return false;
 }

 
 //===----------------------------------------------------------------------===//
 // Operation conversion
 //===----------------------------------------------------------------------===//
 
 namespace {
 
 class BitFieldInsertPattern
     : public SPIRVToLLVMConversion<spirv::BitFieldInsertOp> {
 public:
   using SPIRVToLLVMConversion<spirv::BitFieldInsertOp>::SPIRVToLLVMConversion;
 
   LogicalResult
   matchAndRewrite(spirv::BitFieldInsertOp op, ArrayRef<Value> operands,
                   ConversionPatternRewriter &rewriter) const override {
     auto srcType = op.getType();
     auto dstType = this->typeConverter.convertType(srcType);
     if (!dstType)
       return failure();
     Location loc = op.getLoc();
 
     // Broadcast `Offset` and `Count` to match the type of `Base` and `Insert`.
     // If `Base` is of a vector type, construct a vector that has:
     //  - same number of elements as `Base`
     //  - each element has the type that is the same as the type of `Offset` or
     //    `Count`
     //  - each element has the same value as `Offset` or `Count`
     Value offset;
     Value count;
     if (auto vectorType = srcType.dyn_cast<VectorType>()) {
       unsigned numElements = vectorType.getNumElements();
       offset =
           broadcast(loc, op.offset(), numElements, typeConverter, rewriter);
       count = broadcast(loc, op.count(), numElements, typeConverter, rewriter);
     } else {
       offset = op.offset();
       count = op.count();
     }
 
     // Create a mask with all bits set of the same type as `srcType`
     Value minusOne = createConstantAllBitsSet(loc, srcType, dstType, rewriter);
 
     // Need to cast `Offset` and `Count` if their bit width is different
     // from `Base` bit width.
     Value optionallyCastedCount =
         optionallyTruncateOrExtend(loc, count, dstType, rewriter);
     Value optionallyCastedOffset =
         optionallyTruncateOrExtend(loc, offset, dstType, rewriter);
 
     // Create a mask with bits set outside [Offset, Offset + Count - 1].
     Value maskShiftedByCount = rewriter.create<LLVM::ShlOp>(
         loc, dstType, minusOne, optionallyCastedCount);
     Value negated = rewriter.create<LLVM::XOrOp>(loc, dstType,
                                                  maskShiftedByCount, minusOne);
     Value maskShiftedByCountAndOffset = rewriter.create<LLVM::ShlOp>(
         loc, dstType, negated, optionallyCastedOffset);
     Value mask = rewriter.create<LLVM::XOrOp>(
         loc, dstType, maskShiftedByCountAndOffset, minusOne);
 
     // Extract unchanged bits from the `Base`  that are outside of
     // [Offset, Offset + Count - 1]. Then `or` with shifted `Insert`.
     Value baseAndMask =
         rewriter.create<LLVM::AndOp>(loc, dstType, op.base(), mask);
     Value insertShiftedByOffset = rewriter.create<LLVM::ShlOp>(
         loc, dstType, op.insert(), optionallyCastedOffset);
     rewriter.replaceOpWithNewOp<LLVM::OrOp>(op, dstType, baseAndMask,
                                             insertShiftedByOffset);
     return success();
   }
 };
 
+/// Converts SPIR-V ConstantOp with scalar or vector type.
+class ConstantScalarAndVectorPattern
+    : public SPIRVToLLVMConversion<spirv::ConstantOp> {
+public:
+  using SPIRVToLLVMConversion<spirv::ConstantOp>::SPIRVToLLVMConversion;
+
+  LogicalResult
+  matchAndRewrite(spirv::ConstantOp constOp, ArrayRef<Value> operands,
+                  ConversionPatternRewriter &rewriter) const override {
+    auto srcType = constOp.getType();
+    if (!srcType.isa<VectorType>() && !srcType.isIntOrFloat())
+      return failure();
+
+    auto dstType = typeConverter.convertType(srcType);
+    if (!dstType)
+      return failure();
+
+    // SPIR-V constant can be a signed/unsigned integer, which has to be
+    // casted to signless integer when converting to LLVM dialect. Removing the
+    // sign bit may have unexpected behaviour. However, it is better to handle
+    // it case-by-case, given that the purpose of the conversion is not to
+    // cover all possible corner cases.
+    if (isSignedIntegerOrVector(srcType) ||
+        isUnsignedIntegerOrVector(srcType)) {
+      auto *context = rewriter.getContext();
+      auto signlessType = IntegerType::get(getBitWidth(srcType), context);
+
+      if (srcType.isa<VectorType>()) {
+        auto dstElementsAttr = constOp.value().cast<DenseIntElementsAttr>();
+        rewriter.replaceOpWithNewOp<LLVM::ConstantOp>(
+            constOp, dstType,
+            dstElementsAttr.mapValues(
+                signlessType, [&](const APInt &value) { return value; }));
+        return success();
+      }
+      auto srcAttr = constOp.value().cast<IntegerAttr>();
+      auto dstAttr = rewriter.getIntegerAttr(signlessType, srcAttr.getValue());
+      rewriter.replaceOpWithNewOp<LLVM::ConstantOp>(constOp, dstType, dstAttr);
+      return success();
+    }
+    rewriter.replaceOpWithNewOp<LLVM::ConstantOp>(constOp, dstType, operands,
+                                                  constOp.getAttrs());
+    return success();
+  }
+};
+
 /// Converts SPIR-V operations that have straightforward LLVM equivalent
 /// into LLVM dialect operations.
 template <typename SPIRVOp, typename LLVMOp>
 class DirectConversionPattern : public SPIRVToLLVMConversion<SPIRVOp> {
 public:
   using SPIRVToLLVMConversion<SPIRVOp>::SPIRVToLLVMConversion;
 
   LogicalResult
patterns.insert<
       IComparePattern<spirv::SGreaterThanEqualOp, LLVM::ICmpPredicate::sge>,
       IComparePattern<spirv::SLessThanEqualOp, LLVM::ICmpPredicate::sle>,
       IComparePattern<spirv::SLessThanOp, LLVM::ICmpPredicate::slt>,
       IComparePattern<spirv::UGreaterThanOp, LLVM::ICmpPredicate::ugt>,
       IComparePattern<spirv::UGreaterThanEqualOp, LLVM::ICmpPredicate::uge>,
       IComparePattern<spirv::ULessThanEqualOp, LLVM::ICmpPredicate::ule>,
       IComparePattern<spirv::ULessThanOp, LLVM::ICmpPredicate::ult>,
 
+      // Constant op
+      ConstantScalarAndVectorPattern,
+
       // Function Call op
       FunctionCallPattern,
 
       // Logical ops
       DirectConversionPattern<spirv::LogicalAndOp, LLVM::AndOp>,
       DirectConversionPattern<spirv::LogicalOrOp, LLVM::OrOp>,
       IComparePattern<spirv::LogicalEqualOp, LLVM::ICmpPredicate::eq>,
       IComparePattern<spirv::LogicalNotEqualOp, LLVM::ICmpPredicate::ne>,

mlir/test/Conversion/SPIRVToLLVM/constant-op-to-llvm.mlir

+// RUN: mlir-opt -convert-spirv-to-llvm %s | FileCheck %s
+
+//===----------------------------------------------------------------------===//
+// spv.constant
+//===----------------------------------------------------------------------===//
+
+func @bool_constant_scalar() {
+  // CHECK: {{.*}} = llvm.mlir.constant(true) : !llvm.i1
+  %0 = spv.constant true
+  // CHECK: {{.*}} = llvm.mlir.constant(false) : !llvm.i1
+  %1 = spv.constant false
+  return
+}
+
+func @bool_constant_vector() {
+  // CHECK: {{.*}} = llvm.mlir.constant(dense<[true, false]> : vector<2xi1>) : !llvm<"<2 x i1>">
+  %0 = constant dense<[true, false]> : vector<2xi1>
+  // CHECK: {{.*}} = llvm.mlir.constant(dense<false> : vector<3xi1>) : !llvm<"<3 x i1>">
+  %1 = constant dense<false> : vector<3xi1>
+  return
+}
+
+func @integer_constant_scalar() {
+  // CHECK: {{.*}} = llvm.mlir.constant(0 : i8) : !llvm.i8
+  %0 = spv.constant  0 : i8
+  // CHECK: {{.*}} = llvm.mlir.constant(-5 : i64) : !llvm.i64
+  %1 = spv.constant -5 : si64
+  // CHECK: {{.*}} = llvm.mlir.constant(10 : i16) : !llvm.i16
+  %2 = spv.constant  10 : ui16
+  return
+}
+
+func @integer_constant_vector() {
+  // CHECK: {{.*}} = llvm.mlir.constant(dense<[2, 3]> : vector<2xi32>) : !llvm<"<2 x i32>">
+  %0 = spv.constant dense<[2, 3]> : vector<2xi32>
+  // CHECK: {{.*}} = llvm.mlir.constant(dense<-4> : vector<2xi32>) : !llvm<"<2 x i32>">
+  %1 = spv.constant dense<-4> : vector<2xsi32>
+  // CHECK: {{.*}} = llvm.mlir.constant(dense<[2, 3, 4]> : vector<3xi32>) : !llvm<"<3 x i32>">
+  %2 = spv.constant dense<[2, 3, 4]> : vector<3xui32>
+  return
+}
+
+func @float_constant_scalar() {
+  // CHECK: {{.*}} = llvm.mlir.constant(5.000000e+00 : f16) : !llvm.half
+  %0 = spv.constant 5.000000e+00 : f16
+  // CHECK: {{.*}} = llvm.mlir.constant(5.000000e+00 : f64) : !llvm.double
+  %1 = spv.constant 5.000000e+00 : f64
+  return
+}
+
+func @float_constant_vector() {
+  // CHECK: {{.*}} = llvm.mlir.constant(dense<[2.000000e+00, 3.000000e+00]> : vector<2xf32>) : !llvm<"<2 x float>">
+  %0 = spv.constant dense<[2.000000e+00, 3.000000e+00]> : vector<2xf32>
+  return
+}