[MLIR][SPIRVToLLVM] Conversion of SPIR-V variable op

[georgemitenkov]

Summary

The patch introduces the conversion pattern for function-level
spv.Variable. It is modelled as llvm.alloca op. If initialized, then
additional store instruction is used. Note that there is no initialization
for arrays and structs since constants of these types are not supported in
LLVM dialect yet. Also, at the moment initialisation is only possible via
spv.constant (since spv.GlobalVariable conversion is not implemented
yet).

[ftynse]

Just as a warning, memory allocated by alloca is live throughout the execution of a function. In the past, we ran into stack overflows because transforming ops executed in a loop into allocas. If the input code has some scoping, it should be reflected with an AutomaticAllocationScope interface that we can then lower to a stacksave/stackrestore pair (TBD).

Note that there is no initialization for arrays and structs since constants of these types are not supported in LLVM dialect yet.

I think we have array and struct initializers for llvm.mlir.global. It should be possible to port those on llvm.mlir.constant as well. Contributions most welcome!

mlir/lib/Conversion/SPIRVToLLVM/ConvertSPIRVToLLVM.cpp

183

Nit: prefer explicit types to auto when it doesn't hamper readability

751

Should this also have a check for supported types and bail out on unsupported types (arrays, structs)?

[georgemitenkov]

Addressed comments. Scoping will be dealt with separately.

mlir/lib/Conversion/SPIRVToLLVM/ConvertSPIRVToLLVM.cpp

static Type convertStructTypePacked(spirv::StructType type,
   auto elementsVector = llvm::to_vector<8>(
       llvm::map_range(type.getElementTypes(), [&](Type elementType) {
         return converter.convertType(elementType).cast<LLVM::LLVMType>();
       }));
   return LLVM::LLVMType::getStructTy(converter.getDialect(), elementsVector,
                                      /*isPacked=*/true);
 }
 
+/// Creates LLVM dialect constant with the given value.
+static Value createI32ConstantOf(Location loc, PatternRewriter &rewriter,
+                                 LLVMTypeConverter &converter, unsigned value) {
+  return rewriter.create<LLVM::ConstantOp>(
+      loc, LLVM::LLVMType::getInt32Ty(converter.getDialect()),
+      rewriter.getIntegerAttr(rewriter.getI32Type(), value));
+}
+
 //===----------------------------------------------------------------------===//
 // Type conversion
 //===----------------------------------------------------------------------===//
 
 /// Converts SPIR-V array type to LLVM array. There is no modelling of array
 /// stride at the moment.
 static Optional<Type> convertArrayType(spirv::ArrayType type,
                                        TypeConverter &converter) {
matchAndRewrite(SPIRVOp operation, ArrayRef<Value> operands,
     }
     Value result = rewriter.template create<LLVMOp>(
         loc, dstType, operation.operand1(), extended);
     rewriter.replaceOp(operation, result);
     return success();
   }
 };
 
