[MLIR][SPIRVToLLVM] Added a hook for descriptor set / binding encoding

[georgemitenkov]

Summary

This patch introduces a hook to encode descriptor set
and binding number into spv.globalVariable's symbolic name. This
allows to preserve this information, and at the same time legalize
the global variable for the conversion to LLVM dialect.

This is required for mlir-spirv-cpu-runner to convert kernel
arguments into LLVM.

Also, a couple of some nits added:

• removed unused comment
• changed to a capital letter in the comment

[mravishankar]

mlir/lib/Conversion/SPIRVToLLVM/ConvertSPIRVToLLVM.cpp

1406

If you rebase this on top of the change to add SymbolName trait to spv.module you can add that to the name as well.

[georgemitenkov]

Rebased on top of the SPIR-V module name patch to include the symbol name of the module into the new name of the global variable.

mlir/include/mlir/Conversion/SPIRVToLLVM/ConvertSPIRVToLLVM.h

SPIRVToLLVMConversion(MLIRContext *context, LLVMTypeConverter &typeConverter,
                         PatternBenefit benefit = 1)
       : OpConversionPattern<SPIRVOp>(context, benefit),
         typeConverter(typeConverter) {}
 
 protected:
   LLVMTypeConverter &typeConverter;
 };
 
+/// Encodes global variable's descriptor set and binding into its name if they
+/// both exist.
+void encodeBindAttribute(ModuleOp module);
+
 /// Populates type conversions with additional SPIR-V types.
 void populateSPIRVToLLVMTypeConversion(LLVMTypeConverter &typeConverter);
 
 /// Populates the given list with patterns that convert from SPIR-V to LLVM.
 void populateSPIRVToLLVMConversionPatterns(MLIRContext *context,
                                            LLVMTypeConverter &typeConverter,
                                            OwningRewritePatternList &patterns);
 

mlir/lib/Conversion/SPIRVToLLVM/ConvertSPIRVToLLVM.cpp

 #include "mlir/Dialect/SPIRV/SPIRVDialect.h"
 #include "mlir/Dialect/SPIRV/SPIRVOps.h"
 #include "mlir/Dialect/StandardOps/IR/Ops.h"
 #include "mlir/IR/Module.h"
 #include "mlir/IR/PatternMatch.h"
 #include "mlir/Support/LogicalResult.h"
 #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;
 
 //===----------------------------------------------------------------------===//
 // Utility functions
 //===----------------------------------------------------------------------===//
patterns.insert<
 
       // 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>,
 
-      // Function Call op
-
       // 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>,
       DirectConversionPattern<spirv::GLSLFMinOp, LLVM::MinNumOp>,

 }
 
 void mlir::populateSPIRVToLLVMModuleConversionPatterns(
     MLIRContext *context, LLVMTypeConverter &typeConverter,
     OwningRewritePatternList &patterns) {
   patterns.insert<ModuleConversionPattern, ModuleEndConversionPattern>(
       context, typeConverter);
 }
+
+//===----------------------------------------------------------------------===//
+// Pre-conversion hooks
+//===----------------------------------------------------------------------===//
+
+/// Hook for descriptor set and binding number encoding.
+static constexpr StringRef kBinding = "binding";
+static constexpr StringRef kDescriptorSet = "descriptor_set";
+void mlir::encodeBindAttribute(ModuleOp module) {
+  auto spvModules = module.getOps<spirv::ModuleOp>();
+  for (auto spvModule : spvModules) {
+    spvModule.walk([&](spirv::GlobalVariableOp op) {
+      IntegerAttr descriptorSet = op.getAttrOfType<IntegerAttr>(kDescriptorSet);
+      IntegerAttr binding = op.getAttrOfType<IntegerAttr>(kBinding);
+      // For every global variable in the module, get the ones with descriptor
+      // set and binding numbers.
+      if (descriptorSet && binding) {
+        // Encode these numbers into the variable's symbolic name. If the
+        // SPIR-V module has a name, add it at the beginning.
+        auto moduleAndName = spvModule.getName().hasValue()
+                                 ? spvModule.getName().getValue().str() + "_" +
+                                       op.sym_name().str()
+                                 : op.sym_name().str();
+        std::string name =
+            llvm::formatv("{0}_descriptor_set{1}_binding{2}", moduleAndName,
+                          std::to_string(descriptorSet.getInt()),
+                          std::to_string(binding.getInt()));
+
+        // Replace all symbol uses and set the new symbol name. Finally, remove
+        // descriptor set and binding attributes.
+        if (failed(SymbolTable::replaceAllSymbolUses(op, name, spvModule)))
+          op.emitError("unable to replace all symbol uses for ") << name;
+        SymbolTable::setSymbolName(op, name);
+        op.removeAttr(kDescriptorSet);
+        op.removeAttr(kBinding);
+      }
+    });
+  }
+}

