[mlir][spirv] Retain nontemporal attribute when converting memref load/store (#82119)

Fixes #77156.

Author: agentcooper

mlir/lib/Conversion/MemRefToSPIRV/MemRefToSPIRV.cpp

@@ -445,15 +445,30 @@ DeallocOpPattern::matchAndRewrite(memref::DeallocOp operation,
 // LoadOp
 //===----------------------------------------------------------------------===//
 
-using AlignmentRequirements =
-    FailureOr<std::pair<spirv::MemoryAccessAttr, IntegerAttr>>;
+struct MemoryRequirements {
+  spirv::MemoryAccessAttr memoryAccess;
+  IntegerAttr alignment;
+};
 
 /// Given an accessed SPIR-V pointer, calculates its alignment requirements, if
 /// any.
-static AlignmentRequirements calculateRequiredAlignment(Value accessedPtr) {
+static FailureOr<MemoryRequirements>
+calculateMemoryRequirements(Value accessedPtr, bool isNontemporal) {
+  MLIRContext *ctx = accessedPtr.getContext();
+
+  auto memoryAccess = spirv::MemoryAccess::None;
+  if (isNontemporal) {
+    memoryAccess = spirv::MemoryAccess::Nontemporal;
+  }
+
   auto ptrType = cast<spirv::PointerType>(accessedPtr.getType());
-  if (ptrType.getStorageClass() != spirv::StorageClass::PhysicalStorageBuffer)
-    return std::pair{spirv::MemoryAccessAttr{}, IntegerAttr{}};
+  if (ptrType.getStorageClass() != spirv::StorageClass::PhysicalStorageBuffer) {
+    if (memoryAccess == spirv::MemoryAccess::None) {
+      return MemoryRequirements{spirv::MemoryAccessAttr{}, IntegerAttr{}};
+    }
+    return MemoryRequirements{spirv::MemoryAccessAttr::get(ctx, memoryAccess),
+                              IntegerAttr{}};
+  }
 
   // PhysicalStorageBuffers require the `Aligned` attribute.
   auto pointeeType = dyn_cast<spirv::ScalarType>(ptrType.getPointeeType());
@@ -465,30 +480,32 @@ static AlignmentRequirements calculateRequiredAlignment(Value accessedPtr) {
   if (!sizeInBytes.has_value())
     return failure();
 
-  MLIRContext *ctx = accessedPtr.getContext();
-  auto memAccessAttr =
-      spirv::MemoryAccessAttr::get(ctx, spirv::MemoryAccess::Aligned);
+  memoryAccess = memoryAccess | spirv::MemoryAccess::Aligned;
+  auto memAccessAttr = spirv::MemoryAccessAttr::get(ctx, memoryAccess);
   auto alignment = IntegerAttr::get(IntegerType::get(ctx, 32), *sizeInBytes);
-  return std::pair{memAccessAttr, alignment};
+  return MemoryRequirements{memAccessAttr, alignment};
 }
 
 /// Given an accessed SPIR-V pointer and the original memref load/store
 /// `memAccess` op, calculates the alignment requirements, if any. Takes into
 /// account the alignment attributes applied to the load/store op.
-static AlignmentRequirements
-calculateRequiredAlignment(Value accessedPtr, Operation *memrefAccessOp) {
-  assert(memrefAccessOp);
-  assert((isa<memref::LoadOp, memref::StoreOp>(memrefAccessOp)) &&
-         "Bad op type");
-
+template <class LoadOrStoreOp>
+static FailureOr<MemoryRequirements>
+calculateMemoryRequirements(Value accessedPtr, LoadOrStoreOp loadOrStoreOp) {
+  static_assert(
+      llvm::is_one_of<LoadOrStoreOp, memref::LoadOp, memref::StoreOp>::value,
+      "Must be called on either memref::LoadOp or memref::StoreOp");
+
+  Operation *memrefAccessOp = loadOrStoreOp.getOperation();
   auto memrefMemAccess = memrefAccessOp->getAttrOfType<spirv::MemoryAccessAttr>(
       spirv::attributeName<spirv::MemoryAccess>());
   auto memrefAlignment =
       memrefAccessOp->getAttrOfType<IntegerAttr>("alignment");
   if (memrefMemAccess && memrefAlignment)
-    return std::pair{memrefMemAccess, memrefAlignment};
+    return MemoryRequirements{memrefMemAccess, memrefAlignment};
 
-  return calculateRequiredAlignment(accessedPtr);
+  return calculateMemoryRequirements(accessedPtr,
+                                     loadOrStoreOp.getNontemporal());
 }
 
 LogicalResult
@@ -538,13 +555,12 @@ IntLoadOpPattern::matchAndRewrite(memref::LoadOp loadOp, OpAdaptor adaptor,
   // If the rewritten load op has the same bit width, use the loading value
   // directly.
   if (srcBits == dstBits) {
-    AlignmentRequirements alignmentRequirements =
-        calculateRequiredAlignment(accessChain, loadOp);
-    if (failed(alignmentRequirements))
+    auto memoryRequirements = calculateMemoryRequirements(accessChain, loadOp);
+    if (failed(memoryRequirements))
       return rewriter.notifyMatchFailure(
-          loadOp, "failed to determine alignment requirements");
+          loadOp, "failed to determine memory requirements");
 
-    auto [memoryAccess, alignment] = *alignmentRequirements;
+    auto [memoryAccess, alignment] = *memoryRequirements;
     Value loadVal = rewriter.create<spirv::LoadOp>(loc, accessChain,
                                                    memoryAccess, alignment);
     if (isBool)
@@ -568,13 +584,12 @@ IntLoadOpPattern::matchAndRewrite(memref::LoadOp loadOp, OpAdaptor adaptor,
   assert(accessChainOp.getIndices().size() == 2);
   Value adjustedPtr = adjustAccessChainForBitwidth(typeConverter, accessChainOp,
                                                    srcBits, dstBits, rewriter);
-  AlignmentRequirements alignmentRequirements =
-      calculateRequiredAlignment(adjustedPtr, loadOp);
-  if (failed(alignmentRequirements))
+  auto memoryRequirements = calculateMemoryRequirements(adjustedPtr, loadOp);
+  if (failed(memoryRequirements))
     return rewriter.notifyMatchFailure(
-        loadOp, "failed to determine alignment requirements");
+        loadOp, "failed to determine memory requirements");
 
-  auto [memoryAccess, alignment] = *alignmentRequirements;
+  auto [memoryAccess, alignment] = *memoryRequirements;
   Value spvLoadOp = rewriter.create<spirv::LoadOp>(loc, dstType, adjustedPtr,
                                                    memoryAccess, alignment);
 
@@ -623,13 +638,13 @@ LoadOpPattern::matchAndRewrite(memref::LoadOp loadOp, OpAdaptor adaptor,
   if (!loadPtr)
     return failure();
 
-  AlignmentRequirements requiredAlignment = calculateRequiredAlignment(loadPtr);
-  if (failed(requiredAlignment))
+  auto memoryRequirements = calculateMemoryRequirements(loadPtr, loadOp);
+  if (failed(memoryRequirements))
     return rewriter.notifyMatchFailure(
-        loadOp, "failed to determine alignment requirements");
+        loadOp, "failed to determine memory requirements");
 
-  auto [memAccessAttr, alignment] = *requiredAlignment;
-  rewriter.replaceOpWithNewOp<spirv::LoadOp>(loadOp, loadPtr, memAccessAttr,
+  auto [memoryAccess, alignment] = *memoryRequirements;
+  rewriter.replaceOpWithNewOp<spirv::LoadOp>(loadOp, loadPtr, memoryAccess,
                                              alignment);
   return success();
 }
@@ -689,18 +704,17 @@ IntStoreOpPattern::matchAndRewrite(memref::StoreOp storeOp, OpAdaptor adaptor,
   assert(dstBits % srcBits == 0);
 
   if (srcBits == dstBits) {
-    AlignmentRequirements requiredAlignment =
-        calculateRequiredAlignment(accessChain);
-    if (failed(requiredAlignment))
+    auto memoryRequirements = calculateMemoryRequirements(accessChain, storeOp);
+    if (failed(memoryRequirements))
       return rewriter.notifyMatchFailure(
-          storeOp, "failed to determine alignment requirements");
+          storeOp, "failed to determine memory requirements");
 
-    auto [memAccessAttr, alignment] = *requiredAlignment;
+    auto [memoryAccess, alignment] = *memoryRequirements;
     Value storeVal = adaptor.getValue();
     if (isBool)
       storeVal = castBoolToIntN(loc, storeVal, dstType, rewriter);
     rewriter.replaceOpWithNewOp<spirv::StoreOp>(storeOp, accessChain, storeVal,
-                                                memAccessAttr, alignment);
+                                                memoryAccess, alignment);
     return success();
   }
 
@@ -847,15 +861,14 @@ StoreOpPattern::matchAndRewrite(memref::StoreOp storeOp, OpAdaptor adaptor,
   if (!storePtr)
     return rewriter.notifyMatchFailure(storeOp, "type conversion failed");
 
-  AlignmentRequirements requiredAlignment =
-      calculateRequiredAlignment(storePtr, storeOp);
-  if (failed(requiredAlignment))
+  auto memoryRequirements = calculateMemoryRequirements(storePtr, storeOp);
+  if (failed(memoryRequirements))
     return rewriter.notifyMatchFailure(
-        storeOp, "failed to determine alignment requirements");
+        storeOp, "failed to determine memory requirements");
 
-  auto [memAccessAttr, alignment] = *requiredAlignment;
+  auto [memoryAccess, alignment] = *memoryRequirements;
   rewriter.replaceOpWithNewOp<spirv::StoreOp>(
-      storeOp, storePtr, adaptor.getValue(), memAccessAttr, alignment);
+      storeOp, storePtr, adaptor.getValue(), memoryAccess, alignment);
   return success();
 }
 

mlir/test/Conversion/MemRefToSPIRV/memref-to-spirv.mlir

@@ -431,3 +431,33 @@ func.func @cast_to_static_zero_elems(%arg: memref<?xf32, #spirv.storage_class<Cr
 }
 
 }
+
+// -----
+
+// Check nontemporal attribute
+
+module attributes {
+  spirv.target_env = #spirv.target_env<#spirv.vce<v1.0, [
+    Shader,
+    PhysicalStorageBufferAddresses
+  ], [
+    SPV_KHR_storage_buffer_storage_class,
+    SPV_KHR_physical_storage_buffer
+  ]>, #spirv.resource_limits<>>
+} {
+  func.func @load_nontemporal(%arg0: memref<f32, #spirv.storage_class<StorageBuffer>>) {
+    %0 = memref.load %arg0[] {nontemporal = true} : memref<f32, #spirv.storage_class<StorageBuffer>>
+//       CHECK:  spirv.Load "StorageBuffer" %{{.+}} ["Nontemporal"] : f32
+    memref.store %0, %arg0[] {nontemporal = true} : memref<f32, #spirv.storage_class<StorageBuffer>>
+//       CHECK:  spirv.Store "StorageBuffer" %{{.+}}, %{{.+}} ["Nontemporal"] : f32
+    return
+  }
+
+  func.func @load_nontemporal_aligned(%arg0: memref<f32, #spirv.storage_class<PhysicalStorageBuffer>>) {
+    %0 = memref.load %arg0[] {nontemporal = true} : memref<f32, #spirv.storage_class<PhysicalStorageBuffer>>
+//       CHECK:  spirv.Load "PhysicalStorageBuffer" %{{.+}} ["Aligned|Nontemporal", 4] : f32
+    memref.store %0, %arg0[] {nontemporal = true} : memref<f32, #spirv.storage_class<PhysicalStorageBuffer>>
+//       CHECK:  spirv.Store "PhysicalStorageBuffer" %{{.+}}, %{{.+}} ["Aligned|Nontemporal", 4] : f32
+    return
+  }
+}