[MLIR][SPIRVToLLVM] Implementation of SPIR-V module conversion pattern

[georgemitenkov]

Summary

This patch introduces conversion patterns for spv.module and spv._module_end. SPIR-V module is converted into ModuleOp. This will play a role of enclosing scope to LLVM ops. At the moment, SPIR-V module attributes (such as memory model, etc) are ignored.

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

void populateSPIRVToLLVMConversionPatterns(MLIRContext *context,
                                            OwningRewritePatternList &patterns);
 
 /// Populates the given list with patterns for function conversion from SPIR-V
 /// to LLVM.
 void populateSPIRVToLLVMFunctionConversionPatterns(
     MLIRContext *context, LLVMTypeConverter &typeConverter,
     OwningRewritePatternList &patterns);
 
+/// Populates the given patterns for module conversion from SPIR-V to LLVM.
+void populateSPIRVToLLVMModuleConversionPatterns(
+    MLIRContext *context, LLVMTypeConverter &typeConverter,
+    OwningRewritePatternList &patterns);
+
 } // namespace mlir
 
 #endif // MLIR_CONVERSION_SPIRVTOLLVM_CONVERTSPIRVTOLLVM_H

mlir/lib/Conversion/SPIRVToLLVM/ConvertSPIRVToLLVM.cpp

#undef DISPATCH
 
     rewriter.inlineRegionBefore(funcOp.getBody(), newFuncOp.getBody(),
                                 newFuncOp.end());
     rewriter.applySignatureConversion(&newFuncOp.getBody(), signatureConverter);
     rewriter.eraseOp(funcOp);
     return success();
   }
 };
+
+//===----------------------------------------------------------------------===//
+// ModuleOp conversion
+//===----------------------------------------------------------------------===//
+
+class ModuleConversionPattern : public SPIRVToLLVMConversion<spirv::ModuleOp> {
+public:
+  using SPIRVToLLVMConversion<spirv::ModuleOp>::SPIRVToLLVMConversion;
+
+  LogicalResult
+  matchAndRewrite(spirv::ModuleOp spvModuleOp, ArrayRef<Value> operands,
+                  ConversionPatternRewriter &rewriter) const override {
+
+    auto newModuleOp = rewriter.create<ModuleOp>(spvModuleOp.getLoc());
+    rewriter.inlineRegionBefore(spvModuleOp.body(), newModuleOp.getBody());
+
+    // Remove the terminator block that was automatically added by builder
+    rewriter.eraseBlock(&newModuleOp.getBodyRegion().back());
+    rewriter.eraseOp(spvModuleOp);
+    return success();
+  }
+};
+
+class ModuleEndConversionPattern
+    : public SPIRVToLLVMConversion<spirv::ModuleEndOp> {
+public:
+  using SPIRVToLLVMConversion<spirv::ModuleEndOp>::SPIRVToLLVMConversion;
+
+  LogicalResult
+  matchAndRewrite(spirv::ModuleEndOp moduleEndOp, ArrayRef<Value> operands,
+                  ConversionPatternRewriter &rewriter) const override {
+
+    rewriter.replaceOpWithNewOp<ModuleTerminatorOp>(moduleEndOp);
+    return success();
+  }
+};
+
 } // namespace
 
 //===----------------------------------------------------------------------===//
 // Pattern population
 //===----------------------------------------------------------------------===//
 
 void mlir::populateSPIRVToLLVMConversionPatterns(
     MLIRContext *context, LLVMTypeConverter &typeConverter,
patterns.insert<
       ReturnPattern, ReturnValuePattern>(context, typeConverter);
 }
 
 void mlir::populateSPIRVToLLVMFunctionConversionPatterns(
     MLIRContext *context, LLVMTypeConverter &typeConverter,
     OwningRewritePatternList &patterns) {
   patterns.insert<FuncConversionPattern>(context, typeConverter);
 }
+
+void mlir::populateSPIRVToLLVMModuleConversionPatterns(
+    MLIRContext *context, LLVMTypeConverter &typeConverter,
+    OwningRewritePatternList &patterns) {
+  patterns.insert<ModuleConversionPattern, ModuleEndConversionPattern>(
+      context, typeConverter);
+}

mlir/lib/Conversion/SPIRVToLLVM/ConvertSPIRVToLLVMPass.cpp

 } // namespace
 
 void ConvertSPIRVToLLVMPass::runOnOperation() {
   MLIRContext *context = &getContext();
   ModuleOp module = getOperation();
   LLVMTypeConverter converter(&getContext());
 
   OwningRewritePatternList patterns;
+  populateSPIRVToLLVMModuleConversionPatterns(context, converter, patterns);
   populateSPIRVToLLVMConversionPatterns(context, converter, patterns);
   populateSPIRVToLLVMFunctionConversionPatterns(context, converter, patterns);
 
   // Currently pulls in Std to LLVM conversion patterns
   // that help with testing. This allows to convert
   // function arguments to LLVM.
   populateStdToLLVMConversionPatterns(converter, patterns);
 
   ConversionTarget target(getContext());
   target.addIllegalDialect<spirv::SPIRVDialect>();
   target.addLegalDialect<LLVM::LLVMDialect>();
+
+  // 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/module-ops-to-llvm.mlir

+// RUN: mlir-opt -convert-spirv-to-llvm %s | FileCheck %s
+
+//===----------------------------------------------------------------------===//
+// spv.module
+//===----------------------------------------------------------------------===//
+
+// CHECK: module
+spv.module Logical GLSL450 {}
+
+// CHECK: module
+spv.module Logical GLSL450 requires #spv.vce<v1.0, [Shader], [SPV_KHR_16bit_storage]> {}
+
+// CHECK: module
+spv.module Logical GLSL450 {
+  // CHECK: }
+  spv._module_end
+}
+
+// CHECK: module
+spv.module Logical GLSL450 {
+  // CHECK-LABEL: llvm.func @empty()
+  spv.func @empty() -> () "None" {
+    // CHECK: llvm.return
+    spv.Return
+  }
+}