[MLIR][SPIRVToLLVM] Enhanced conversion for entry point ops

[georgemitenkov]

Summary

This patch introduces a new conversion pattern for spv.EntryPoint
and spv.ExecutionMode. Before, these ops were just erased. Now, an
approach is to encode the information from these ops in a global struct
variable of the following form:

struct {
  int32_t; // specify execution model
  int32_t; // specify execution mode
  // TODO: interface vars & mode's values
}

At the moment, no interface variables info is passed to the global, as
well as execution mode's optional value. This is a work in progress.

[mravishankar]

Curious what you need this information for?

mlir/lib/Conversion/SPIRVToLLVM/ConvertSPIRVToLLVM.cpp

660

This logic seems fairly involved. Probably doesnt need to be. You should be able to

• get the current return value which is in the terminator, this is previousStruct
• update the entry point as you have done here to get newStruct, except that setEntryPointInfo collects all the operations that uses previousStruct, say validUses
• and then do previousStruct->replaceAllUsesExcept(newStruct, validUses).

[georgemitenkov]

Execution Mode ops may contain important information about the entry point, which we may want to preserve. An example can be LocalSize which indicates the work-group size. This can be reused when handling multi-threaded code conversion.

This patch is more of a skeleton to relate to during the discussion. Maybe it will be better to also have a post on the discourse to see what other people think. What do you think?

I will also address comments, thanks for the feedback!

[mravishankar]

mlir/lib/Conversion/SPIRVToLLVM/ConvertSPIRVToLLVM.cpp

663

COuldnt this all be simplified with setInsertionPoint(terminator)?

[georgemitenkov]

After some more thinking and further reading, I realise that entry point information is not really needed. However, it is important to keep execution mode info. This changes the conversion pattern as follows:

• spv.EntryPoint is simply removed
• Info from spv.ExecutionMode is used to create a global variable, which now looks like:

struct {
   int32_t executionMode;
   int32_t values[];          // optional values
};

This is a better approach in my opinion since we

• Remove unnecessary complexity compared to previous version of the patch
• Handle actually useful info - e.g. local sizes, etc.
• We have a 1-to-1 op mapping: spv.ExecutionMode -> llvm.mlir.global

[georgemitenkov]

Fixed a different global naming when the parent module does not have a name.

[mravishankar]

Seems like this changed to rely on having a single execution mode op per function. I think that is true. Approving this for landing.

mlir/lib/Conversion/SPIRVToLLVM/ConvertSPIRVToLLVM.cpp

657

nit : add '_' between {0} and {1}

mlir/lib/Conversion/SPIRVToLLVM/ConvertSPIRVToLLVM.cpp

matchAndRewrite(SPIRVOp operation, ArrayRef<Value> operands,
     if (!dstType)
       return failure();
     rewriter.template replaceOpWithNewOp<LLVMOp>(operation, dstType, operands,
                                                  operation.getAttrs());
     return success();
   }
 };
 