mlir/lib/Conversion/SPIRVToLLVM/ConvertSPIRVToLLVMPass.cpp

 };
 } // namespace
 
 void ConvertSPIRVToLLVMPass::runOnOperation() {
   MLIRContext *context = &getContext();
   ModuleOp module = getOperation();
   LLVMTypeConverter converter(&getContext());
 
+  // Encode global variable's descriptor set and binding if they exist.
+  encodeBindAttribute(module);
+
   OwningRewritePatternList patterns;
 
   populateSPIRVToLLVMTypeConversion(converter);
 
   populateSPIRVToLLVMModuleConversionPatterns(context, converter, patterns);
   populateSPIRVToLLVMConversionPatterns(context, converter, patterns);
   populateSPIRVToLLVMFunctionConversionPatterns(context, converter, patterns);
 
   ConversionTarget target(getContext());
   target.addIllegalDialect<spirv::SPIRVDialect>();
   target.addLegalDialect<LLVM::LLVMDialect>();
 
-  // set `ModuleOp` and `ModuleTerminatorOp` as legal for `spv.module`
+  // Set `ModuleOp` and `ModuleTerminatorOp` as legal for `spv.module`
   // conversion.
   target.addLegalOp<ModuleOp>();
   target.addLegalOp<ModuleTerminatorOp>();
   if (failed(applyPartialConversion(module, target, patterns)))
     signalPassFailure();
 }
 
 std::unique_ptr<OperationPass<ModuleOp>> mlir::createConvertSPIRVToLLVMPass() {
   return std::make_unique<ConvertSPIRVToLLVMPass>();
 }

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

 spv.module Logical GLSL450 {
   // CHECK: llvm.mlir.global external constant @var() : !llvm.float
   spv.globalVariable @var : !spv.ptr<f32, Input>
 }
 
 spv.module Logical GLSL450 {
   //       CHECK: llvm.mlir.global private @struct() : !llvm.struct<packed (float, array<10 x float>)>
   // CHECK-LABEL: @func
-  //       CHECK: llvm.mlir.addressof @struct : !llvm.ptr<struct<packed (float, array<10 x float>)>>
+  //       CHECK:   llvm.mlir.addressof @struct : !llvm.ptr<struct<packed (float, array<10 x float>)>>
   spv.globalVariable @struct : !spv.ptr<!spv.struct<f32, !spv.array<10xf32>>, Private>
   spv.func @func() "None" {
     %0 = spv._address_of @struct : !spv.ptr<!spv.struct<f32, !spv.array<10xf32>>, Private>
     spv.Return
   }
 }
 
+spv.module Logical GLSL450 {
+  //       CHECK: llvm.mlir.global external @bar_descriptor_set0_binding0() : !llvm.i32
+  // CHECK-LABEL: @foo
+  //       CHECK:   llvm.mlir.addressof @bar_descriptor_set0_binding0 : !llvm.ptr<i32>
+  spv.globalVariable @bar bind(0, 0) : !spv.ptr<i32, StorageBuffer>
+  spv.func @foo() "None" {
+    %0 = spv._address_of @bar : !spv.ptr<i32, StorageBuffer>
+    spv.Return
+  }
+}
+
+spv.module @name Logical GLSL450 {
+  //       CHECK: llvm.mlir.global external @name_bar_descriptor_set0_binding0() : !llvm.i32
+  // CHECK-LABEL: @foo
+  //       CHECK:   llvm.mlir.addressof @name_bar_descriptor_set0_binding0 : !llvm.ptr<i32>
+  spv.globalVariable @bar bind(0, 0) : !spv.ptr<i32, StorageBuffer>
+  spv.func @foo() "None" {
+    %0 = spv._address_of @bar : !spv.ptr<i32, StorageBuffer>
+    spv.Return
+  }
+}
+
 //===----------------------------------------------------------------------===//
 // spv.Load
 //===----------------------------------------------------------------------===//
 
 // CHECK-LABEL: @load
 spv.func @load() "None" {
   %0 = spv.Variable : !spv.ptr<f32, Function>
   //  CHECK: llvm.load %{{.*}} : !llvm.ptr<float>