[mlir][spirv] Handle failed conversions of struct elements

[Pierre-vh]

LLVMStructTypes could be emitted with some null elements. This caused a crash later in the LLVMDialect verifier.

Now, properly check that all struct elements were successfully converted before passing them to the LLVMStructType ctor.

See #59990

[llvmbot]

@llvm/pr-subscribers-mlir

@llvm/pr-subscribers-mlir-spirv

Author: Pierre van Houtryve (Pierre-vh)

Changes

LLVMStructTypes could be emitted with some null elements. This caused a crash later in the LLVMDialect verifier.

Now, properly check that all struct elements were successfully converted before passing them to the LLVMStructType ctor.

See #59990

---

Full diff: https://github.com/llvm/llvm-project/pull/70005.diff

2 Files Affected:

• (modified) mlir/lib/Conversion/SPIRVToLLVM/SPIRVToLLVM.cpp (+17-9)
• (modified) mlir/test/Conversion/SPIRVToLLVM/spirv-types-to-llvm-invalid.mlir (+7)

diff --git a/mlir/lib/Conversion/SPIRVToLLVM/SPIRVToLLVM.cpp b/mlir/lib/Conversion/SPIRVToLLVM/SPIRVToLLVM.cpp
index 60f34f413f587d4..5f752765f6d7f20 100644
--- a/mlir/lib/Conversion/SPIRVToLLVM/SPIRVToLLVM.cpp
+++ b/mlir/lib/Conversion/SPIRVToLLVM/SPIRVToLLVM.cpp
@@ -199,6 +199,16 @@ static Value processCountOrOffset(Location loc, Value value, Type srcType,
   return optionallyTruncateOrExtend(loc, broadcasted, dstType, rewriter);
 }
 
+static bool convertTypes(LLVMTypeConverter &converter, const spirv::StructType::ElementTypeRange &types, SmallVectorImpl<Type> &out) {
+  for(const auto &type: types) {
+    if(auto convertedType = converter.convertType(type))
+      out.push_back(convertedType);
+    else
+      return false;
+  }
+  return true;
+}
+
 /// Converts SPIR-V struct with a regular (according to `VulkanLayoutUtils`)
 /// offset to LLVM struct. Otherwise, the conversion is not supported.
 static std::optional<Type>
@@ -207,21 +217,19 @@ convertStructTypeWithOffset(spirv::StructType type,
   if (type != VulkanLayoutUtils::decorateType(type))
     return std::nullopt;
 
-  auto elementsVector = llvm::to_vector<8>(
-      llvm::map_range(type.getElementTypes(), [&](Type elementType) {
-        return converter.convertType(elementType);
-      }));
+  SmallVector<Type> elementsVector;
+  if(!convertTypes(converter, type.getElementTypes(), elementsVector))
+    return std::nullopt;
   return LLVM::LLVMStructType::getLiteral(type.getContext(), elementsVector,
                                           /*isPacked=*/false);
 }
 
 /// Converts SPIR-V struct with no offset to packed LLVM struct.
-static Type convertStructTypePacked(spirv::StructType type,
+static std::optional<Type> convertStructTypePacked(spirv::StructType type,
                                     LLVMTypeConverter &converter) {
-  auto elementsVector = llvm::to_vector<8>(
-      llvm::map_range(type.getElementTypes(), [&](Type elementType) {
-        return converter.convertType(elementType);
-      }));
+  SmallVector<Type> elementsVector;
+  if(!convertTypes(converter, type.getElementTypes(), elementsVector))
+    return std::nullopt;
   return LLVM::LLVMStructType::getLiteral(type.getContext(), elementsVector,
                                           /*isPacked=*/true);
 }
diff --git a/mlir/test/Conversion/SPIRVToLLVM/spirv-types-to-llvm-invalid.mlir b/mlir/test/Conversion/SPIRVToLLVM/spirv-types-to-llvm-invalid.mlir
index 3965c47ec199fcb..438c90205abedc4 100644
--- a/mlir/test/Conversion/SPIRVToLLVM/spirv-types-to-llvm-invalid.mlir
+++ b/mlir/test/Conversion/SPIRVToLLVM/spirv-types-to-llvm-invalid.mlir
@@ -7,6 +7,13 @@ spirv.func @array_with_unnatural_stride(%arg: !spirv.array<4 x f32, stride=8>) -
 
 // -----
 
+// expected-error@+1 {{failed to legalize operation 'spirv.func' that was explicitly marked illegal}}
+spirv.func @struct_array_with_unnatural_stride(%arg: !spirv.struct<(!spirv.array<4 x f32, stride=8>)>) -> () "None" {
+  spirv.Return
+}
+
+// -----
+
 // expected-error@+1 {{failed to legalize operation 'spirv.func' that was explicitly marked illegal}}
 spirv.func @struct_with_unnatural_offset(%arg: !spirv.struct<(i32[0], i32[8])>) -> () "None" {
   spirv.Return

[github-actions]

✅ With the latest revision this PR passed the C/C++ code formatter.

[victor-eds]

This is a nice catch. I see convertArrayType, convertPointerType and convertRuntimeArrayType have the same issue. Can you please also address in this PR?

[kuhar]

+1 for using TypeConverter

[victor-eds]

Change LGTM. Also checking element types for other compound types (ArrayType, PointerType and RuntimeArrayType) would be a plus.

[Pierre-vh]

Change LGTM. Also checking element types for other compound types (ArrayType, PointerType and RuntimeArrayType) would be a plus.

I'm not too familiar with SPIRV semantics so I just want to focus on avoiding the crash in this patch.
@kuhar is this good to land?

[kuhar]

Looks good % the failure checks

[kuhar]

LGTM!

[antiagainst]

Nice, thanks for the fix!