[Mlir-commits] [mlir] ffd3a19 - Fix crash in Spirv -lower-host-to-llvm pass

[Mehdi Amini]

Author: Mehdi Amini
Date: 2023-01-16T21:06:49Z
New Revision: ffd3a19e653f9356ec0be65d4e1c38562573f6c7

URL: https://github.com/llvm/llvm-project/commit/ffd3a19e653f9356ec0be65d4e1c38562573f6c7
DIFF: https://github.com/llvm/llvm-project/commit/ffd3a19e653f9356ec0be65d4e1c38562573f6c7.diff

LOG: Fix crash in Spirv -lower-host-to-llvm pass

When providing with a spirv module as input where no conversion happens
the code didn't defend against broken invariant.

We'll fail the pass here, but it's not clear if it is the right thing
or if the module should just be ignored.

Fixes #59971

Reviewed By: kuhar

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

Added: 
    mlir/test/Conversion/SPIRVToLLVM/lower-host-to-llvm-calls_fail.mlir

Modified: 
    mlir/lib/Conversion/SPIRVToLLVM/ConvertLaunchFuncToLLVMCalls.cpp

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Conversion/SPIRVToLLVM/ConvertLaunchFuncToLLVMCalls.cpp b/mlir/lib/Conversion/SPIRVToLLVM/ConvertLaunchFuncToLLVMCalls.cpp
index a58072c0d9d67..2017d553b3788 100644
--- a/mlir/lib/Conversion/SPIRVToLLVM/ConvertLaunchFuncToLLVMCalls.cpp
+++ b/mlir/lib/Conversion/SPIRVToLLVM/ConvertLaunchFuncToLLVMCalls.cpp
@@ -131,7 +131,12 @@ static LogicalResult encodeKernelName(spirv::ModuleOp module) {
   // based on `getKernelGlobalVariables()` call. Update this function's name
   // to:
   //   {spv_module_name}_{function_name}
-  auto entryPoint = *module.getOps<spirv::EntryPointOp>().begin();
+  auto entryPoints = module.getOps<spirv::EntryPointOp>();
+  if (!llvm::hasSingleElement(entryPoints)) {
+    return module.emitError(
+        "The module must contain exactly one entry point function");
+  }
+  spirv::EntryPointOp entryPoint = *entryPoints.begin();
   StringRef funcName = entryPoint.getFn();
   auto funcOp = module.lookupSymbol<spirv::FuncOp>(entryPoint.getFnAttr());
   StringAttr newFuncName =
@@ -313,8 +318,12 @@ class LowerHostCodeToLLVM
 
     // Finally, modify the kernel function in SPIR-V modules to avoid symbolic
     // conflicts.
-    for (auto spvModule : module.getOps<spirv::ModuleOp>())
-      (void)encodeKernelName(spvModule);
+    for (auto spvModule : module.getOps<spirv::ModuleOp>()) {
+      if (failed(encodeKernelName(spvModule))) {
+        signalPassFailure();
+        return;
+      }
+    }
   }
 };
 } // namespace

diff  --git a/mlir/test/Conversion/SPIRVToLLVM/lower-host-to-llvm-calls_fail.mlir b/mlir/test/Conversion/SPIRVToLLVM/lower-host-to-llvm-calls_fail.mlir
new file mode 100644
index 0000000000000..e36d30b434fc7
--- /dev/null
+++ b/mlir/test/Conversion/SPIRVToLLVM/lower-host-to-llvm-calls_fail.mlir
@@ -0,0 +1,7 @@
+// RUN: mlir-opt --lower-host-to-llvm %s -verify-diagnostics
+
+module {
+// expected-error @+1 {{The module must contain exactly one entry point function}}
+  spirv.module Logical GLSL450 {
+  }
+}