[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

736

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

739

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

733

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

mlir/lib/Conversion/SPIRVToLLVM/ConvertSPIRVToLLVM.cpp

 #include "mlir/Transforms/DialectConversion.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/FormatVariadic.h"
 
 #define DEBUG_TYPE "spirv-to-llvm-pattern"
 
 using namespace mlir;
 
+static constexpr const char kSPIRVModule[] = "__spv__";
+
 //===----------------------------------------------------------------------===//
 // 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;
static LogicalResult replaceWithLoadOrStore(Operation *op,
   }
   auto storeOp = cast<spirv::StoreOp>(op);
   rewriter.replaceOpWithNewOp<LLVM::StoreOp>(storeOp, storeOp.value(),
                                              storeOp.ptr(), alignment,
                                              isVolatile, isNonTemporal);
   return success();
 }
 
+/// Creates a name for the given entry point op or execution mode op and a given
+/// module. The name is of the following form:
+///   __spv__{module_name}_{fn_name}_entry_point_info
+/// This is used to avoid symbolic conflicts, as well as to encode entry-point
+/// function name.
+template <typename SPIRVOp>
+static std::string getEntryPointInfoName(SPIRVOp op, ModuleOp module) {
+  std::string moduleName;
+  if (module.getName().hasValue())
+    moduleName = kSPIRVModule + module.getName().getValue().str() + "_";
+  else
+    moduleName = kSPIRVModule;
+  return moduleName + op.fn().str() + "_entry_point_info";
+}
+
+/// Creates a global struct containing entry point information.
+static LLVM::GlobalOp createEntryPointInfo(Operation *op, std::string name,
+                                           ModuleOp module,
+                                           PatternRewriter &builder) {
+  MLIRContext *context = builder.getContext();
+  OpBuilder::InsertionGuard guard(builder);
+  builder.setInsertionPointToStart(module.getBody());
+
+  // Create a struct type, corresponding to the C struct below.
+  // struct {
+  //   int32_t executionModel;
+  //   int32_t executionMode;
+  // };
+  // TODO: support interface variables and execution mode values.
+  auto llvmI32Type = LLVM::LLVMType::getInt32Ty(context);
+  SmallVector<LLVM::LLVMType, 2> fields = {llvmI32Type, llvmI32Type};
+  auto structType = LLVM::LLVMType::getStructTy(context, fields);
+
+  // Create `llvm.mlir.global` with initializer region containing one block.
+  auto dstGlobal = builder.create<LLVM::GlobalOp>(
+      UnknownLoc::get(context), structType, /*isConstant=*/true,
+      LLVM::Linkage::External, name, Attribute());
+  Location loc = dstGlobal.getLoc();
+  Region &region = dstGlobal.getInitializerRegion();
+  Block *block = builder.createBlock(&region);
+
+  // Initialize the struct and add a terminator to the block.
+  builder.setInsertionPoint(block, block->begin());
+  Value structValue = builder.create<LLVM::UndefOp>(loc, structType);
+  builder.create<LLVM::ReturnOp>(loc, ArrayRef<Value>({structValue}));
+  return dstGlobal;
+}
+
+/// Fills entry point info struct, extracting information from the given op.
+static Value setEntryPointInfo(Operation *op, Location loc, Value pred,
+                               PatternRewriter &rewriter) {
+  MLIRContext *context = rewriter.getContext();
+  auto llvmI32Type = LLVM::LLVMType::getInt32Ty(context);
+
+  if (auto entryPointOp = dyn_cast<spirv::EntryPointOp>(op)) {
+    IntegerAttr executionModelAttr = entryPointOp.execution_modelAttr();
+    Value executionModel =
+        rewriter.create<LLVM::ConstantOp>(loc, llvmI32Type, executionModelAttr);
+    return rewriter.create<LLVM::InsertValueOp>(
+        loc, pred.getType(), pred, executionModel,
+        ArrayAttr::get({rewriter.getIntegerAttr(rewriter.getI32Type(), 0)},
+                       context));
+  }
+
+  auto executionModeOp = cast<spirv::ExecutionModeOp>(op);
+  IntegerAttr executionModeAttr = executionModeOp.execution_modeAttr();
+  Value executionMode =
+      rewriter.create<LLVM::ConstantOp>(loc, llvmI32Type, executionModeAttr);
+  return rewriter.create<LLVM::InsertValueOp>(
+      loc, pred.getType(), pred, executionMode,
+      ArrayAttr::get({rewriter.getIntegerAttr(rewriter.getI32Type(), 1)},
+                     context));
+}
+
 //===----------------------------------------------------------------------===//
 // Type conversion
 //===----------------------------------------------------------------------===//
 
 /// Converts SPIR-V array type to LLVM array. Natural stride (according to
 /// `VulkanLayoutUtils`) is also mapped to LLVM array. This has to be respected
 /// when converting ops that manipulate array types.
 static Optional<Type> convertArrayType(spirv::ArrayType type,
matchAndRewrite(SPIRVOp operation, ArrayRef<Value> operands,
     if (!dstType)
       return failure();
     rewriter.template replaceOpWithNewOp<LLVMOp>(operation, dstType, operands,
                                                  operation.getAttrs());
     return success();
   }
 };
 