+/// Converts `spv.ExecutionMode` into a global struct constant that holds
+/// execution mode information.
+class ExecutionModePattern
+    : public SPIRVToLLVMConversion<spirv::ExecutionModeOp> {
+public:
+  using SPIRVToLLVMConversion<spirv::ExecutionModeOp>::SPIRVToLLVMConversion;
+
+  LogicalResult
+  matchAndRewrite(spirv::ExecutionModeOp op, ArrayRef<Value> operands,
+                  ConversionPatternRewriter &rewriter) const override {
+    // First, create the global struct's name that would be associated with
+    // this entry point's execution mode. We set it to be:
+    //   __spv__{SPIR-V module name}_{function name}_execution_mode_info
+    ModuleOp module = op.getParentOfType<ModuleOp>();
+    std::string moduleName;
+    if (module.getName().hasValue())
+      moduleName = "_" + module.getName().getValue().str();
+    else
+      moduleName = "";
+    std::string executionModeInfoName = llvm::formatv(
+        "__spv_{0}_{1}_execution_mode_info", moduleName, op.fn().str());
+
+    MLIRContext *context = rewriter.getContext();
+    OpBuilder::InsertionGuard guard(rewriter);
+    rewriter.setInsertionPointToStart(module.getBody());
+
+    // Create a struct type, corresponding to the C struct below.
+    // struct {
+    //   int32_t executionMode;
+    //   int32_t values[];          // optional values
+    // };
+    auto llvmI32Type = LLVM::LLVMType::getInt32Ty(context);
+    SmallVector<LLVM::LLVMType, 2> fields;
+    fields.push_back(llvmI32Type);
+    ArrayAttr values = op.values();
+    if (!values.empty()) {
+      auto arrayType = LLVM::LLVMType::getArrayTy(llvmI32Type, values.size());
+      fields.push_back(arrayType);
+    }
+    auto structType = LLVM::LLVMType::getStructTy(context, fields);
+
+    // Create `llvm.mlir.global` with initializer region containing one block.
+    auto global = rewriter.create<LLVM::GlobalOp>(
+        UnknownLoc::get(context), structType, /*isConstant=*/true,
+        LLVM::Linkage::External, executionModeInfoName, Attribute());
+    Location loc = global.getLoc();
+    Region &region = global.getInitializerRegion();
+    Block *block = rewriter.createBlock(&region);
+
+    // Initialize the struct and set the execution mode value.
+    rewriter.setInsertionPoint(block, block->begin());
+    Value structValue = rewriter.create<LLVM::UndefOp>(loc, structType);
+    IntegerAttr executionModeAttr = op.execution_modeAttr();
+    Value executionMode =
+        rewriter.create<LLVM::ConstantOp>(loc, llvmI32Type, executionModeAttr);
+    structValue = rewriter.create<LLVM::InsertValueOp>(
+        loc, structType, structValue, executionMode,
+        ArrayAttr::get({rewriter.getIntegerAttr(rewriter.getI32Type(), 0)},
+                       context));
+
+    // Insert extra operands if they exist into execution mode info struct.
+    for (unsigned i = 0, e = values.size(); i < e; ++i) {
+      auto attr = values.getValue()[i];
+      Value entry = rewriter.create<LLVM::ConstantOp>(loc, llvmI32Type, attr);
+      structValue = rewriter.create<LLVM::InsertValueOp>(
+          loc, structType, structValue, entry,
+          ArrayAttr::get({rewriter.getIntegerAttr(rewriter.getI32Type(), 1),
+                          rewriter.getIntegerAttr(rewriter.getI32Type(), i)},
+                         context));
+    }
+    rewriter.create<LLVM::ReturnOp>(loc, ArrayRef<Value>({structValue}));
+    rewriter.eraseOp(op);
+    return success();
+  }
+};
+
 /// Converts `spv.globalVariable` to `llvm.mlir.global`. Note that SPIR-V global
 /// returns a pointer, whereas in LLVM dialect the global holds an actual value.
 /// This difference is handled by `spv._address_of` and `llvm.mlir.addressof`ops
 /// that both return a pointer.
 class GlobalVariablePattern
     : public SPIRVToLLVMConversion<spirv::GlobalVariableOp> {
 public:
   using SPIRVToLLVMConversion<spirv::GlobalVariableOp>::SPIRVToLLVMConversion;
patterns.insert<
       // Constant op
       ConstantScalarAndVectorPattern,
 
       // Control Flow ops
       BranchConversionPattern, BranchConditionalConversionPattern,
       FunctionCallPattern, LoopPattern, SelectionPattern,
       ErasePattern<spirv::MergeOp>,
 
-      // Entry points and execution mode
-      // Module generated from SPIR-V could have other "internal" functions, so
-      // having entry point and execution mode metadata can be useful. For now,
-      // simply remove them.
-      // TODO: Support EntryPoint/ExecutionMode properly.
-      ErasePattern<spirv::EntryPointOp>, ErasePattern<spirv::ExecutionModeOp>,
+      // Entry points and execution mode are handled separately.
+      ErasePattern<spirv::EntryPointOp>, ExecutionModePattern,
 
       // GLSL extended instruction set ops
       DirectConversionPattern<spirv::GLSLCeilOp, LLVM::FCeilOp>,
       DirectConversionPattern<spirv::GLSLCosOp, LLVM::CosOp>,
       DirectConversionPattern<spirv::GLSLExpOp, LLVM::ExpOp>,
       DirectConversionPattern<spirv::GLSLFAbsOp, LLVM::FAbsOp>,
       DirectConversionPattern<spirv::GLSLFloorOp, LLVM::FFloorOp>,
       DirectConversionPattern<spirv::GLSLFMaxOp, LLVM::MaxNumOp>,

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

spv.func @select_vector(%arg0: vector<2xi1>, %arg1: vector<2xi32>) "None" {
   spv.Return
 }
 
 //===----------------------------------------------------------------------===//
 // spv.EntryPoint and spv.ExecutionMode
 //===----------------------------------------------------------------------===//
 
 //      CHECK: module {
+// CHECK-NEXT:   llvm.mlir.global external constant @{{.*}}() : !llvm.struct<(i32)> {
+// CHECK-NEXT:     %[[UNDEF:.*]] = llvm.mlir.undef : !llvm.struct<(i32)>
+// CHECK-NEXT:     %[[VAL:.*]] = llvm.mlir.constant(31 : i32) : !llvm.i32
+// CHECK-NEXT:     %[[RET:.*]] = llvm.insertvalue %[[VAL]], %[[UNDEF]][0 : i32] : !llvm.struct<(i32)>
+// CHECK-NEXT:     llvm.return %[[RET]] : !llvm.struct<(i32)>
+// CHECK-NEXT:   }
 // CHECK-NEXT:   llvm.func @empty
 // CHECK-NEXT:     llvm.return
 // CHECK-NEXT:   }
 // CHECK-NEXT: }
-spv.module Logical GLSL450 {
+spv.module Logical OpenCL {
   spv.func @empty() "None" {
     spv.Return
   }
-  spv.EntryPoint "GLCompute" @empty
-  spv.ExecutionMode @empty "LocalSize", 1, 1, 1
+  spv.EntryPoint "Kernel" @empty
+  spv.ExecutionMode @empty "ContractionOff"
+}
+
+//      CHECK: module {
+// CHECK-NEXT:   llvm.mlir.global external constant @{{.*}}() : !llvm.struct<(i32, array<3 x i32>)> {
+// CHECK-NEXT:     %[[UNDEF:.*]] = llvm.mlir.undef : !llvm.struct<(i32, array<3 x i32>)>
+// CHECK-NEXT:     %[[EM:.*]] = llvm.mlir.constant(18 : i32) : !llvm.i32
+// CHECK-NEXT:     %[[T0:.*]] = llvm.insertvalue %[[EM]], %[[UNDEF]][0 : i32] : !llvm.struct<(i32, array<3 x i32>)>
+// CHECK-NEXT:     %[[C0:.*]] = llvm.mlir.constant(32 : i32) : !llvm.i32
+// CHECK-NEXT:     %[[T1:.*]] = llvm.insertvalue %[[C0]], %[[T0]][1 : i32, 0 : i32] : !llvm.struct<(i32, array<3 x i32>)>
+// CHECK-NEXT:     %[[C1:.*]] = llvm.mlir.constant(1 : i32) : !llvm.i32
+// CHECK-NEXT:     %[[T2:.*]] = llvm.insertvalue %[[C1]], %[[T1]][1 : i32, 1 : i32] : !llvm.struct<(i32, array<3 x i32>)>
+// CHECK-NEXT:     %[[C2:.*]] = llvm.mlir.constant(1 : i32) : !llvm.i32
+// CHECK-NEXT:     %[[RET:.*]] = llvm.insertvalue %[[C2]], %[[T2]][1 : i32, 2 : i32] : !llvm.struct<(i32, array<3 x i32>)>
+// CHECK-NEXT:     llvm.return %[[RET]] : !llvm.struct<(i32, array<3 x i32>)>
+// CHECK-NEXT:   }
+// CHECK-NEXT:   llvm.func @bar
+// CHECK-NEXT:     llvm.return
+// CHECK-NEXT:   }
+// CHECK-NEXT: }
+spv.module Logical OpenCL {
+  spv.func @bar() "None" {
+    spv.Return
+  }
+  spv.EntryPoint "Kernel" @bar
+  spv.ExecutionMode @bar "LocalSizeHint", 32, 1, 1
 }
 
 //===----------------------------------------------------------------------===//
 // spv.Undef
 //===----------------------------------------------------------------------===//
 
 // CHECK-LABEL: @undef_scalar
 spv.func @undef_scalar() "None" {