[Mlir-commits] [mlir] def16bc - [mlir][spirv] Retain nontemporal attribute when converting memref load/store (#82119)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Sat Mar 2 15:49:22 PST 2024
Author: Artem Tyurin
Date: 2024-03-02T15:49:18-08:00
New Revision: def16bca8133a2f49127a382061aa36406862a15
URL: https://github.com/llvm/llvm-project/commit/def16bca8133a2f49127a382061aa36406862a15
DIFF: https://github.com/llvm/llvm-project/commit/def16bca8133a2f49127a382061aa36406862a15.diff
LOG: [mlir][spirv] Retain nontemporal attribute when converting memref load/store (#82119)
Fixes #77156.
Added:
Modified:
mlir/lib/Conversion/MemRefToSPIRV/MemRefToSPIRV.cpp
mlir/test/Conversion/MemRefToSPIRV/memref-to-spirv.mlir
Removed:
################################################################################
diff --git a/mlir/lib/Conversion/MemRefToSPIRV/MemRefToSPIRV.cpp b/mlir/lib/Conversion/MemRefToSPIRV/MemRefToSPIRV.cpp
index 57d8e894a24b0e..0acb2142f3f68a 100644
--- a/mlir/lib/Conversion/MemRefToSPIRV/MemRefToSPIRV.cpp
+++ b/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();
}
diff --git a/mlir/test/Conversion/MemRefToSPIRV/memref-to-spirv.mlir b/mlir/test/Conversion/MemRefToSPIRV/memref-to-spirv.mlir
index aa05fd9bc8ca89..feb6d4e924015f 100644
--- a/mlir/test/Conversion/MemRefToSPIRV/memref-to-spirv.mlir
+++ b/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
+ }
+}
More information about the Mlir-commits
mailing list