[mlir][spirv] Implement missing validation rules for ptr variables

[kuhar]

Variables that point to physical storage buffer require aliasing decorations. This is specified by the SPV_KHR_physical_storage_buffer extension.

Also add an example of a variable with a decoration attribute.

[llvmbot]

@llvm/pr-subscribers-mlir

@llvm/pr-subscribers-mlir-spirv

Changes

Variables that point to physical storage buffer require aliasing decorations. This is specified by the SPV_KHR_physical_storage_buffer extension.

Also add an example of a variable with a decoration attribute.

---

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

3 Files Affected:

• (modified) mlir/include/mlir/Dialect/SPIRV/IR/SPIRVMemoryOps.td (+21-2)
• (modified) mlir/lib/Dialect/SPIRV/IR/MemoryOps.cpp (+45-11)
• (modified) mlir/test/Dialect/SPIRV/IR/memory-ops.mlir (+55)

diff --git a/mlir/include/mlir/Dialect/SPIRV/IR/SPIRVMemoryOps.td b/mlir/include/mlir/Dialect/SPIRV/IR/SPIRVMemoryOps.td
index 41cb4819ea389e5..291f2ef055c8a0c 100644
--- a/mlir/include/mlir/Dialect/SPIRV/IR/SPIRVMemoryOps.td
+++ b/mlir/include/mlir/Dialect/SPIRV/IR/SPIRVMemoryOps.td
@@ -289,7 +289,9 @@ def SPIRV_PtrAccessChainOp : SPIRV_Op<"PtrAccessChain", [Pure]> {
     MinVersion<SPIRV_V_1_0>,
     MaxVersion<SPIRV_V_1_6>,
     Extension<[]>,
-    Capability<[SPIRV_C_Addresses, SPIRV_C_PhysicalStorageBufferAddresses, SPIRV_C_VariablePointers, SPIRV_C_VariablePointersStorageBuffer]>
+    Capability<[
+      SPIRV_C_Addresses, SPIRV_C_PhysicalStorageBufferAddresses,
+      SPIRV_C_VariablePointers, SPIRV_C_VariablePointersStorageBuffer]>
   ];
 
   let arguments = (ins
@@ -375,11 +377,16 @@ def SPIRV_VariableOp : SPIRV_Op<"Variable", []> {
     must be the `Function` Storage Class.
 
     Initializer is optional. If Initializer is present, it will be the
-    initial value of the variable’s memory content. Initializer must be an
+    initial value of the variable's memory content. Initializer must be an
     <id> from a constant instruction or a global (module scope) OpVariable
     instruction. Initializer must have the same type as the type pointed to
     by Result Type.
 
+    From `SPV_KHR_physical_storage_buffer`:
+    If an OpVariable's pointee type is a pointer (or array of pointers) in
+    PhysicalStorageBuffer storage class, then the variable must be decorated
+    with exactly one of AliasedPointer or RestrictPointer.
+
     <!-- End of AutoGen section -->
 
     ```
@@ -396,6 +403,9 @@ def SPIRV_VariableOp : SPIRV_Op<"Variable", []> {
 
     %1 = spirv.Variable : !spirv.ptr<f32, Function>
     %2 = spirv.Variable init(%0): !spirv.ptr<f32, Function>
+
+    %3 = spirv.Variable {aliased_pointer} :
+      !spirv.ptr<!spirv.ptr<f32, PhysicalStorageBuffer>, Function>
     ```
   }];
 
@@ -407,6 +417,15 @@ def SPIRV_VariableOp : SPIRV_Op<"Variable", []> {
   let results = (outs
     SPIRV_AnyPtr:$pointer
   );
+
+  let extraClassDeclaration = [{
+    ::mlir::spirv::PointerType getPointerType() {
+      return ::llvm::cast<::mlir::spirv::PointerType>(getType());
+    }
+    ::mlir::Type getPointeeType() {
+      return getPointerType().getPointeeType();
+    }
+  }];
 }
 
 #endif // MLIR_DIALECT_SPIRV_IR_MEMORY_OPS
diff --git a/mlir/lib/Dialect/SPIRV/IR/MemoryOps.cpp b/mlir/lib/Dialect/SPIRV/IR/MemoryOps.cpp
index 6ee162a761e21e2..0df59e42218b552 100644
--- a/mlir/lib/Dialect/SPIRV/IR/MemoryOps.cpp
+++ b/mlir/lib/Dialect/SPIRV/IR/MemoryOps.cpp
@@ -10,12 +10,16 @@
 //
 //===----------------------------------------------------------------------===//
 
+#include "mlir/Dialect/SPIRV/IR/SPIRVEnums.h"
 #include "mlir/Dialect/SPIRV/IR/SPIRVOps.h"
 
 #include "SPIRVOpUtils.h"
 #include "SPIRVParsingUtils.h"
+#include "mlir/Dialect/SPIRV/IR/SPIRVTypes.h"
+#include "mlir/IR/Diagnostics.h"
 
 #include "llvm/ADT/StringExtras.h"
+#include "llvm/Support/Casting.h"
 
 using namespace mlir::spirv::AttrNames;
 
@@ -730,19 +734,49 @@ LogicalResult VariableOp::verify() {
                          "constant or spirv.GlobalVariable op");
   }
 
+  auto getDecorationAttr = [op = getOperation()](spirv::Decoration decoration) {
+    return op->getAttr(
+        llvm::convertToSnakeFromCamelCase(stringifyDecoration(decoration)));
+  };
+
   // TODO: generate these strings using ODS.
