[Mlir-commits] [mlir] [mlir][SPIR-V] Guard getSizeInBytes() optionals in UnifyAliasedResourcePass rewriters (PR #207325)

[llvmlistbot at llvm.org]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir

Author: Arseniy Obolenskiy (aobolensk)

<details>
<summary>Changes</summary>



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


2 Files Affected:

- (modified) mlir/lib/Dialect/SPIRV/Transforms/UnifyAliasedResourcePass.cpp (+18-6) 
- (modified) mlir/test/Dialect/SPIRV/Transforms/unify-aliased-resource.mlir (+18) 


``````````diff
diff --git a/mlir/lib/Dialect/SPIRV/Transforms/UnifyAliasedResourcePass.cpp b/mlir/lib/Dialect/SPIRV/Transforms/UnifyAliasedResourcePass.cpp
index 692f2e7616e5a..ecfce68db036b 100644
--- a/mlir/lib/Dialect/SPIRV/Transforms/UnifyAliasedResourcePass.cpp
+++ b/mlir/lib/Dialect/SPIRV/Transforms/UnifyAliasedResourcePass.cpp
@@ -371,8 +371,12 @@ struct ConvertAccessChain : public ConvertAliasResource<spirv::AccessChainOp> {
       // them into a buffer with vector element types. We need to scale the last
       // index for the vector as a whole, then add one level of index for inside
       // the vector.
-      int srcNumBytes = *srcElemType.getSizeInBytes();
-      int dstNumBytes = *dstElemType.getSizeInBytes();
+      std::optional<int64_t> srcBytes = srcElemType.getSizeInBytes();
+      std::optional<int64_t> dstBytes = dstElemType.getSizeInBytes();
+      if (!srcBytes || !dstBytes)
+        return rewriter.notifyMatchFailure(acOp, "unknown element byte size");
+      int srcNumBytes = *srcBytes;
+      int dstNumBytes = *dstBytes;
       assert(dstNumBytes >= srcNumBytes && dstNumBytes % srcNumBytes == 0);
 
       auto indices = llvm::to_vector<4>(acOp.getIndices());
@@ -398,8 +402,12 @@ struct ConvertAccessChain : public ConvertAliasResource<spirv::AccessChainOp> {
       // The source indices are for a buffer with larger bitwidth scalar/vector
       // element types. Rewrite them into a buffer with smaller bitwidth element
       // types. We only need to scale the last index.
-      int srcNumBytes = *srcElemType.getSizeInBytes();
-      int dstNumBytes = *dstElemType.getSizeInBytes();
+      std::optional<int64_t> srcBytes = srcElemType.getSizeInBytes();
+      std::optional<int64_t> dstBytes = dstElemType.getSizeInBytes();
+      if (!srcBytes || !dstBytes)
+        return rewriter.notifyMatchFailure(acOp, "unknown element byte size");
+      int srcNumBytes = *srcBytes;
+      int dstNumBytes = *dstBytes;
       assert(srcNumBytes >= dstNumBytes && srcNumBytes % dstNumBytes == 0);
 
       auto indices = llvm::to_vector<4>(acOp.getIndices());
@@ -455,8 +463,12 @@ struct ConvertLoad : public ConvertAliasResource<spirv::LoadOp> {
       // vector types of different component counts. For such cases, we load
       // multiple smaller bitwidth values and construct a larger bitwidth one.
 
-      int srcNumBytes = *srcElemType.getSizeInBytes();
-      int dstNumBytes = *dstElemType.getSizeInBytes();
+      std::optional<int64_t> srcBytes = srcElemType.getSizeInBytes();
+      std::optional<int64_t> dstBytes = dstElemType.getSizeInBytes();
+      if (!srcBytes || !dstBytes)
+        return rewriter.notifyMatchFailure(loadOp, "unknown element byte size");
+      int srcNumBytes = *srcBytes;
+      int dstNumBytes = *dstBytes;
       assert(srcNumBytes > dstNumBytes && srcNumBytes % dstNumBytes == 0);
       int ratio = srcNumBytes / dstNumBytes;
       if (ratio > 4)
diff --git a/mlir/test/Dialect/SPIRV/Transforms/unify-aliased-resource.mlir b/mlir/test/Dialect/SPIRV/Transforms/unify-aliased-resource.mlir
index f5cd490c164d7..d7cc1d0cc1f84 100644
--- a/mlir/test/Dialect/SPIRV/Transforms/unify-aliased-resource.mlir
+++ b/mlir/test/Dialect/SPIRV/Transforms/unify-aliased-resource.mlir
@@ -549,6 +549,24 @@ spirv.module Logical GLSL450 {
 
 // -----
 
+// Make sure we do not crash on element types without a defined byte size (i1).
+
+spirv.module Logical GLSL450 {
+  spirv.GlobalVariable @var01_i1 bind(0, 1) {aliased} : !spirv.ptr<!spirv.struct<(!spirv.rtarray<i1, stride=1> [0])>, StorageBuffer>
+  spirv.GlobalVariable @var01_i32 bind(0, 1) {aliased} : !spirv.ptr<!spirv.struct<(!spirv.rtarray<i32, stride=4> [0])>, StorageBuffer>
+
+  spirv.func @unknown_byte_size(%index: i32) -> i32 "None" {
+    %c0 = spirv.Constant 0 : i32
+    %addr = spirv.mlir.addressof @var01_i32 : !spirv.ptr<!spirv.struct<(!spirv.rtarray<i32, stride=4> [0])>, StorageBuffer>
+    // expected-error at +1 {{failed to legalize operation 'spirv.AccessChain'}}
+    %ac = spirv.AccessChain %addr[%c0, %index] : !spirv.ptr<!spirv.struct<(!spirv.rtarray<i32, stride=4> [0])>, StorageBuffer>, i32, i32 -> !spirv.ptr<i32, StorageBuffer>
+    %value = spirv.Load "StorageBuffer" %ac : i32
+    spirv.ReturnValue %value : i32
+  }
+}
+
+// -----
+
 // Make sure we do not crash on function arguments.
 
 spirv.module Logical GLSL450 {

``````````

</details>


https://github.com/llvm/llvm-project/pull/207325