+/// Converts `spv.EntryPoint` and `spv.ExecutionMode` into a global struct
+/// constant that holds all entry point information.
+template <typename SPIRVOp>
+class EntryPointPattern : public SPIRVToLLVMConversion<SPIRVOp> {
+public:
+  using SPIRVToLLVMConversion<SPIRVOp>::SPIRVToLLVMConversion;
+
+  LogicalResult
+  matchAndRewrite(SPIRVOp op, ArrayRef<Value> operands,
+                  ConversionPatternRewriter &rewriter) const override {
+    // First, check if the global struct has been created.
+    ModuleOp module = op.template getParentOfType<ModuleOp>();
+    auto entryPointInfoName = getEntryPointInfoName<SPIRVOp>(op, module);
+    auto entryPointInfo =
+        module.lookupSymbol<LLVM::GlobalOp>(entryPointInfoName);
+
+    // If entry point info struct has not been created, create it.
+    if (!entryPointInfo)
+      entryPointInfo =
+          createEntryPointInfo(op, entryPointInfoName, module, rewriter);
+
+    // Get last two operations (current struct value and the return operation)
+    // to insert in between them.
+    Block *block = &entryPointInfo.getInitializerRegion().front();
+    auto iter = Block::reverse_iterator(block->back());
+    Operation &insertBefore = *iter;
+    Operation &insertAfter = *std::next(iter);
+
+    OpBuilder::InsertionGuard guard(rewriter);
+    rewriter.setInsertionPoint(&insertBefore);
+    Location loc = insertBefore.getLoc();
+    Value previousStruct = insertAfter.getResult(0);
+    Value newStruct = setEntryPointInfo(op, loc, previousStruct, rewriter);
+
+    // Update initializer's block terminator with a new struct value.
+    Operation *terminator = block->getTerminator();
+    terminator->replaceUsesOfWith(previousStruct, newStruct);
+    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 metadat 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.
+      EntryPointPattern<spirv::EntryPointOp>,
+      EntryPointPattern<spirv::ExecutionModeOp>,
 
       // 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, i32)> {
+// CHECK-NEXT:     %[[UNDEF:.*]] = llvm.mlir.undef : !llvm.struct<(i32, i32)>
+// CHECK-NEXT:     %[[VAL1:.*]] = llvm.mlir.constant(6 : i32) : !llvm.i32
+// CHECK-NEXT:     %[[T0:.*]] = llvm.insertvalue %[[VAL1]], %[[UNDEF]][0 : i32] : !llvm.struct<(i32, i32)>
+// CHECK-NEXT:     %[[VAL2:.*]] = llvm.mlir.constant(31 : i32) : !llvm.i32
+// CHECK-NEXT:     %[[T1:.*]] = llvm.insertvalue %[[VAL2]], %[[T0]][1 : i32] : !llvm.struct<(i32, i32)>
+// CHECK-NEXT:     llvm.return %[[T1]] : !llvm.struct<(i32, 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"
 }
 
 //===----------------------------------------------------------------------===//
 // spv.Undef
 //===----------------------------------------------------------------------===//
 
 // CHECK-LABEL: @undef_scalar
 spv.func @undef_scalar() "None" {