-  auto *op = getOperation();
-  auto descriptorSetName = llvm::convertToSnakeFromCamelCase(
-      stringifyDecoration(spirv::Decoration::DescriptorSet));
-  auto bindingName = llvm::convertToSnakeFromCamelCase(
-      stringifyDecoration(spirv::Decoration::Binding));
-  auto builtInName = llvm::convertToSnakeFromCamelCase(
-      stringifyDecoration(spirv::Decoration::BuiltIn));
-
-  for (const auto &attr : {descriptorSetName, bindingName, builtInName}) {
-    if (op->getAttr(attr))
+  for (auto decoration :
+       {spirv::Decoration::DescriptorSet, spirv::Decoration::Binding,
+        spirv::Decoration::BuiltIn}) {
+    if (auto attr = getDecorationAttr(decoration))
       return emitOpError("cannot have '")
-             << attr << "' attribute (only allowed in spirv.GlobalVariable)";
+             << llvm::convertToSnakeFromCamelCase(
+                    stringifyDecoration(decoration))
+             << "' attribute (only allowed in spirv.GlobalVariable)";
+  }
+
+  // From SPV_KHR_physical_storage_buffer:
+  // > If an OpVariable's pointee type is a pointer (or array of pointers) in
+  // > PhysicalStorageBuffer storage class, then the variable must be decorated
+  // > with exactly one of AliasedPointer or RestrictPointer.
+  auto pointeePtrType = dyn_cast<spirv::PointerType>(getPointeeType());
+  if (!pointeePtrType) {
+    if (auto pointeeArrayType = dyn_cast<spirv::ArrayType>(getPointeeType())) {
+      pointeePtrType =
+          dyn_cast<spirv::PointerType>(pointeeArrayType.getElementType());
+    }
+  }
+
+  if (pointeePtrType && pointeePtrType.getStorageClass() ==
+                            spirv::StorageClass::PhysicalStorageBuffer) {
+    bool hasAliasedPtr =
+        getDecorationAttr(spirv::Decoration::AliasedPointer) != nullptr;
+    bool hasRestrictPtr =
+        getDecorationAttr(spirv::Decoration::RestrictPointer) != nullptr;
+
+    if (!hasAliasedPtr && !hasRestrictPtr)
+      return emitOpError() << " with physical buffer pointer must be decorated "
+                              "either 'AliasedPointer' or 'RestrictPointer'";
+
+    if (hasAliasedPtr && hasRestrictPtr)
+      return emitOpError()
+             << " with physical buffer pointer must have exactly one "
+                "aliasing decoration";
   }
 
   return success();
diff --git a/mlir/test/Dialect/SPIRV/IR/memory-ops.mlir b/mlir/test/Dialect/SPIRV/IR/memory-ops.mlir
index 60af49262b7c9ee..0cae7f3b8ade5e8 100644
--- a/mlir/test/Dialect/SPIRV/IR/memory-ops.mlir
+++ b/mlir/test/Dialect/SPIRV/IR/memory-ops.mlir
@@ -525,6 +525,61 @@ spirv.module Logical GLSL450 {
 
 // -----
 
+func.func @variable_ptr_physical_buffer() -> () {
+  %0 = spirv.Variable {aliased_pointer} :
+    !spirv.ptr<!spirv.ptr<f32, PhysicalStorageBuffer>, Function>
+  %1 = spirv.Variable {restrict_pointer} :
+    !spirv.ptr<!spirv.ptr<f32, PhysicalStorageBuffer>, Function>
+  return
+}
+
+// -----
+
+func.func @variable_ptr_physical_buffer_no_decoration() -> () {
+  // expected-error @+1 {{must be decorated either 'AliasedPointer' or 'RestrictPointer'}}
+  %0 = spirv.Variable : !spirv.ptr<!spirv.ptr<f32, PhysicalStorageBuffer>, Function>
+  return
+}
+
+// -----
+
+func.func @variable_ptr_physical_buffer_two_alias_decorations() -> () {
+  // expected-error @+1 {{must have exactly one aliasing decoration}}
+  %0 = spirv.Variable {aliased_pointer, restrict_pointer} :
+    !spirv.ptr<!spirv.ptr<f32, PhysicalStorageBuffer>, Function>
+  return
+}
+
+// -----
+
+func.func @variable_ptr_array_physical_buffer() -> () {
+  %0 = spirv.Variable {aliased_pointer} :
+    !spirv.ptr<!spirv.array<4x!spirv.ptr<f32, PhysicalStorageBuffer>>, Function>
+  %1 = spirv.Variable {restrict_pointer} :
+    !spirv.ptr<!spirv.array<4x!spirv.ptr<f32, PhysicalStorageBuffer>>, Function>
+  return
+}
+
+// -----
+
+func.func @variable_ptr_physical_buffer_no_decoration() -> () {
+  // expected-error @+1 {{must be decorated either 'AliasedPointer' or 'RestrictPointer'}}
+  %0 = spirv.Variable :
+    !spirv.ptr<!spirv.array<4x!spirv.ptr<f32, PhysicalStorageBuffer>>, Function>
+  return
+}
+
+// -----
+
+func.func @variable_ptr_physical_buffer_two_alias_decorations() -> () {
+  // expected-error @+1 {{must have exactly one aliasing decoration}}
+  %0 = spirv.Variable {aliased_pointer, restrict_pointer} :
+    !spirv.ptr<!spirv.array<4x!spirv.ptr<f32, PhysicalStorageBuffer>>, Function>
+  return
+}
+
+// -----
+
 func.func @variable_bind() -> () {
   // expected-error @+1 {{cannot have 'descriptor_set' attribute (only allowed in spirv.GlobalVariable)}}
   %0 = spirv.Variable bind(1, 2) : !spirv.ptr<f32, Function>