[Mlir-commits] [mlir] [mlir][SPIR-V] Derive Aligned memory operand for PhysicalStorageBuffer in vector.load/store (PR #214200)

[Arseniy Obolenskiy]

================
@@ -53,6 +53,50 @@ static int getNumBits(Type type) {
   return type.getIntOrFloatBitWidth();
 }
 
+namespace {
+struct MemoryRequirements {
+  spirv::MemoryAccessAttr memoryAccess;
+  IntegerAttr alignment;
+};
+} // namespace
+
+/// PhysicalStorageBuffer pointers always need Aligned, derived from pointee
+/// size unless preferredAlignment is set; other storage classes may omit it.
+static FailureOr<MemoryRequirements>
+calculateMemoryRequirements(Value accessedPtr, uint64_t preferredAlignment) {
+  if (preferredAlignment >= std::numeric_limits<uint32_t>::max())
+    return failure();
+
+  MLIRContext *ctx = accessedPtr.getContext();
+  auto ptrType = cast<spirv::PointerType>(accessedPtr.getType());
+  bool mayOmitAlignment =
+      !preferredAlignment &&
+      ptrType.getStorageClass() != spirv::StorageClass::PhysicalStorageBuffer;
+  if (mayOmitAlignment)
+    return MemoryRequirements{spirv::MemoryAccessAttr{}, IntegerAttr{}};
+
+  // PhysicalStorageBuffer pointers require Aligned.
+  std::optional<int64_t> sizeInBytes;
+  Type pointeeType = ptrType.getPointeeType();
+  if (auto scalarType = dyn_cast<spirv::ScalarType>(pointeeType)) {
+    sizeInBytes = scalarType.getSizeInBytes();
+  } else if (auto vecType = dyn_cast<VectorType>(pointeeType)) {
+    if (auto scalarElem = dyn_cast<spirv::ScalarType>(vecType.getElementType()))
+      if (auto elemSize = scalarElem.getSizeInBytes())
+        sizeInBytes = *elemSize * vecType.getNumElements();
+  }
+
+  if (!sizeInBytes)
+    return failure();
+
+  auto memoryAccess =
+      spirv::MemoryAccessAttr::get(ctx, spirv::MemoryAccess::Aligned);
+  uint64_t alignmentValue =
+      preferredAlignment ? preferredAlignment : *sizeInBytes;
----------------
aobolensk wrote:

Thanks, switched to the element size

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