[MLIR][SPIRVToLLVM] Enhanced conversion for execution mode

This patch introduces a new conversion pattern for `spv.ExecutionMode`.
`spv.ExecutionMode` may contain important information about the entry
point, which we want to preserve. For example, `LocalSize` provides
information about the work-group size that can be reused. Hence, the
pattern for entry-point ops changes to the following:
- `spv.EntryPoint` is still simply removed
- Info from `spv.ExecutionMode` is used to create a global struct variable,
  which looks like:

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

Reviewed By: mravishankar

Differential Revision: https://reviews.llvm.org/D89989

mlir/lib/Conversion/SPIRVToLLVM/ConvertSPIRVToLLVM.cpp

@@ -634,6 +634,82 @@ public:
   }
 };
 
+/// 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
@@ -1386,12 +1462,8 @@ void mlir::populateSPIRVToLLVMConversionPatterns(
       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>,

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

@@ -63,16 +63,47 @@ spv.func @select_vector(%arg0: vector<2xi1>, %arg1: vector<2xi32>) "None" {
 //===----------------------------------------------------------------------===//
 
 //      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
 }
 
 //===----------------------------------------------------------------------===//