+class VariablePattern : public SPIRVToLLVMConversion<spirv::VariableOp> {
+public:
+  using SPIRVToLLVMConversion<spirv::VariableOp>::SPIRVToLLVMConversion;
+
+  LogicalResult
+  matchAndRewrite(spirv::VariableOp varOp, ArrayRef<Value> operands,
+                  ConversionPatternRewriter &rewriter) const override {
+    auto srcType = varOp.getType();
+    // Initialization is supported for scalars and vectors only.
+    auto pointerTo = srcType.cast<spirv::PointerType>().getPointeeType();
+    auto init = varOp.initializer();
+    if (init && !pointerTo.isIntOrFloat() && !pointerTo.isa<VectorType>())
+      return failure();
+
+    auto dstType = typeConverter.convertType(srcType);
+    if (!dstType)
+      return failure();
+
+    Location loc = varOp.getLoc();
+    Value size = createI32ConstantOf(loc, rewriter, typeConverter, 1);
+    if (!init) {
+      rewriter.replaceOpWithNewOp<LLVM::AllocaOp>(varOp, dstType, size);
+      return success();
+    }
+    Value allocated = rewriter.create<LLVM::AllocaOp>(loc, dstType, size);
+    rewriter.create<LLVM::StoreOp>(loc, init, allocated);
+    rewriter.replaceOp(varOp, allocated);
+    return success();
+  }
+};
+
 //===----------------------------------------------------------------------===//
 // FuncOp conversion
 //===----------------------------------------------------------------------===//
 
 class FuncConversionPattern : public SPIRVToLLVMConversion<spirv::FuncOp> {
 public:
   using SPIRVToLLVMConversion<spirv::FuncOp>::SPIRVToLLVMConversion;
 
patterns.insert<
 
       // Logical ops
       DirectConversionPattern<spirv::LogicalAndOp, LLVM::AndOp>,
       DirectConversionPattern<spirv::LogicalOrOp, LLVM::OrOp>,
       IComparePattern<spirv::LogicalEqualOp, LLVM::ICmpPredicate::eq>,
       IComparePattern<spirv::LogicalNotEqualOp, LLVM::ICmpPredicate::ne>,
       NotPattern<spirv::LogicalNotOp>,
 
+      // Memory ops
+      VariablePattern,
+
       // Miscellaneous ops
       DirectConversionPattern<spirv::SelectOp, LLVM::SelectOp>,
       DirectConversionPattern<spirv::UndefOp, LLVM::UndefOp>,
 
       // Shift ops
       ShiftPattern<spirv::ShiftRightArithmeticOp, LLVM::AShrOp>,
       ShiftPattern<spirv::ShiftRightLogicalOp, LLVM::LShrOp>,
       ShiftPattern<spirv::ShiftLeftLogicalOp, LLVM::ShlOp>,

mlir/test/Conversion/SPIRVToLLVM/memory-ops-to-llvm.mlir

+// RUN: mlir-opt -convert-spirv-to-llvm %s | FileCheck %s
+
+//===----------------------------------------------------------------------===//
+// spv.Variable
+//===----------------------------------------------------------------------===//
+
+func @variable_scalar() {
+  // CHECK: %[[SIZE1:.*]] = llvm.mlir.constant(1 : i32) : !llvm.i32
+  // CHECK: %{{.*}} = llvm.alloca %[[SIZE1]] x !llvm.float : (!llvm.i32) -> !llvm<"float*">
+  %0 = spv.Variable : !spv.ptr<f32, Function>
+  // CHECK: %[[SIZE2:.*]] = llvm.mlir.constant(1 : i32) : !llvm.i32
+  // CHECK: %{{.*}} = llvm.alloca %[[SIZE2]] x !llvm.i8 : (!llvm.i32) -> !llvm<"i8*">
+  %1 = spv.Variable : !spv.ptr<i8, Function>
+  return
+}
+
+func @variable_scalar_with_initialization() {
+  // CHECK: %[[VALUE:.*]] = llvm.mlir.constant(0 : i64) : !llvm.i64
+  // CHECK: %[[SIZE:.*]] = llvm.mlir.constant(1 : i32) : !llvm.i32
+  // CHECK: %[[ALLOCATED:.*]] = llvm.alloca %[[SIZE]] x !llvm.i64 : (!llvm.i32) -> !llvm<"i64*">
+  // CHECK: llvm.store %[[VALUE]], %[[ALLOCATED]] : !llvm<"i64*">
+  %c = spv.constant 0 : i64
+  %0 = spv.Variable init(%c) : !spv.ptr<i64, Function>
+  return
+}
+
+func @variable_vector() {
+  // CHECK: %[[SIZE:.*]] = llvm.mlir.constant(1 : i32) : !llvm.i32
+  // CHECK: %{{.*}} = llvm.alloca  %[[SIZE]] x !llvm<"<3 x float>"> : (!llvm.i32) -> !llvm<"<3 x float>*">
+  %0 = spv.Variable : !spv.ptr<vector<3xf32>, Function>
+  return
+}
+
+func @variable_vector_with_initialization() {
+  // CHECK: %[[VALUE:.*]] = llvm.mlir.constant(dense<false> : vector<3xi1>) : !llvm<"<3 x i1>">
+  // CHECK: %[[SIZE:.*]] = llvm.mlir.constant(1 : i32) : !llvm.i32
+  // CHECK: %[[ALLOCATED:.*]] = llvm.alloca %[[SIZE]] x !llvm<"<3 x i1>"> : (!llvm.i32) -> !llvm<"<3 x i1>*">
+  // CHECK: llvm.store %[[VALUE]], %[[ALLOCATED]] : !llvm<"<3 x i1>*">
+  %c = spv.constant dense<false> : vector<3xi1>
+  %0 = spv.Variable init(%c) : !spv.ptr<vector<3xi1>, Function>
+  return
+}
+
+func @variable_array() {
+  // CHECK: %[[SIZE:.*]] = llvm.mlir.constant(1 : i32) : !llvm.i32
+  // CHECK: %{{.*}} = llvm.alloca %[[SIZE]] x !llvm<"[10 x i32]"> : (!llvm.i32) -> !llvm<"[10 x i32]*">
+  %0 = spv.Variable : !spv.ptr<!spv.array<10 x i32>, Function>
+  return
+}