From 31c35285d27ed507ae758aefdca0d9cd05c7f21d Mon Sep 17 00:00:00 2001
From: Jakub Kuderski <kubak@google.com>
Date: Tue, 14 Feb 2023 16:11:34 -0500
Subject: [PATCH] [mlir][spirv] Fix crash when decorating physical storage
 buffer pointers

Add a comment explaining `PhysicalStorageBufferAddresses` are not
supported yet.

Fixes: https://github.com/llvm/llvm-project/issues/60196

Reviewed By: antiagainst

Differential Revision: https://reviews.llvm.org/D144039
---
 mlir/include/mlir/Dialect/SPIRV/Utils/LayoutUtils.h              | 2 +-
 .../Dialect/SPIRV/Transforms/DecorateCompositeTypeLayoutPass.cpp | 9 ++++++---
 mlir/lib/Dialect/SPIRV/Utils/LayoutUtils.cpp                     | 4 ++++
 mlir/test/Dialect/SPIRV/Transforms/layout-decoration.mlir        | 9 +++++++++
 4 files changed, 20 insertions(+), 4 deletions(-)

diff --git a/mlir/include/mlir/Dialect/SPIRV/Utils/LayoutUtils.h b/mlir/include/mlir/Dialect/SPIRV/Utils/LayoutUtils.h
index 1f1fa79..0c61f7e 100644
--- a/mlir/include/mlir/Dialect/SPIRV/Utils/LayoutUtils.h
+++ b/mlir/include/mlir/Dialect/SPIRV/Utils/LayoutUtils.h
@@ -26,7 +26,7 @@ class RuntimeArrayType;
 class StructType;
 } // namespace spirv
 
-/// According to the Vulkan spec "14.5.4. Offset and Stride Assignment":
+/// According to the Vulkan spec "15.6.4. Offset and Stride Assignment":
 /// "There are different alignment requirements depending on the specific
 /// resources and on the features enabled on the device."
 ///
diff --git a/mlir/lib/Dialect/SPIRV/Transforms/DecorateCompositeTypeLayoutPass.cpp b/mlir/lib/Dialect/SPIRV/Transforms/DecorateCompositeTypeLayoutPass.cpp
index 42679c1..f154840 100644
--- a/mlir/lib/Dialect/SPIRV/Transforms/DecorateCompositeTypeLayoutPass.cpp
+++ b/mlir/lib/Dialect/SPIRV/Transforms/DecorateCompositeTypeLayoutPass.cpp
@@ -21,6 +21,8 @@
 #include "mlir/Dialect/SPIRV/Utils/LayoutUtils.h"
 #include "mlir/Transforms/DialectConversion.h"
 
+#include "llvm/Support/FormatVariadic.h"
+
 using namespace mlir;
 
 namespace mlir {
@@ -41,11 +43,12 @@ public:
     SmallVector<NamedAttribute, 4> globalVarAttrs;
 
     auto ptrType = op.getType().cast<spirv::PointerType>();
-    auto structType = VulkanLayoutUtils::decorateType(
-        ptrType.getPointeeType().cast<spirv::StructType>());
+    auto pointeeType = ptrType.getPointeeType().cast<spirv::StructType>();
+    spirv::StructType structType = VulkanLayoutUtils::decorateType(pointeeType);
 
     if (!structType)
-      return failure();
+      return op->emitError(llvm::formatv(
+          "failed to decorate (unsuported pointee type: '{0}')", pointeeType));
 
     auto decoratedType =
         spirv::PointerType::get(structType, ptrType.getStorageClass());
diff --git a/mlir/lib/Dialect/SPIRV/Utils/LayoutUtils.cpp b/mlir/lib/Dialect/SPIRV/Utils/LayoutUtils.cpp
index b5a38c3..67d61f8 100644
--- a/mlir/lib/Dialect/SPIRV/Utils/LayoutUtils.cpp
+++ b/mlir/lib/Dialect/SPIRV/Utils/LayoutUtils.cpp
@@ -95,6 +95,10 @@ Type VulkanLayoutUtils::decorateType(Type type, VulkanLayoutUtils::Size &size,
     size = std::numeric_limits<Size>().max();
     return decorateType(arrayType, alignment);
   }
+  if (type.isa<spirv::PointerType>()) {
+    // TODO: Add support for `PhysicalStorageBufferAddresses`.
+    return nullptr;
+  }
   llvm_unreachable("unhandled SPIR-V type");
 }
 
diff --git a/mlir/test/Dialect/SPIRV/Transforms/layout-decoration.mlir b/mlir/test/Dialect/SPIRV/Transforms/layout-decoration.mlir
index f6be114..d2c9f83 100644
--- a/mlir/test/Dialect/SPIRV/Transforms/layout-decoration.mlir
+++ b/mlir/test/Dialect/SPIRV/Transforms/layout-decoration.mlir
@@ -97,3 +97,12 @@ spirv.module Logical GLSL450 {
   // CHECK: spirv.GlobalVariable @var1 : !spirv.ptr<!spirv.struct<(i32 [0])>, PhysicalStorageBuffer>
   spirv.GlobalVariable @var1 : !spirv.ptr<!spirv.struct<(i32)>, PhysicalStorageBuffer>
 }
+
+// -----
+
+spirv.module Physical64 GLSL450 {
+  // expected-error @+2 {{failed to decorate (unsuported pointee type: '!spirv.struct<rec, (!spirv.ptr<!spirv.struct<rec>, StorageBuffer>)>')}}
+  // expected-error @+1 {{failed to legalize operation 'spirv.GlobalVariable'}}
+  spirv.GlobalVariable @recursive:
+    !spirv.ptr<!spirv.struct<rec, (!spirv.ptr<!spirv.struct<rec>, StorageBuffer>)>, StorageBuffer>
+}
-- 
2.7.4