[Mlir-commits] [mlir] [mlir][spirv] Support function argument decorations for ptr in the PhysicalStorageBuffer (PR #76353)

[Jakub Kuderski]

================
@@ -369,6 +370,46 @@ LogicalResult spirv::Deserializer::processMemberName(ArrayRef<uint32_t> words) {
   return success();
 }
 
+void spirv::Deserializer::setArgAttrs(uint32_t argID) {
+  if (!decorations.count(argID)) {
+    argAttrs.push_back(DictionaryAttr::get(context, {}));
+    return;
+  }
+
+  // Replace a decoration as UnitAttr with DecorationAttr for the physical
+  // buffer pointer in the function parameter.
+  // e.g. "aliased" -> "spirv.decoration = #spirv.decoration<Aliased>").
+  for (auto decAttr : decorations[argID]) {
+    if (decAttr.getName() ==
+        getSymbolDecoration(stringifyDecoration(spirv::Decoration::Aliased))) {
+      decorations[argID].erase(decAttr.getName());
+      decorations[argID].set(
+          spirv::DecorationAttr::name,
+          spirv::DecorationAttr::get(context, spirv::Decoration::Aliased));
+    } else if (decAttr.getName() == getSymbolDecoration(stringifyDecoration(
+                                        spirv::Decoration::Restrict))) {
+      decorations[argID].erase(decAttr.getName());
+      decorations[argID].set(
+          spirv::DecorationAttr::name,
+          spirv::DecorationAttr::get(context, spirv::Decoration::Restrict));
+    } else if (decAttr.getName() == getSymbolDecoration(stringifyDecoration(
+                                        spirv::Decoration::AliasedPointer))) {
+      decorations[argID].erase(decAttr.getName());
+      decorations[argID].set(spirv::DecorationAttr::name,
+                             spirv::DecorationAttr::get(
+                                 context, spirv::Decoration::AliasedPointer));
+    } else if (decAttr.getName() == getSymbolDecoration(stringifyDecoration(
+                                        spirv::Decoration::RestrictPointer))) {
+      decorations[argID].erase(decAttr.getName());
+      decorations[argID].set(spirv::DecorationAttr::name,
+                             spirv::DecorationAttr::get(
+                                 context, spirv::Decoration::RestrictPointer));
+    }
+  }
----------------
kuhar wrote:

```suggestion
  for (NamedAttribute decAttr : decorations[argID]) {
    for (auto decoration :
         {spirv::Decoration::Aliased, spirv::Decoration::Restrict,
          spirv::Decoration::AliasedPointer,
          spirv::Decoration::RestrictPointer}) {
      if (decAttr.getName() ==
          getSymbolDecoration(stringifyDecoration(decoration))) {
        decorations[argID].erase(decAttr.getName());
        decorations[argID].set(spirv::DecorationAttr::name,
                               spirv::DecorationAttr::get(context, decoration));
      }
    }
  }
```

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