[Mlir-commits] [mlir] [mlir][xegpu] Support N-D block transfers in VectorToXeGPU (PR #210527)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Sat Jul 18 10:46:01 PDT 2026
llvmorg-github-actions[bot] wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir
@llvm/pr-subscribers-mlir-gpu
Author: Jianhui Li (Jianhui-Li)
<details>
<summary>Changes</summary>
Extend the vector.transfer_read/transfer_write lowerings so they can produce N-D xegpu.load_nd/store_nd, not just 1D/2D, and relax the out-of-bounds handling to match load_nd's implicit-zero padding.
Restructure both patterns as "block first, then scatter as fallback.
---
Patch is 44.73 KiB, truncated to 20.00 KiB below, full version: https://github.com/llvm/llvm-project/pull/210527.diff
3 Files Affected:
- (modified) mlir/lib/Conversion/VectorToXeGPU/VectorToXeGPU.cpp (+149-119)
- (modified) mlir/test/Conversion/VectorToXeGPU/transfer-read-to-xegpu.mlir (+117-52)
- (modified) mlir/test/Conversion/VectorToXeGPU/transfer-write-to-xegpu.mlir (+114-111)
``````````diff
diff --git a/mlir/lib/Conversion/VectorToXeGPU/VectorToXeGPU.cpp b/mlir/lib/Conversion/VectorToXeGPU/VectorToXeGPU.cpp
index 9a994e87697f6..306ce6f1e8008 100644
--- a/mlir/lib/Conversion/VectorToXeGPU/VectorToXeGPU.cpp
+++ b/mlir/lib/Conversion/VectorToXeGPU/VectorToXeGPU.cpp
@@ -15,6 +15,7 @@
#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/Dialect/MemRef/IR/MemRef.h"
+#include "mlir/Dialect/UB/IR/UBOps.h"
#include "mlir/Dialect/Utils/IndexingUtils.h"
#include "mlir/Dialect/Utils/StructuredOpsUtils.h"
#include "mlir/Dialect/Vector/IR/VectorOps.h"
@@ -49,14 +50,53 @@ static bool isZeroConstant(Value val) {
.Default(false);
}
+// Return true if the transfer padding value is compatible with the implicit
+// padding of an nd block load. LoadNdOp fills out-of-bounds elements with zero,
+// so a zero constant matches its semantics exactly. A poison padding means the
+// out-of-bounds elements are "don't care", so any implicit padding (including
+// zero) is also acceptable.
+static bool isZeroOrPoisonPadding(Value val) {
+ return isZeroConstant(val) || val.getDefiningOp<ub::PoisonOp>();
+}
+
+// Return true if the permutation map keeps every dimension in place except the
+// innermost two, which are swapped, e.g.:
+// (d0, d1) -> (d1, d0)
+// (d0, d1, d2) -> (d2, d1)
+// (d0, d1, d2, d3) -> (d0, d1, d3, d2)
+// This is the only non-identity permutation an nd block load can realize (by
+// loading the untransposed block and applying a trailing vector.transpose).
+static bool isInnermostTwoDimsTransposed(AffineMap map) {
+ unsigned numResults = map.getNumResults();
+ if (numResults < 2)
+ return false;
+ MLIRContext *ctx = map.getContext();
+ unsigned numInputs = map.getNumInputs();
+ // All but the innermost two results must match the minor-identity map.
+ for (unsigned i = 0; i + 2 < numResults; ++i)
+ if (map.getResult(i) != getAffineDimExpr(numInputs - numResults + i, ctx))
+ return false;
+ // The innermost two results must be the last two input dims, swapped.
+ return map.getResult(numResults - 2) == getAffineDimExpr(numInputs - 1, ctx) &&
+ map.getResult(numResults - 1) == getAffineDimExpr(numInputs - 2, ctx);
+}
+
static LogicalResult storeLoadPreconditions(PatternRewriter &rewriter,
Operation *op, VectorType vecTy,
- MemRefType memTy) {
+ MemRefType memTy,
+ bool allowHighDimVector = false) {
// Validate only vector as the basic vector store and load ops guarantee
// XeGPU-compatible memref source.
unsigned vecRank = vecTy.getRank();
- if (!(vecRank == 1 || vecRank == 2))
+ // nd block loads/stores support N-D block transfers, so callers that lower to
+ // them (transfer_read/transfer_write) allow any non-zero rank. The plain
+ // vector.load/store lowering keeps the 1D/2D restriction.
+ if (allowHighDimVector) {
+ if (vecRank == 0)
+ return rewriter.notifyMatchFailure(op, "Expects non-0D vector");
+ } else if (!(vecRank == 1 || vecRank == 2)) {
return rewriter.notifyMatchFailure(op, "Expects 1D or 2D vector");
+ }
if (!vecTy.getElementType().isIntOrFloat())
return rewriter.notifyMatchFailure(
@@ -594,92 +634,78 @@ struct TransferReadLowering : public OpRewritePattern<vector::TransferReadOp> {
return success();
}
- // TODO:This check needs to be replaced with proper uArch capability check
+ // TODO: This check needs to be replaced with proper uArch capability check.
auto chip = xegpu::getChipStr(readOp);
- // Lower to scattered load Op if the target HW doesn't have 2d block load
- // support and the load is not from shared memory.
- if ((chip != "pvc" && chip != "bmg" && chip != "cri") ||
- readOp.getVectorType().getRank() > 2) {
-
- // TODO: add support for OutOfBound access
- if (isOutOfBounds)
- return failure();
- return lowerToScatteredLoadOp(readOp, rewriter);
- }
-
- // Handle the 1D non-SLM case using load.gather.
- if (loadedVecTy.getRank() == 1 && !isOutOfBounds)
- return lowerToScatteredLoadOp(readOp, rewriter);
-
- // Perform common data transfer checks.
- // TODO: Maybe too strict for SLM case.
- if (failed(
- storeLoadPreconditions(rewriter, readOp, loadedVecTy, readMemTy)))
- return failure();
-
- if (isOutOfBounds && !isZeroConstant(readOp.getPadding()))
- return rewriter.notifyMatchFailure(
- readOp, "Unsupported non-zero padded out-of-bounds read");
+ bool hasBlockLoadSupport =
+ (chip == "pvc" || chip == "bmg" || chip == "cri");
+ // An nd block load can realize a minor-identity map directly, or an
+ // innermost-two-dims transpose via a trailing vector.transpose. Any other
+ // permutation (e.g. a mid-vector transpose of a high-dim load) is left to
+ // the scattered path, which permutes strides explicitly.
AffineMap readMap = readOp.getPermutationMap();
- // Check if this is a transpose: the map must have exactly 2 results,
- // and those 2 results must be the last 2 input dimensions interchanged.
- // Examples:
- // (d0, d1) -> (d1, d0) // transpose
- // (d0, d1) -> (d0, d1) // not a transpose
- // (d0, d1, d2) -> (d2, d1) // transpose (last 2 dims swapped)
- bool isTransposeLoad = false;
- if (readMap.getNumResults() == 2) {
- auto results = readMap.getResults();
- unsigned numInputs = readMap.getNumInputs();
- if (numInputs >= 2) {
- auto lastDim = getAffineDimExpr(numInputs - 1, readMap.getContext());
- auto secondLastDim =
- getAffineDimExpr(numInputs - 2, readMap.getContext());
- isTransposeLoad =
- (results[0] == lastDim && results[1] == secondLastDim);
+ bool isTransposeLoad = isInnermostTwoDimsTransposed(readMap);
+
+ // Prefer an nd block load. It requires HW block-load support, a >1D vector
+ // of a scalar element type backed by a scalar-element memref, and a map the
+ // block load can realize. Out-of-bounds reads are allowed as long as the
+ // padding matches load_nd's implicit zero padding.
+ bool canLowerToLoadNd =
+ hasBlockLoadSupport && loadedVecTy.getRank() > 1 &&
+ (readMap.isMinorIdentity() || isTransposeLoad) &&
+ loadedVecTy.getElementType().isIntOrFloat() &&
+ readMemTy.getElementType().isIntOrFloat() &&
+ (!isOutOfBounds || isZeroOrPoisonPadding(readOp.getPadding()));
+
+ if (canLowerToLoadNd) {
+ auto elementType = loadedVecTy.getElementType();
+
+ SmallVector<int64_t> descShape(loadedVecTy.getShape());
+ if (isTransposeLoad) {
+ // If load is transposed, simply swap the last two dimensions of the
+ // loaded vector type to get the descriptor shape.
+ size_t rank = descShape.size();
+ assert(rank >= 2 && "Transpose requires at least 2 dimensions");
+ std::swap(descShape[rank - 1], descShape[rank - 2]);
+ loadedVecTy = VectorType::get(descShape, elementType);
}
+ auto descType = xegpu::TensorDescType::get(
+ descShape, elementType, /*array_length=*/1,
+ /*boundary_check=*/isOutOfBounds, xegpu::MemorySpace::Global);
+ auto [src, indices] = convertMemrefAndOffsetsToTargetRank(
+ rewriter, loc, readOp.getBase(),
+ getAsOpFoldResult(readOp.getIndices()), loadedVecTy.getRank());
+ // By default, no specific caching policy is assigned.
+ xegpu::CachePolicyAttr hint = nullptr;
+ xegpu::CreateNdDescOp ndDesc = createNdDescriptor(
+ rewriter, loc, descType, dyn_cast<TypedValue<MemRefType>>(src));
+
+ Operation *loadedOp =
+ xegpu::LoadNdOp::create(rewriter, loc, loadedVecTy, ndDesc, indices,
+ /*packed=*/nullptr, /*transpose=*/nullptr,
+ /*l1_hint=*/hint,
+ /*l2_hint=*/hint, /*l3_hint=*/hint,
+ /*layout=*/nullptr);
+ if (isTransposeLoad) {
+ // Undo the innermost-two-dims swap with a trailing vector.transpose:
+ // keep the leading dimensions in place and interchange only the last
+ // two.
+ int64_t rank = loadedVecTy.getRank();
+ SmallVector<int64_t> perm(llvm::to_vector(llvm::seq<int64_t>(0, rank)));
+ std::swap(perm[rank - 1], perm[rank - 2]);
+ loadedOp = vector::TransposeOp::create(rewriter, loc,
+ loadedOp->getResult(0), perm);
+ }
+ rewriter.replaceOp(readOp, loadedOp);
+ return success();
}
- auto elementType = loadedVecTy.getElementType();
-
- SmallVector<int64_t> descShape(loadedVecTy.getShape());
- if (isTransposeLoad) {
- // If load is transposed, simply swap the last two dimensions of the
- // loaded vector type to get the descriptor shape.
- size_t rank = descShape.size();
- assert(rank >= 2 && "Transpose requires at least 2 dimensions");
- std::swap(descShape[rank - 1], descShape[rank - 2]);
- loadedVecTy = VectorType::get(descShape, elementType);
- }
- auto descType = xegpu::TensorDescType::get(
- descShape, elementType, /*array_length=*/1,
- /*boundary_check=*/isOutOfBounds, xegpu::MemorySpace::Global);
- auto [src, indices] = convertMemrefAndOffsetsToTargetRank(
- rewriter, loc, readOp.getBase(), getAsOpFoldResult(readOp.getIndices()),
- loadedVecTy.getRank());
- // By default, no specific caching policy is assigned.
- xegpu::CachePolicyAttr hint = nullptr;
- xegpu::CreateNdDescOp ndDesc = createNdDescriptor(
- rewriter, loc, descType, dyn_cast<TypedValue<MemRefType>>(src));
-
- Operation *loadedOp =
- xegpu::LoadNdOp::create(rewriter, loc, loadedVecTy, ndDesc, indices,
- /*packed=*/nullptr, /*transpose=*/nullptr,
- /*l1_hint=*/hint,
- /*l2_hint=*/hint, /*l3_hint=*/hint,
- /*layout=*/nullptr);
- if (isTransposeLoad) {
- // Transposing the loaded vector with a separate vector.transpose
- // operation
- auto range = llvm::seq<int64_t>(0, readMap.getResults().size());
- SmallVector<int64_t> perm(
- range.rbegin(), range.rend()); // reverse the range for transpose
- loadedOp = vector::TransposeOp::create(rewriter, loc,
- loadedOp->getResult(0), perm);
- }
- rewriter.replaceOp(readOp, loadedOp);
- return success();
+ // Fall back to a scattered load. It supports arbitrary permutations and any
+ // rank, but cannot express out-of-bounds accesses.
+ // TODO: add support for OutOfBound access.
+ if (isOutOfBounds)
+ return failure();
+ return lowerToScatteredLoadOp(readOp, rewriter);
}
};
@@ -727,47 +753,51 @@ struct TransferWriteLowering
return success();
}
- // TODO:This check needs to be replaced with proper uArch capability check
+ // TODO: This check needs to be replaced with proper uArch capability check.
auto chip = xegpu::getChipStr(writeOp);
- // Lower to scattered store Op if the target HW doesn't have 2d block
- // store support and the memref is not SLM.
- if ((chip != "pvc" && chip != "bmg" && chip != "cri") ||
- writeOp.getVectorType().getRank() > 2) {
-
- // TODO: add support for OutOfBound access
- if (writeOp.hasOutOfBoundsDim())
- return failure();
- return lowerToScatteredStoreOp(writeOp, rewriter);
- }
-
- if (failed(storeLoadPreconditions(rewriter, writeOp, vecTy, writeMemTy)))
- return failure();
+ bool hasBlockStoreSupport =
+ (chip == "pvc" || chip == "bmg" || chip == "cri");
+ // Prefer an nd block store. It requires HW block-store support, a >1D
+ // vector of a scalar element type backed by a scalar-element memref, and a
+ // minor-identity map (block stores have no transpose support). Out-of-bounds
+ // writes are handled by the descriptor's boundary check.
AffineMap map = writeOp.getPermutationMap();
- if (!map.isMinorIdentity())
- return rewriter.notifyMatchFailure(writeOp, "Expects identity map");
-
- auto [src, indices] = convertMemrefAndOffsetsToTargetRank(
- rewriter, loc, writeOp.getBase(),
- getAsOpFoldResult(writeOp.getIndices()), vecTy.getRank());
-
- auto descType = xegpu::TensorDescType::get(
- vecTy.getShape(), vecTy.getElementType(),
- /*array_length=*/1, /*boundary_check=*/writeOp.hasOutOfBoundsDim(),
- xegpu::MemorySpace::Global);
- // By default, no specific caching policy is assigned.
- xegpu::CachePolicyAttr hint = nullptr;
- xegpu::CreateNdDescOp ndDesc = createNdDescriptor(
- rewriter, loc, descType, dyn_cast<TypedValue<MemRefType>>(src));
-
- auto storeOp = xegpu::StoreNdOp::create(rewriter, loc, writeOp.getVector(),
- ndDesc, indices,
- /*l1_hint=*/hint,
- /*l2_hint=*/hint, /*l3_hint=*/hint,
- /*layout=*/nullptr);
- rewriter.replaceOp(writeOp, storeOp);
+ bool canLowerToStoreNd =
+ hasBlockStoreSupport && vecTy.getRank() > 1 && map.isMinorIdentity() &&
+ vecTy.getElementType().isIntOrFloat() &&
+ writeMemTy.getElementType().isIntOrFloat();
+
+ if (canLowerToStoreNd) {
+ auto [src, indices] = convertMemrefAndOffsetsToTargetRank(
+ rewriter, loc, writeOp.getBase(),
+ getAsOpFoldResult(writeOp.getIndices()), vecTy.getRank());
+
+ auto descType = xegpu::TensorDescType::get(
+ vecTy.getShape(), vecTy.getElementType(),
+ /*array_length=*/1, /*boundary_check=*/writeOp.hasOutOfBoundsDim(),
+ xegpu::MemorySpace::Global);
+ // By default, no specific caching policy is assigned.
+ xegpu::CachePolicyAttr hint = nullptr;
+ xegpu::CreateNdDescOp ndDesc = createNdDescriptor(
+ rewriter, loc, descType, dyn_cast<TypedValue<MemRefType>>(src));
+
+ auto storeOp =
+ xegpu::StoreNdOp::create(rewriter, loc, writeOp.getVector(), ndDesc,
+ indices,
+ /*l1_hint=*/hint,
+ /*l2_hint=*/hint, /*l3_hint=*/hint,
+ /*layout=*/nullptr);
+ rewriter.replaceOp(writeOp, storeOp);
+ return success();
+ }
- return success();
+ // Fall back to a scattered store. It supports arbitrary permutations and any
+ // rank, but cannot express out-of-bounds accesses.
+ // TODO: add support for OutOfBound access.
+ if (writeOp.hasOutOfBoundsDim())
+ return failure();
+ return lowerToScatteredStoreOp(writeOp, rewriter);
}
};
diff --git a/mlir/test/Conversion/VectorToXeGPU/transfer-read-to-xegpu.mlir b/mlir/test/Conversion/VectorToXeGPU/transfer-read-to-xegpu.mlir
index 3354868cd6cb9..ed495729f9add 100644
--- a/mlir/test/Conversion/VectorToXeGPU/transfer-read-to-xegpu.mlir
+++ b/mlir/test/Conversion/VectorToXeGPU/transfer-read-to-xegpu.mlir
@@ -182,6 +182,33 @@ gpu.func @load_transpose_3d_memref(%source: memref<32x64x128xf32>,
}
+// -----
+// A high-dim load whose innermost two dims are transposed lowers to an
+// (untransposed) nd block load followed by a vector.transpose of the last two
+// dims.
+gpu.module @xevm_module {
+gpu.func @load_high_dim_transposed(%source: memref<2x2x64x128xf16>,
+ %offset: index) -> vector<1x1x64x128xf16> {
+ %c0 = arith.constant 0.0 : f16
+ %0 = vector.transfer_read %source[%offset, %offset, %offset, %offset], %c0
+ {permutation_map = affine_map<(d0, d1, d2, d3) -> (d0, d1, d3, d2)>,
+ in_bounds = [true, true, true, true]}
+ : memref<2x2x64x128xf16>, vector<1x1x64x128xf16>
+ gpu.return %0 : vector<1x1x64x128xf16>
+}
+
+// LOAD-ND-LABEL: @load_high_dim_transposed(
+// LOAD-ND-SAME: %[[SRC:.+]]: memref<2x2x64x128xf16>,
+// LOAD-ND: %[[DESC:.+]] = xegpu.create_nd_tdesc %[[SRC]] : memref<2x2x64x128xf16>
+// LOAD-ND-SAME: -> !xegpu.tensor_desc<1x1x128x64xf16, #xegpu.block_tdesc_attr<boundary_check = false>>
+// LOAD-ND: %[[VEC:.+]] = xegpu.load_nd %[[DESC]]
+// LOAD-ND-SAME: -> vector<1x1x128x64xf16>
+// LOAD-ND: vector.transpose %[[VEC]], [0, 1, 3, 2] : vector<1x1x128x64xf16> to vector<1x1x64x128xf16>
+
+// LOAD-GATHER-LABEL: @load_high_dim_transposed(
+// LOAD-GATHER: %[[VEC:.+]] = xegpu.load {{.*}} : i64, vector<1x1x64x128xindex>, vector<1x1x64x128xi1> -> vector<1x1x64x128xf16>
+}
+
// -----
gpu.module @xevm_module {
gpu.func @load_dynamic_source(%source: memref<?x?x?xf32>,
@@ -275,21 +302,31 @@ gpu.func @load_dynamic_source3(%source: memref<?x?x?x?x?xf32>,
gpu.return %0 : vector<2x4x8x16xf32>
}
-// CHECK-LABEL: @load_dynamic_source3(
-// CHECK-SAME: %[[SRC:.+]]: memref<?x?x?x?x?xf32>
-// CHECK: %[[CST:.+]] = arith.constant dense<true> : vector<2x4x8x16xi1>
-// CHECK: memref.extract_strided_metadata %[[SRC]] : memref<?x?x?x?x?xf32> -> memref<f32>, index, index, index, index, index, index, index, index, index, index, index
-// CHECK-COUNT4: vector.step
-// CHECK-COUNT3: vector.broadcast
-// CHECK-COUNT4: vector.shape_cast
-// CHECK-COUNT4: vector.broadcast {{.*}} : vector<2x4x8x16xindex>
-// CHECK-COUNT3: arith.addi {{.*}} : vector<2x4x8x16xindex>
-// CHECK: %[[SPLAT:.+]] = vector.broadcast {{.*}} : index to vector<2x4x8x16xindex>
-// CHECK: %[[IDX:.+]] = arith.addi %[[SPLAT]], {{.*}} : vector<2x4x8x16xindex>
-// CHECK: %[[COLLAPSE:.+]] = memref.extract_aligned_pointer_as_index %[[SRC]] : memref<?x?x?x?x?xf32> -> index
-// CHECK: %[[COLLAPSE_I:.+]] = arith.index_cast %[[COLLAPSE]] : index to i64
-// CHECK: %[[VEC:.+]] = xegpu.load %[[COLLAPSE_I]]{{\[}}%[[IDX]]{{\]}}, %[[CST]] : i64, vector<2x4x8x16xindex>, vector<2x4x8x16xi1> -> vector<2x4x8x16xf32>
-// CHECK: return %[[VEC]]
+// LOAD-ND-LABEL: @load_dynamic_source3(
+// LOAD-ND-SAME: %[[SRC:.+]]: memref<?x?x?x?x?xf32>
+// LOAD-ND: %[[SUBVIEW:.+]] = memref.subview %[[SRC]]
+// LOAD-ND: %[[BASE_BUFFER:.*]], %[[OFF1:.*]], %[[SIZES:.*]]:4, %[[STRIDES:.*]]:4 = memref.extract_strided_metadata %[[SUBVIEW]]
+// LOAD-ND: %[[DESC:.+]] = xegpu.create_nd_tdesc
+// LOAD-ND-SAME: -> !xegpu.tensor_desc<2x4x8x16xf32, #xegpu.block_tdesc_attr<boundary_check = false>>
+// LOAD-ND: %[[VEC:.+]] = xegpu.load_nd %[[DESC]]
+// LOAD-ND-SAME: -> vector<2x4x8x16xf32>
+// LOAD-ND: return %[[VEC]]
+
+// LOAD-GATHER-LABEL: @load_dynamic_source3(
+// LOAD-GATHER-SAME: %[[SRC:.+]]: memref<?x?x?x?x?xf32>
+// LOAD-GATHER: %[[CST:.+]] = arith.constant dense<true> : vector<2x4x8x16xi1>
+// LOAD-GATHER: memref.extract_strided_metadata %[[SRC]] : memref<?x?x?x?x?xf32> -> memref<f32>, index, index, index, index, index, index, index, index, index, index, index
+// LOAD-GATHER-COUNT4: vector.step
+// LOAD-GATHER-COUNT3: vector.broadcast
+// LOAD-GATHER-COUNT4: vector.shape_cast
+// LOAD-GATHER-COUNT4: vector.broadcast {{.*}} : vector<2x4x8x16xindex>
+// LOAD-GATHER-COUNT3: arith.addi {{.*}} : vector<2x4x8x16xindex>
+// LOAD-GATHER: %[[SPLAT:.+]] = vector.broadcast {{.*}} : index to vector<2x4x8x16xindex>
+// LOAD-GATHER: %[[IDX:.+]] = arith.addi %[[SPLAT]], {{.*}} : vector<2x4x8x16xindex>
+// LOAD-GATHER: %[[COLLAPSE:.+]] = memref.extract_aligned_pointer_as_index %[[SRC]] : memref<?x?x?x?x?xf32> -> index
+// LOAD-GATHER: %[[COLLAPSE_I:.+]] = arith.index_cast %[[COLLAPSE]] : index to i64
+// LOAD-GATHER: %[[VEC:.+]] = xegpu.load %[[COLLAPSE_I]]{{\[}}%[[IDX]]{{\]}}, %[[CST]] : i64, vector<2x4x8x16xindex>, vector<2x4x8x16xi1> -> vector<2x4x8x16xf32>
+// LOAD-GATHER: return %[[VEC]]
}
// -----
@@ -302,21 +339,27 @@ gpu.func @load_high_dim_vector(%source: memref<16x32x64xf32>,
gpu.return %0 : vector<8x16x32xf32>
}
-// CHECK-LABEL: @load_high_dim_vector(
-// CHECK: %[[CST:.+]] = arith.constant dense<true> : vector<8x16x32xi1>
-// CHECK: %[[CST_0:.+]] = arith.constant dense<64> : vector<16xindex>
-// CHECK: %[[CST_1...
[truncated]
``````````
</details>
https://github.com/llvm/llvm-project/pull/210527
More information about the Mlir-commits
mailing list