[Mlir-commits] [mlir] ebf518b - [MLIR][XeGPU] Refactor XeGPU layout propagation: passing lane_layout/lane_data with inst_data (#203156)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Tue Jun 23 12:58:32 PDT 2026
Author: Jianhui Li
Date: 2026-06-23T12:58:27-07:00
New Revision: ebf518b37a85a418ef887d77ff688d17d2bd8fa5
URL: https://github.com/llvm/llvm-project/commit/ebf518b37a85a418ef887d77ff688d17d2bd8fa5
DIFF: https://github.com/llvm/llvm-project/commit/ebf518b37a85a418ef887d77ff688d17d2bd8fa5.diff
LOG: [MLIR][XeGPU] Refactor XeGPU layout propagation: passing lane_layout/lane_data with inst_data (#203156)
**Motivation**
Enhance setup* rules in layout propagation to pass lane_layout, and
lane_data information during inst_data propagation, so that the
propagation can have lane level information when choosing an optimal
inst_data. This branch makes that relationship explicit and uniform
across all setup rules.
**Invariant**
All setup rules now produce layouts that satisfy:
Nd ops + dpas/dpas_mx: inst_data = k * (lane_layout * lane_data), k ≥ 1
Scatter/matrix ops + non-anchor ops: inst_data = lane_layout * lane_data
**Key changes in XeGPULayoutImpl**
- New per-op anchor setup rules: setupStoreNdAnchorLayout,
setupPrefetchNdAnchorLayout,
setupLoadNdAnchorLayout (replacing a generic block-IO path) — Nd ops
have rigid lane info
and fit inst_data to the lane factorization.
- New complete*LayoutFromInstData helpers: fill in lane info on
user-provided anchors that
specify only inst_data, by re-running the op's Lane-kind setup with
inst_data as the
destination shape. Covers scatter load/store, block store/load, DPAS,
and DPAS_MX.
- BlockIOInstructionInterface (new in uArchBase.h): shared abstraction
over the
load/store/prefetch 2D-block instructions; the three Xe2 instructions
now implement it.
- setupMultiReductionResultLayout reorganized so the InstData and Lane
branches share the
same lane-layout logic (computeReductionLaneLayoutAndData).
- createScaleLayout (dpas_mx): caps the scale lane_layout by inst_data
so the scale
operand's load_nd satisfies the multiple-of invariant.
- getValidLayouts → getSgLayoutCandidates: renamed and generalized from
2D to N-D
factorization.
- inferShapeCastSourceLayout preserves lane info through 1D↔ND
collapse-style casts.
- chunkSize → contiguousChunkSize parameter rename for the scatter
inst_data decision.
- ResolveLayoutConflicts: retargets an existing convert_layout instead
of chaining a second
one; UnrollConvertLayoutOp short-circuits when input == target.
- Sub-byte (4-bit) element support added to the block-load uArch table.
- Tests updated: propagate-layout-inst-data.mlir, propagate-layout.mlir,
resolve-layout-conflicts.mlir.
Assisted-by-claude
---------
Co-authored-by: Claude Opus 4.7 (1M context) <noreply at anthropic.com>
Added:
Modified:
mlir/include/mlir/Dialect/XeGPU/Transforms/XeGPULayoutImpl.h
mlir/include/mlir/Dialect/XeGPU/uArch/IntelGpuXe2.h
mlir/include/mlir/Dialect/XeGPU/uArch/uArchBase.h
mlir/lib/Dialect/XeGPU/IR/XeGPUDialect.cpp
mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
mlir/lib/Dialect/XeGPU/Transforms/XeGPUUnroll.cpp
mlir/test/Dialect/XeGPU/propagate-layout-inst-data.mlir
mlir/test/Dialect/XeGPU/propagate-layout.mlir
mlir/test/Dialect/XeGPU/resolve-layout-conflicts.mlir
Removed:
################################################################################
diff --git a/mlir/include/mlir/Dialect/XeGPU/Transforms/XeGPULayoutImpl.h b/mlir/include/mlir/Dialect/XeGPU/Transforms/XeGPULayoutImpl.h
index 94d1d5aecbe60..bf61f0ced6048 100644
--- a/mlir/include/mlir/Dialect/XeGPU/Transforms/XeGPULayoutImpl.h
+++ b/mlir/include/mlir/Dialect/XeGPU/Transforms/XeGPULayoutImpl.h
@@ -159,6 +159,22 @@ DistributeLayoutAttr
inferSourceLayoutFromResultForNonAnchorOp(OpOperand &operand,
DistributeLayoutAttr resLayout);
+/// Note on the `consumerLayout` argument used by the consumer-driven setup* /
+/// complete* helpers below:
+///
+/// Layout propagation is a backward dataflow analysis, so a producer learns its
+/// consumers' demands one at a time. The `consumerLayout` passed to these
+/// helpers is the *single* layout that the first consumer to reach the producer
+/// has requested (see `getConsumerLayoutAt`); these helpers do not pick among,
+/// or merge, multiple consumers, and they do not reason about cost (e.g. a
+/// consumer inside a loop vs. one outside). If a producer has several consumers
+/// with conflicting layout demands, only the first-arriving one shapes the
+/// producer's anchor layout here; any later, inconsistent consumer is left
+/// as-is and reconciled afterwards by the layout conflict resolution process
+/// (`ResolveLayoutConflicts`), which inserts a `convert_layout` op on that
+/// edge. So these helpers can always assume exactly one (possibly null)
+/// consumer layout to honor.
+
/// Sets up layout for Multi-Reduction operations by creating a SliceAttr for
/// the result.
///
@@ -208,28 +224,95 @@ DistributeLayoutAttr setupInsertStridedSliceResultLayout(
DistributeLayoutAttr consumerLayout, const uArch::uArch *uArch);
/// Sets up the anchor layout for a load gather operation.
-DistributeLayoutAttr
-setupLoadGatherAnchorLayout(LayoutKind layoutKind, VectorType vectorTy,
- int chunkSize, DistributeLayoutAttr consumerLayout,
- const uArch::uArch *uArch);
+DistributeLayoutAttr setupLoadGatherAnchorLayout(
+ LayoutKind layoutKind, VectorType vectorTy, int contigChunkSize,
+ DistributeLayoutAttr consumerLayout, const uArch::uArch *uArch);
/// Sets up the anchor layout for load matrix operation.
-DistributeLayoutAttr
-setupLoadMatrixAnchorLayout(LayoutKind layoutKind, VectorType vectorTy,
- DistributeLayoutAttr consumerLayout,
- const uArch::uArch *uArch);
+DistributeLayoutAttr setupLoadMatrixAnchorLayout(
+ LayoutKind layoutKind, VectorType vectorTy, int contigChunkSize,
+ DistributeLayoutAttr consumerLayout, const uArch::uArch *uArch);
/// Sets up the anchor layout for a store scatter operation.
DistributeLayoutAttr setupStoreScatterAnchorLayout(LayoutKind layoutKind,
VectorType vectorTy,
- int chunkSize,
+ int contigChunkSize,
const uArch::uArch *uArch);
/// Sets up the anchor layout for a store matrix operation.
DistributeLayoutAttr setupStoreMatrixAnchorLayout(LayoutKind layoutKind,
VectorType vectorTy,
+ int contigChunkSize,
const uArch::uArch *uArch);
+/// If the consumer layout has only inst_data (no lane_layout/lane_data),
+/// completes it by running the corresponding scatter-style Lane-kind setup
+/// rule with inst_data as the destination shape. The resulting lane info is
+/// merged with the consumer's inst_data so downstream setup* paths see a
+/// fully-populated layout.
+/// Returns the layout unchanged when it is null, has no inst_data, or already
+/// carries lane info; returns nullopt when the derived lane factorization does
+/// not divide the user's inst_data (an invalid inst_data).
+std::optional<DistributeLayoutAttr> completeScatterLoadLaneLayoutFromInstData(
+ DistributeLayoutAttr userSpecifiedLayout,
+ DistributeLayoutAttr consumerLayout, Type elemTy,
+ const xegpu::uArch::LoadGatherInstructionInterface *uArchInstruction,
+ const int subgroupSize);
+
+/// Like completeScatterLoadLaneLayoutFromInstData, but for scatter stores
+/// (store_scatter / store_matrix). A store is a data sink: lane info is derived
+/// purely from inst_data using the uArch's StoreScatter per-lane store width,
+/// with no consumer layout to reuse.
+std::optional<DistributeLayoutAttr> completeScatterStoreLaneLayoutFromInstData(
+ DistributeLayoutAttr specifiedLayout, Type elemTy,
+ const xegpu::uArch::StoreScatterInstructionInterface *uArchInstruction,
+ const int subgroupSize);
+
+/// Completes a user-provided 2D-block store_nd / prefetch_nd anchor that has
+/// only inst_data. These ops are data sinks, so lane info is derived purely
+/// from inst_data using the shared BlockIOInstructionInterface; one helper
+/// serves both store_nd and prefetch_nd.
+std::optional<DistributeLayoutAttr> completeBlockStoreLaneLayoutFromInstData(
+ DistributeLayoutAttr specifiedLayout, Type elemTy,
+ const xegpu::uArch::BlockIOInstructionInterface *uArchInstruction,
+ const int subgroupSize);
+
+/// Like completeBlockStoreLaneLayoutFromInstData, but for load_nd. The consumer
+/// layout supplies the transform / transpose / packing properties; the lane
+/// factorization is recomputed from inst_data (load-side lane counts
diff er
+/// from the consumer's).
+std::optional<DistributeLayoutAttr> completeBlockLoadLaneLayoutFromInstData(
+ DistributeLayoutAttr specifiedLayout, DistributeLayoutAttr consumerLayout,
+ Type elemTy,
+ const xegpu::uArch::BlockIOInstructionInterface *uArchInstruction,
+ const int subgroupSize);
+
+/// Sets up the anchor layout for a store_nd operation. StoreNd does not
+/// consider a consumer layout (it is a data sink), and picks its layout from
+/// uArch block parameters. `numSg` is only used for Subgroup-kind layouts.
+DistributeLayoutAttr setupStoreNdAnchorLayout(LayoutKind layoutKind,
+ VectorType vectorTy, int numSg,
+ const uArch::uArch *uArch);
+
+/// Sets up the anchor layout for a prefetch_nd operation. PrefetchNd has no
+/// value result and thus no consumer; it picks its layout from uArch block
+/// parameters. `numSg` is only used for Subgroup-kind layouts.
+DistributeLayoutAttr setupPrefetchNdAnchorLayout(LayoutKind layoutKind,
+ TensorDescType tdescTy,
+ int numSg,
+ const uArch::uArch *uArch);
+
+/// Sets up the anchor layout for a load_nd operation. LoadNd takes a
+/// (downstream) consumer layout and validates it against uArch constraints;
+/// when valid, the consumer's `inst_data` / `sg_layout` are honored.
+/// Otherwise defaults derived from uArch block parameters are used.
+/// `consumerLayout` must be presented. `numSg` is only used for Subgroup-kind
+/// layouts when the consumer does not already provide an sg_layout.
+DistributeLayoutAttr
+setupLoadNdAnchorLayout(LayoutKind layoutKind, VectorType vectorTy,
+ DistributeLayoutAttr consumerLayout, int numSg,
+ const uArch::uArch *uArch);
+
/// Sets up the anchor layouts for a dpas operands (A, B, and C/D).
/// The numSg and consumerLayout (optional) are only used by sg layout creation.
std::optional<std::tuple<DistributeLayoutAttr, DistributeLayoutAttr,
@@ -249,6 +332,31 @@ setupDpasMxLayout(LayoutKind layoutKind, VectorType aTy, VectorType bTy,
DistributeLayoutAttr consumerLayout, int numSg,
const uArch::uArch *uArch);
+/// Completes user-provided DPAS A/B/C-D anchors that carry only inst_data by
+/// filling in lane_layout / lane_data derived from the operand shapes (mirrors
+/// the InstData branch of setupDpasLayout). Returns nullopt if the uArch lacks
+/// the matmul instruction.
+std::optional<std::tuple<DistributeLayoutAttr, DistributeLayoutAttr,
+ DistributeLayoutAttr>>
+completeDpasLaneLayoutFromInstData(DistributeLayoutAttr aLayout,
+ DistributeLayoutAttr bLayout,
+ DistributeLayoutAttr cdLayout,
+ VectorType aTy, VectorType bTy,
+ VectorType cdTy, const uArch::uArch *uArch);
+
+/// Like completeDpasLaneLayoutFromInstData, but for dpas_mx: additionally
+/// re-derives the A_scale / B_scale layouts from the completed A / B layouts.
+std::optional<
+ std::tuple<DistributeLayoutAttr, DistributeLayoutAttr, DistributeLayoutAttr,
+ DistributeLayoutAttr, DistributeLayoutAttr>>
+completeDpasMxLaneLayoutFromInstData(DistributeLayoutAttr aLayout,
+ DistributeLayoutAttr bLayout,
+ DistributeLayoutAttr cdLayout,
+ VectorType aTy, VectorType bTy,
+ VectorType cdTy, VectorType aScaleTy,
+ VectorType bScaleTy,
+ const uArch::uArch *uArch);
+
/// Gets the expected layout for a given consumer operand. This will check if
/// the owning operation of the consumer operand is one of the special layout
/// users and determine the expected layout accordingly.
diff --git a/mlir/include/mlir/Dialect/XeGPU/uArch/IntelGpuXe2.h b/mlir/include/mlir/Dialect/XeGPU/uArch/IntelGpuXe2.h
index eeb1100cc8eab..ff80a77b28d37 100644
--- a/mlir/include/mlir/Dialect/XeGPU/uArch/IntelGpuXe2.h
+++ b/mlir/include/mlir/Dialect/XeGPU/uArch/IntelGpuXe2.h
@@ -45,7 +45,8 @@ struct Xe2Plus : public uArch {
//===----------------------------------------------------------------------===//
// uArch instructions
//===----------------------------------------------------------------------===//
-struct Subgroup2DBlockStoreInstruction : public Instruction {
+struct Subgroup2DBlockStoreInstruction : public Instruction,
+ public BlockIOInstructionInterface {
Subgroup2DBlockStoreInstruction()
: Instruction(InstructionKind::Subgroup2DBlockStore,
InstructionScope::Subgroup) {}
@@ -54,9 +55,12 @@ struct Subgroup2DBlockStoreInstruction : public Instruction {
}
// Source :
// https://registry.khronos.org/OpenCL/extensions/intel/cl_intel_subgroup_2d_block_io.html#_add_a_new_section_5_2_x_cl_intel_subgroup_2d_block_io
+ // Stores ignore the transform / transpose / upConv flags.
std::optional<
std::tuple<llvm::ArrayRef<int>, llvm::ArrayRef<int>, llvm::ArrayRef<int>>>
- getBlockWidthHeightCount(Type elemTy) const {
+ getBlockWidthHeightCount(Type elemTy, bool /*hasTransform*/ = false,
+ bool /*hasTranspose*/ = false,
+ bool /*upConv*/ = false) const override {
const static int kHeight[] = {1, 2, 4, 8};
const static int kWidth16[] = {16};
const static int kWidth32[] = {16};
@@ -73,10 +77,11 @@ struct Subgroup2DBlockStoreInstruction : public Instruction {
return std::nullopt;
}
- int32_t getPackedFormatBitSize() const { return 16; }
+ int32_t getPackedFormatBitSize() const override { return 16; }
};
-struct Subgroup2DBlockLoadInstruction : public Instruction {
+struct Subgroup2DBlockLoadInstruction : public Instruction,
+ public BlockIOInstructionInterface {
Subgroup2DBlockLoadInstruction()
: Instruction(InstructionKind::Subgroup2DBlockLoad,
InstructionScope::Subgroup) {}
@@ -88,48 +93,63 @@ struct Subgroup2DBlockLoadInstruction : public Instruction {
// https://registry.khronos.org/OpenCL/extensions/intel/cl_intel_subgroup_2d_block_io.html#_add_a_new_section_5_2_x_cl_intel_subgroup_2d_block_io
std::optional<
std::tuple<llvm::ArrayRef<int>, llvm::ArrayRef<int>, llvm::ArrayRef<int>>>
- getBlockWidthHeightCount(Type elemTy, bool hasTransform, bool hasTranspose,
- bool upConv = false) const {
+ getBlockWidthHeightCount(Type elemTy, bool hasTransform = false,
+ bool hasTranspose = false,
+ bool upConv = false) const override {
static const int kHeightAtLeast1[] = {1, 2, 4, 8, 16, 32};
static const int kHeightAtLeast8[] = {8, 16, 32};
static const int kHeightAtLeast16[] = {16, 32};
- static const int kHeightAtLeast32[] = {32};
+ static const int kHeight32[] = {32};
+ static const int kHeight64[] = {64};
+ static const int kWidth64[] = {64};
static const int kWidth32[] = {32};
static const int kWidth16[] = {16};
+ static const int kWidthAtLeast16[] = {16, 32};
+ static const int kWidthAtLeast32[] = {32, 64};
static const int kWidth8[] = {8};
static const int32_t kCount1[] = {1};
static const int32_t kCount2[] = {1, 2};
static const int32_t kCount4[] = {1, 2, 4};
static const int32_t kCount4Only[] = {4};
- // (elemBytes, transform, transpose, upConvert)
+ // (elemBits, transform, transpose, upConvert)
using Key = std::tuple<int, uint8_t, uint8_t, uint8_t>;
// (widths, heights, counts)
using Value = std::tuple<llvm::ArrayRef<int32_t>, llvm::ArrayRef<int32_t>,
llvm::ArrayRef<int32_t>>;
+ // The table is keyed on element bit width so sub-byte elements can be
+ // expressed directly. 4-bit elements are packed two-per-byte, so their
+ // widths (or heights, when transformed) are double the 8-bit rows.
static const llvm::DenseMap<Key, Value> kMap = {
- {{1, false, false, false}, {kWidth32, kHeightAtLeast1, kCount2}},
- {{1, false, false, true}, {kWidth16, kHeightAtLeast8, kCount4Only}},
- {{2, false, false, false}, {kWidth16, kHeightAtLeast1, kCount2}},
- {{4, false, false, false}, {kWidth16, kHeightAtLeast1, kCount1}},
+ {{8, false, false, false}, {kWidthAtLeast16, kHeightAtLeast1, kCount2}},
+ {{8, false, false, true}, {kWidth16, kHeightAtLeast8, kCount4Only}},
+ {{16, false, false, false}, {kWidth16, kHeightAtLeast1, kCount2}},
+ {{32, false, false, false}, {kWidth16, kHeightAtLeast1, kCount1}},
// Block Loads with Transform:
- {{1, true, false, false}, {kWidth16, kHeightAtLeast32, kCount4}},
- {{2, true, false, false}, {kWidth16, kHeightAtLeast16, kCount2}},
+ {{8, true, false, false}, {kWidth16, kHeight32, kCount4}},
+ {{16, true, false, false}, {kWidth16, kHeightAtLeast16, kCount2}},
// Block Loads with Transpose:
- {{4, false, true, false}, {kWidth8, kHeightAtLeast16, kCount1}},
- };
- const int elemByteSize = elemTy.getIntOrFloatBitWidth() / 8;
- auto it = kMap.find({elemByteSize, hasTransform, hasTranspose, upConv});
+ {{8, false, true, false}, {kWidth32, kHeightAtLeast16, kCount1}},
+ {{16, false, true, false}, {kWidth16, kHeightAtLeast16, kCount1}},
+ {{32, false, true, false}, {kWidth8, kHeightAtLeast16, kCount1}},
+ // 4-bit elements (sub-byte):
+ {{4, false, false, false}, {kWidthAtLeast32, kHeightAtLeast1, kCount2}},
+ {{4, false, false, true}, {kWidth32, kHeightAtLeast8, kCount4Only}},
+ {{4, true, false, false}, {kWidth16, kHeight64, kCount4}},
+ {{4, false, true, false}, {kWidth64, kHeightAtLeast16, kCount1}}};
+ int elemBitSize = elemTy.getIntOrFloatBitWidth();
+ auto it = kMap.find({elemBitSize, hasTransform, hasTranspose, upConv});
if (it != kMap.end())
return it->second;
return std::nullopt;
}
- int32_t getPackedFormatBitSize() const { return 16; }
+ int32_t getPackedFormatBitSize() const override { return 16; }
};
-struct Subgroup2DBlockPrefetchInstruction : public Instruction {
+struct Subgroup2DBlockPrefetchInstruction : public Instruction,
+ public BlockIOInstructionInterface {
Subgroup2DBlockPrefetchInstruction()
: Instruction(InstructionKind::Subgroup2DBlockPrefetch,
InstructionScope::Subgroup) {}
@@ -138,9 +158,12 @@ struct Subgroup2DBlockPrefetchInstruction : public Instruction {
}
// Source :
// https://registry.khronos.org/OpenCL/extensions/intel/cl_intel_subgroup_buffer_prefetch.html#_add_a_new_section_6_15_x_sub_group_prefetch_functions
+ // Prefetches ignore the transform / transpose / upConv flags.
std::optional<
std::tuple<llvm::ArrayRef<int>, llvm::ArrayRef<int>, llvm::ArrayRef<int>>>
- getBlockWidthHeightCount(Type elemTy) const {
+ getBlockWidthHeightCount(Type elemTy, bool /*hasTransform*/ = false,
+ bool /*hasTranspose*/ = false,
+ bool /*upConv*/ = false) const override {
static const int kHeightAtLeast1[] = {1, 2, 4, 8, 16, 32};
static const int kWidth32[] = {32};
@@ -164,7 +187,7 @@ struct Subgroup2DBlockPrefetchInstruction : public Instruction {
return it->second;
return std::nullopt;
}
- int32_t getPackedFormatBitSize() const { return 16; }
+ int32_t getPackedFormatBitSize() const override { return 16; }
};
struct SubgroupMatrixMultiplyAcc : public Instruction,
diff --git a/mlir/include/mlir/Dialect/XeGPU/uArch/uArchBase.h b/mlir/include/mlir/Dialect/XeGPU/uArch/uArchBase.h
index 147a56a52c188..61db4605e85fa 100644
--- a/mlir/include/mlir/Dialect/XeGPU/uArch/uArchBase.h
+++ b/mlir/include/mlir/Dialect/XeGPU/uArch/uArchBase.h
@@ -19,6 +19,7 @@
#include <iostream>
#include <map>
#include <mutex>
+#include <optional>
#include <shared_mutex>
#include <tuple>
@@ -255,6 +256,24 @@ struct MMAInstructionInterface {
virtual ~MMAInstructionInterface() = default;
};
+// Interface for subgroup-level 2D block instructions (load / store / prefetch).
+// All three describe the set of hardware-supported block shapes via
+// (width, height, count) tuples and share a packed-format bit size. The
+// transform / transpose / upConv flags are only meaningful for loads; store
+// and prefetch implementations ignore them.
+struct BlockIOInstructionInterface {
+ // Returns the supported (widths, heights, counts) for the given element
+ // type, or std::nullopt if the element type is unsupported.
+ virtual std::optional<
+ std::tuple<llvm::ArrayRef<int>, llvm::ArrayRef<int>, llvm::ArrayRef<int>>>
+ getBlockWidthHeightCount(Type elemTy, bool hasTransform = false,
+ bool hasTranspose = false,
+ bool upConv = false) const = 0;
+ // Bit size of the packed format used by this block instruction.
+ virtual int32_t getPackedFormatBitSize() const = 0;
+ virtual ~BlockIOInstructionInterface() = default;
+};
+
//===----------------------------------------------------------------------===//
// Common instructions (shared across architectures)
//===----------------------------------------------------------------------===//
diff --git a/mlir/lib/Dialect/XeGPU/IR/XeGPUDialect.cpp b/mlir/lib/Dialect/XeGPU/IR/XeGPUDialect.cpp
index 4aa1c0d666a94..311cf9a64c0c4 100644
--- a/mlir/lib/Dialect/XeGPU/IR/XeGPUDialect.cpp
+++ b/mlir/lib/Dialect/XeGPU/IR/XeGPUDialect.cpp
@@ -822,8 +822,8 @@ DistributeLayoutAttr LayoutAttr::expandDim(int64_t dim,
}
// Per-sg view used as the base for lane_layout / lane_data / inst_data:
- // targetShape[i] / sg_layout[i] when sg_layout is present, else
- // targetShape itself.
+ // targetShape[i] / sg_layout[i] when sg_layout is present (and not
+ // replicated), else targetShape itself.
SmallVector<int64_t> perSgShape(targetShape.begin(), targetShape.end());
if (hasSgLayout && !sgDataReplicated)
for (int64_t i = 0; i < expCount; ++i)
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
index ca766f67c2583..2a13997aa181f 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
@@ -370,6 +370,86 @@ void xegpu::removeTemporaryLayoutAttrs(Operation *op) {
});
}
+/// Returns true if every dimension of `shape` except the innermost
+/// `numInnerDims` is a unit (size-1) dimension.
+static bool leadingDimsAreUnit(ArrayRef<int64_t> shape, int numInnerDims) {
+ int numLeading = static_cast<int>(shape.size()) - numInnerDims;
+ if (numLeading <= 0)
+ return true;
+ return llvm::all_of(shape.take_front(numLeading),
+ [](int64_t dim) { return dim == 1; });
+}
+
+static xegpu::LayoutAttr buildInstDataLayoutWithLane(
+ mlir::MLIRContext *context, ArrayRef<int64_t> instData,
+ ArrayRef<int64_t> laneLayout, ArrayRef<int64_t> laneData,
+ DenseI32ArrayAttr orderAttr = nullptr) {
+ auto toI32Attr = [&](auto range) {
+ SmallVector<int32_t> v(range.begin(), range.end());
+ return DenseI32ArrayAttr::get(context, v);
+ };
+ return xegpu::LayoutAttr::get(context, /*sg_layout=*/nullptr,
+ /*sg_data=*/nullptr, toI32Attr(instData),
+ toI32Attr(laneLayout), toI32Attr(laneData),
+ orderAttr);
+}
+
+static bool isValidLaneLayout(ArrayRef<int64_t> dataShape,
+ ArrayRef<int64_t> laneLayout,
+ ArrayRef<int64_t> laneData) {
+ return !llvm::any_of(llvm::seq<int>(0, dataShape.size()), [&](int dim) {
+ return dataShape[dim] % (laneLayout[dim] * laneData[dim]) != 0;
+ });
+}
+
+static xegpu::LayoutAttr
+buildLaneLayout(mlir::MLIRContext *context, ArrayRef<int64_t> laneLayout,
+ ArrayRef<int64_t> laneData,
+ DenseI32ArrayAttr orderAttr = nullptr) {
+ auto toI32Attr = [&](auto range) {
+ SmallVector<int32_t> v(range.begin(), range.end());
+ return DenseI32ArrayAttr::get(context, v);
+ };
+ return xegpu::LayoutAttr::get(context, /*sg_layout=*/nullptr,
+ /*sg_data=*/nullptr,
+ /*inst_data=*/nullptr, toI32Attr(laneLayout),
+ toI32Attr(laneData), orderAttr);
+}
+
+static xegpu::LayoutAttr
+buildLayout(mlir::MLIRContext *context, ArrayRef<int64_t> sgLayout,
+ ArrayRef<int64_t> sgData, ArrayRef<int64_t> instData,
+ ArrayRef<int64_t> laneLayout, ArrayRef<int64_t> laneData,
+ DenseI32ArrayAttr orderAttr = nullptr) {
+ auto toI32Attr = [&](auto range) {
+ SmallVector<int32_t> v(range.begin(), range.end());
+ return DenseI32ArrayAttr::get(context, v);
+ };
+ return xegpu::LayoutAttr::get(
+ context, sgLayout.empty() ? nullptr : toI32Attr(sgLayout),
+ sgData.empty() ? nullptr : toI32Attr(sgData),
+ instData.empty() ? nullptr : toI32Attr(instData),
+ laneLayout.empty() ? nullptr : toI32Attr(laneLayout),
+ laneData.empty() ? nullptr : toI32Attr(laneData), orderAttr);
+}
+
+static xegpu::LayoutAttr buildSgLayout(mlir::MLIRContext *context,
+ ArrayRef<int64_t> wgTileShape,
+ ArrayRef<int64_t> sgLayout,
+ int dimK = -1,
+ DenseI32ArrayAttr orderAttr = nullptr) {
+ SmallVector<int64_t> sgData(sgLayout.size());
+ for (int dim = 0; dim < (int)sgLayout.size(); ++dim) {
+ if (dim == dimK)
+ sgData[dim] = wgTileShape[dim];
+ else
+ sgData[dim] = wgTileShape[dim] / sgLayout[dim];
+ }
+ return buildLayout(context, sgLayout, sgData,
+ /*inst_data=*/{}, /*lane_layout=*/{},
+ /*lane_data=*/{}, /*order=*/nullptr);
+}
+
/// Infers the source layout attribute for a broadcast operation given the
/// result layout attribute, result shape, source shape.
xegpu::DistributeLayoutAttr
@@ -666,21 +746,13 @@ xegpu::inferExtractSourceLayout(xegpu::DistributeLayoutAttr resLayout,
order.push_back(dimDiff - 1 - i);
}
- DenseI32ArrayAttr orderAttr = resLayout ? resLayout.getOrder() : nullptr;
- auto toAttr = [&](ArrayRef<int64_t> v) -> DenseI32ArrayAttr {
- if (v.empty())
- return DenseI32ArrayAttr();
- SmallVector<int32_t> v32(v.begin(), v.end());
- return DenseI32ArrayAttr::get(context, v32);
- };
- auto srcLayout = xegpu::LayoutAttr::get(
- context, sgLayout.empty() ? nullptr : toAttr(sgLayout),
- sgData.empty() ? nullptr : toAttr(sgData),
- instData.empty() ? nullptr : toAttr(instData),
- laneLayout.empty() ? nullptr : toAttr(laneLayout),
- laneData.empty() ? nullptr : toAttr(laneData),
- (!orderAttr || orderAttr.empty()) ? nullptr : toAttr(order));
- return srcLayout;
+ DenseI32ArrayAttr orderAttr = DenseI32ArrayAttr::get(
+ context, SmallVector<int32_t>(order.begin(), order.end()));
+ if (!resLayout.getOrder())
+ orderAttr = nullptr;
+
+ return buildLayout(context, sgLayout, sgData, instData, laneLayout,
+ laneData, orderAttr);
}
return resLayout;
}
@@ -727,18 +799,6 @@ xegpu::inferShapeCastSourceLayout(xegpu::DistributeLayoutAttr resLayout,
// Use case 3: General dim collapse, for cross-sg reduction to SLM and other
// shape casts where consecutive src dims fold into a single dst dim.
- //
- // Mirrors use case 2's elegant shape: walk the dst-side groups and call
- // a single layout-attribute primitive per group. Here the primitive is
- // `expandDim(dim, targetShape)`, the inverse of `collapseDims`. It applies
- // the per-field distribution policy required for a no-data-movement collapse
- // (sg_layout/lane_layout spread outer-to-inner; sg_data/lane_data/inst_data
- // fill innermost-first; inst_data is seeded from lane_layout * lane_data).
- // See LayoutAttr::expandDim for the full policy.
- //
- // Iteration goes innermost-first (reverse dst order) so that each
- // expandDim/dropDims call only mutates dst positions whose indices are
- // unaffected by earlier calls.
SmallVector<SmallVector<int64_t>> collapseDims;
if (xegpu::matchDimCollapse(srcShape, resShape, collapseDims)) {
auto srcLayout = resLayout;
@@ -746,12 +806,10 @@ xegpu::inferShapeCastSourceLayout(xegpu::DistributeLayoutAttr resLayout,
dstIdx >= 0; --dstIdx) {
ArrayRef<int64_t> srcDims = collapseDims[dstIdx];
if (srcDims.empty()) {
- // Unit dst dim with no backing src dim: drop it.
srcLayout = srcLayout.dropDims({dstIdx});
continue;
}
if (srcDims.size() == 1)
- // 1:1 mapping, nothing to do for this dim.
continue;
SmallVector<int64_t> targetShape;
targetShape.reserve(srcDims.size());
@@ -761,7 +819,6 @@ xegpu::inferShapeCastSourceLayout(xegpu::DistributeLayoutAttr resLayout,
}
return srcLayout;
}
- llvm_unreachable("running into unsupported shape cast scenarios");
return nullptr;
}
@@ -776,460 +833,877 @@ xegpu::DistributeLayoutAttr xegpu::inferMaskOffsetLayoutForScatterIO(
return payloadLayout;
}
-/// Sets up layout for reduction operations by creating a SliceAttr for the
-/// result.
-///
-/// Algorithm Overview:
-/// This function attempts to construct a source layout that, when sliced along
-/// reduction dimensions, produces a result layout compatible with the
-/// consumer layout.
-///
-/// For subgroup layouts, it first tries to align the source layout's subgroup
-/// layout and data with the consumer's layout on non-reduction dimensions.
-/// Then, it distributes remaining subgroups across reduction dimensions. This
-/// avoids subgroup data redistribution overhead between the reduced result and
-/// its consumer. When the consumer layout is a slice layout, it attempts to
-/// reuse the slice layout's parent layout for the source to further minimize
-/// potential data redistribution.
-///
-/// InstData requries {1, ..., min(maxReduceVectorSize, srcShape),subgroupSize}
-/// Lane Layout requires {1, ..., 1, subgroupSize}
-/// Lane data requires {1, ..., min(maxReduceVectorSize, srcShape), 1}
-///
-/// Examples:
-/// 1. Subgroup layout - Row reduction on 2D tensor:
-/// srcShape=[32, 128], reductionDims=[1], resShape=[32], subgroupSize=16,
-/// NumSg=32
-/// * Consumer Layout:
-/// #xegpu.slice<#xegpu.layout<sg_layout=[4, 8], sg_data=[8, 8]>, dims =
-/// [1]>}
-//// * Result Layout:
-/// #xegpu.slice<#xegpu.layout<sg_layout=[4, 8],sg_data=[8, 16]>, dims =
-/// [1]>}
-/// Note that the sg_layout is reused but sg_data needs to be adjusted to
-/// evenly distribute the source tensor tile among the reduction dim.
-///
-/// 2. Subgroup layout - Same example above but consumer doesn't have a
-/// reusable slice layout.
-/// * Consumer Layout:
-/// #xegpu.layout<sgLayout=[32], sgData=[1]>
-/// * Result Layout:
-/// #xegpu.slice<#xegpu.layout<sgLayout=[32,1], sgData=[1, 64]>, dims =
-/// [1]>}
-/// * Consumer Layout:
-/// #xegpu.slice<#xegpu.layout<sgLayout=[8, 2, 4], sgData=[4, 64, 32]>,
-/// dims = [1, 2]>}
-/// * Result Layout:
-/// #xegpu.slice<#xegpu.layout<sgLayout=[8,4], sgData=[4, 32]>, dims =
-/// [1]>}
-/// Note that the consumer's layout can't be directly reused as is.
-/// So the algorithm distributes all subgroups on non reduction dimensions
-/// first and then distribute remaining subgroups on the reduction
-/// dimension.
-///
-/// 2. InstData layout - Column reduction:
-/// srcShape=[32, 64], reductionDims=[0], subgroupSize=16
-/// Result: instData=[1, 16] (maxReduceVectorSize=1, subgroupSize on
-/// innermost)
-///
-/// 3. Lane layout - Multi-dimensional reduction:
-/// srcShape=[16, 32, 64], reductionDims=[1], subgroupSize=16
-/// Result: laneLayout=[1, 1, 16], laneData=[1, 1, 1]
-/// (subgroupSize on innermost dim, max vector size on reduction dim)
-
-xegpu::SliceAttr xegpu::setupMultiReductionResultLayout(
- xegpu::LayoutKind layoutKind, VectorType srcVecTy,
- DistributeLayoutAttr consumerLayout, SmallVector<int64_t> reductionDims,
- int numSg, const xegpu::uArch::uArch *uArch) {
+//===----------------------------------------------------------------------===//
+// Layout derivation helpers: factorize sgCount into
+// sg_layout candidates, then
+// compute per-subgroup (sgData) and per-lane
+// (lane_layout/lane_data/inst_data).
+//===----------------------------------------------------------------------===//
- auto srcShape = srcVecTy.getShape();
- int srcRank = srcShape.size();
- auto context = srcVecTy.getContext();
+using LayoutRepresentation = SmallVector<int64_t>;
+
+/// Enumerates all ways to split `total` into `rank` factors whose product
+/// equals `total`. Returns the list of all such factorizations.
+static SmallVector<LayoutRepresentation> enumerateFactorizations(int64_t total,
+ int64_t rank) {
+ SmallVector<LayoutRepresentation> results;
+ SmallVector<int64_t> current(rank, 0);
+
+ // Returns all divisors of `n` in ascending order.
+ auto getDivisors = [](int64_t n) {
+ SmallVector<int64_t> divs;
+ for (int64_t i = 1; i * i <= n; ++i) {
+ if (n % i == 0) {
+ divs.push_back(i);
+ if (i != n / i)
+ divs.push_back(n / i);
+ }
+ }
+ llvm::sort(divs);
+ return divs;
+ };
- // Helper lambda to convert int64 vectors to int32 DenseArrayAttr
- auto toInt32Attr = [&](ArrayRef<int64_t> vec) {
- SmallVector<int32_t> vec32(vec.begin(), vec.end());
- return DenseI32ArrayAttr::get(context, vec32);
+ std::function<void(int64_t, int64_t)> generate = [&](int64_t dim,
+ int64_t remaining) {
+ if (dim == rank - 1) {
+ current[dim] = remaining;
+ results.push_back(LayoutRepresentation(current));
+ return;
+ }
+ for (int64_t factor : getDivisors(remaining)) {
+ current[dim] = factor;
+ generate(dim + 1, remaining / factor);
+ }
};
- const int subgroupSize = uArch->getSubgroupSize();
- int64_t maxReduceVectorSize = 1; // could extend to spirv vector Size
- xegpu::DistributeLayoutAttr srcLayout;
- if (layoutKind == xegpu::LayoutKind::Subgroup) {
- xegpu::SliceAttr consumerSliceLayout =
- dyn_cast_if_present<xegpu::SliceAttr>(consumerLayout);
- if (consumerSliceLayout &&
- consumerSliceLayout.getDims().asArrayRef().equals(reductionDims)) {
- srcLayout = consumerSliceLayout.getParent();
- SmallVector<int64_t> sgLayoutFromConsumer =
- srcLayout.getEffectiveSgLayoutAsInt();
- auto srcSgData = computeShapeRatio(srcShape, sgLayoutFromConsumer);
- if (srcSgData)
- for (int dim = 0; dim < srcRank; dim++) {
- if (llvm::is_contained(reductionDims, dim))
- srcLayout =
- srcLayout.setDimData(dim, srcSgData.value()[dim], -1, -1);
- }
- } else {
- SmallVector<int64_t> consumerSgLayout =
- consumerLayout ? consumerLayout.getEffectiveSgLayoutAsInt()
- : SmallVector<int64_t>();
- SmallVector<int64_t> consumerSgData =
- consumerLayout ? consumerLayout.getEffectiveSgDataAsInt()
- : SmallVector<int64_t>();
- SmallVector<int64_t> consumerOrder =
- consumerLayout ? consumerLayout.getEffectiveOrderAsInt()
- : SmallVector<int64_t>();
- DenseI32ArrayAttr orderAttr =
- consumerLayout ? consumerLayout.getOrder() : nullptr;
- SmallVector<int64_t> sgLayout(srcRank), sgData(srcRank), order(srcRank);
- int remainingSgCount =
- consumerLayout ? consumerLayout.getNumSubgroups() : numSg;
- int consumerIdx = 0;
+ generate(0, total);
+ return results;
+}
- // First pass: Match consumer's layout on non-reduction dimensions
- for (int i = 0; i < srcRank; i++) {
- if (!llvm::is_contained(reductionDims, i) &&
- consumerIdx < static_cast<int>(consumerSgLayout.size())) {
- sgLayout[i] = consumerSgLayout[consumerIdx];
- sgData[i] = consumerSgData[consumerIdx];
- remainingSgCount /= sgLayout[i];
- order[i] = consumerOrder[consumerIdx];
- consumerIdx++;
- }
- }
+// Computes all valid N-dimensional sg_layout candidates for the given
+// sgCount, whose sgData (= wgShape / sgLayout):
+// 1. Evenly divides wgShape (i.e., wgShape[d] % sgLayout[d] == 0).
+// 2. Is a multiple of instData (i.e., sgData[d] % instData[d] == 0).
+// Results are sorted by balance (smallest max-min spread first), with
+// lexicographic order as a tiebreaker.
+//
+// Example (2D):
+// wgShape = [128, 64], instData = [8, 16], sgCount = 32
+// Returns: [[8,4], [16,2]], corresponding to sgData [16,16] and [8,32].
+static SmallVector<LayoutRepresentation>
+getSgLayoutCandidates(ArrayRef<int64_t> wgShape, ArrayRef<int64_t> instData,
+ int64_t sgCount) {
+ int64_t rank = wgShape.size();
+ assert(rank > 0 && "wgShape must be non-empty");
+ assert(static_cast<int64_t>(instData.size()) == rank &&
+ "instData rank must match wgShape rank");
- // Second pass: Distribute remaining subgroups across reduction dimensions
- // the reduction to scalar case is handled only by this loop
- int64_t remainOrder = consumerSgLayout.size();
- for (int i = 0; i < srcRank; i++) {
- if (llvm::is_contained(reductionDims, i)) {
- sgLayout[i] =
- std::min(srcShape[i], static_cast<int64_t>(remainingSgCount));
- assert((srcShape[i] % sgLayout[i] == 0) &&
- "source shape not divisible by sg_layout");
- sgData[i] = srcShape[i] / sgLayout[i];
- remainingSgCount /= sgLayout[i];
- order[i] = remainOrder++;
- }
- }
+ // Step 1: Get all N-D factorizations of sgCount.
+ auto allFactorizations = enumerateFactorizations(sgCount, rank);
- assert(remainingSgCount == 1 && "not all subgroups distributed");
- srcLayout = xegpu::LayoutAttr::get(
- context, toInt32Attr(sgLayout), toInt32Attr(sgData),
- /*inst_data =*/nullptr, /*lane_layout =*/nullptr,
- /*lane_data =*/nullptr, /*order =*/
- (!orderAttr || orderAttr.empty()) ? nullptr : toInt32Attr(order));
+ // Step 2: Filter to keep only valid candidates.
+ SmallVector<LayoutRepresentation> candidates;
+ for (const auto &sgLayout : allFactorizations) {
+ bool valid = true;
+ for (int64_t dim = 0; dim < rank; ++dim) {
+ if (wgShape[dim] % sgLayout[dim] != 0) {
+ valid = false;
+ break;
+ }
+ int64_t sgData = wgShape[dim] / sgLayout[dim];
+ if (sgData % instData[dim] != 0) {
+ valid = false;
+ break;
+ }
}
- } else if (layoutKind == xegpu::LayoutKind::InstData) {
-
- SmallVector<int64_t> instData(srcRank, 1);
- if (srcRank >= 2)
- instData[srcRank - 2] =
- std::min(maxReduceVectorSize, srcShape[srcRank - 2]);
- instData[srcRank - 1] =
- std::min(static_cast<int64_t>(subgroupSize), srcShape[srcRank - 1]);
- srcLayout = xegpu::LayoutAttr::get(context, toInt32Attr(instData));
- } else if (layoutKind == xegpu::LayoutKind::Lane) {
-
- SmallVector<int64_t> laneLayout(srcRank, 1), laneData(srcRank, 1);
- laneLayout[srcRank - 1] =
- std::min(static_cast<int64_t>(subgroupSize), srcShape[srcRank - 1]);
- if (srcRank >= 2)
- laneData[srcRank - 2] =
- std::min(maxReduceVectorSize, srcShape[srcRank - 2]);
- srcLayout = xegpu::LayoutAttr::get(context, toInt32Attr(laneLayout),
- toInt32Attr(laneData));
+ if (valid)
+ candidates.push_back(sgLayout);
}
- return xegpu::SliceAttr::get(context, srcLayout,
- DenseI64ArrayAttr::get(context, reductionDims));
+ // Step 3: Sort by balance (smallest max-min spread), then lexicographic.
+ llvm::sort(candidates, [](const LayoutRepresentation &lhs,
+ const LayoutRepresentation &rhs) {
+ int64_t spreadLhs = *llvm::max_element(lhs) - *llvm::min_element(lhs);
+ int64_t spreadRhs = *llvm::max_element(rhs) - *llvm::min_element(rhs);
+ if (spreadLhs != spreadRhs)
+ return spreadLhs < spreadRhs;
+ return lhs < rhs;
+ });
+ return candidates;
}
-/// Sets up layout for Reduction operations by creating a SliceAttr for the
-/// result.
-xegpu::SliceAttr
-xegpu::setupReductionResultLayout(xegpu::LayoutKind layoutKind,
- VectorType srcVecTy,
- const xegpu::uArch::uArch *uArch) {
+/// Helper function to compute inst_data vectors for DPAS operands A, B, and
+/// C/D.
+static std::optional<SmallVector<int64_t>> get2DBlockIOInstDataLayout(
+ ArrayRef<int64_t> dataShape, Type elemTy,
+ const xegpu::uArch::BlockIOInstructionInterface *uArchInstruction,
+ bool transform = false, bool transpose = false) {
+ int rank = dataShape.size();
+ auto blockWHC =
+ uArchInstruction->getBlockWidthHeightCount(elemTy, transform, transpose);
+ if (!blockWHC)
+ return std::nullopt;
+ auto [bWidths, bHeights, bCounts] = blockWHC.value();
+ // Compute inst_data from hardware block params. For Nd ops, the lane
+ // factorization above (laneLayout / laneData) is rigid; inst_data must be
+ // a multiple of lane_layout * lane_data on each dim (Category A
+ // invariant).
+ SmallVector<int64_t> instData(rank, 1);
+ assert(rank >= 2 && "dataShape must be at least 2D for 2D-block IO");
+ int instWidth =
+ xegpu::getLargestDivisor(static_cast<int>(dataShape.back()), bWidths);
+ int instHeight =
+ xegpu::getLargestDivisor(static_cast<int>(dataShape[rank - 2]), bHeights);
+ instData.back() = instWidth;
+ instData[rank - 2] = instHeight;
+
+ return instData;
+}
- auto srcShape = srcVecTy.getShape();
- auto context = srcVecTy.getContext();
- auto subgroupSize = uArch->getSubgroupSize();
- xegpu::LayoutAttr srcLayout;
+/// Helper function to compute inst_data vectors for DPAS operands A, B, and
+/// C/D. Look up the uArch table and search for the largest supported block size
+/// that divides the data shape
+static std::optional<std::tuple<SmallVector<int64_t>, SmallVector<int64_t>,
+ SmallVector<int64_t>>>
+getDpasInstDataLayouts(
+ VectorType aTy, VectorType bTy, VectorType cdTy,
+ const xegpu::uArch::MMAInstructionInterface *uArchInstruction) {
- if (layoutKind == xegpu::LayoutKind::Subgroup) {
- assert(true && "subgroup layout assignment not supported for reduction (op "
- "is not expected at this level).");
- } else if (layoutKind == xegpu::LayoutKind::InstData) {
- assert(true && "instData layout assignment not supported for reduction (op "
- "is not expected at this level).");
- } else if (layoutKind == xegpu::LayoutKind::Lane) {
- SmallVector<int32_t> laneLayout(1), laneData(1);
- laneLayout[0] = std::min(subgroupSize, static_cast<int32_t>(srcShape[0]));
- laneData[0] = 1;
- srcLayout = xegpu::LayoutAttr::get(
- context, DenseI32ArrayAttr::get(context, laneLayout),
- DenseI32ArrayAttr::get(context, laneData));
- }
+ // M dimension is the second-to-last dim of A (handles batch dims).
+ const unsigned dataALen = aTy.getShape()[aTy.getRank() - 2];
+ auto supportedALen = uArchInstruction->getSupportedM(aTy.getElementType());
+ const int maxALen =
+ xegpu::getLargestDivisor(dataALen, ArrayRef<unsigned>(supportedALen));
- auto result = xegpu::SliceAttr::get(context, srcLayout,
- DenseI64ArrayAttr::get(context, 0));
- return result;
-}
+ // N dimension is the last dim of B.
+ const unsigned dataBLen = bTy.getShape().back();
+ auto supportedBLen = uArchInstruction->getSupportedN(bTy.getElementType());
+ const int maxBLen =
+ xegpu::getLargestDivisor(dataBLen, ArrayRef<unsigned>(supportedBLen));
-/// Sets up the result layout for a bitcast operation.
-/// When casting to a smaller bitwidth, adjusts the layout dimensions (sgData,
-/// instData, or laneData) by multiplying by the bitwidth ratio to ensure the
-/// result layout can be correctly divided back to the source layout during
-/// inference.
-///
-/// Examples:
-/// 1. Casting f32 -> f16 (32-bit to 16-bit, bitWidthRatio = 2):
-/// Consumer layout: instData=[1, 16], subgroupSize=16
-/// Source shape: [8, 32]
-/// Result layout: instData=[1, 32] (16 * 2)
-/// The innermost dimension is multiplied by 2 to maintain consistency.
-///
-/// 2. Casting f32 -> i8 (32-bit to 8-bit, bitWidthRatio = 4):
-/// Consumer instData=[1, 16], subgroupSize=16
-/// Source shape: [4, 128]
-/// adjust the instData from [1, 16] to [1, 16 * 4 = 64]
-///
-/// 3. Casting i8 -> i32 (8-bit to 32-bit, bitWidthRatio = 1/4):
-/// Consumer layout: laneLayout=[1, 16], laneData=[1, 4]
-/// No adjustment needed - returns consumer layout directly.
-///
-xegpu::DistributeLayoutAttr xegpu::setupBitCastResultLayout(
- xegpu::LayoutKind layoutKind, VectorType srcVecTy, VectorType resVecTy,
- DistributeLayoutAttr consumerLayout, const xegpu::uArch::uArch *uArch) {
+ auto supportedCLen = uArchInstruction->getSupportedN(cdTy.getElementType());
+ const int maxCLen =
+ xegpu::getLargestDivisor(dataBLen, ArrayRef<unsigned>(supportedCLen));
+ if (maxALen == -1 || maxBLen == -1 || maxCLen == -1)
+ return std::nullopt;
- int srcElemTyBitWidth = srcVecTy.getElementType().getIntOrFloatBitWidth();
- int resElemTyBitWidth = resVecTy.getElementType().getIntOrFloatBitWidth();
+ auto supportedKLen = uArchInstruction->getSupportedK(aTy.getElementType());
+ if (supportedKLen.empty())
+ return std::nullopt;
+ auto kDimSize = supportedKLen[0];
- ArrayRef<int64_t> srcShape = srcVecTy.getShape();
- ArrayRef<int64_t> resShape = resVecTy.getShape();
- SmallVector<int64_t> sgData = consumerLayout.getEffectiveSgDataAsInt();
- SmallVector<int64_t> instData = consumerLayout.getEffectiveInstDataAsInt();
- SmallVector<int64_t> laneData = consumerLayout.getEffectiveLaneDataAsInt();
- SmallVector<int64_t> laneLayout =
- consumerLayout.getEffectiveLaneLayoutAsInt();
+ SmallVector<int64_t> instDataA(aTy.getRank(), 1);
+ instDataA[aTy.getRank() - 2] = maxALen;
+ instDataA[aTy.getRank() - 1] = kDimSize;
+ SmallVector<int64_t> instDataB(bTy.getRank(), 1);
+ instDataB[bTy.getRank() - 2] = kDimSize;
+ instDataB[bTy.getRank() - 1] = maxBLen;
+ SmallVector<int64_t> instDataCD(cdTy.getRank(), 1);
+ instDataCD[cdTy.getRank() - 2] = maxALen;
+ instDataCD[cdTy.getRank() - 1] = maxCLen;
+ return std::make_tuple(instDataA, instDataB, instDataCD);
+}
- assert(consumerLayout.getRank() == static_cast<int64_t>(srcShape.size()) &&
- "laneData must be available for all dimensions");
- size_t innerMostDim = srcShape.size() - 1;
- int64_t sgDataValue = -1;
- int64_t instDataValue = -1;
- int64_t laneDataValue = -1;
- if (srcElemTyBitWidth > resElemTyBitWidth) {
- // When casting to a smaller bitwidth, multiply the result layout
- // accordingly to ensure it can be divided by the ratio back to the
- // source layout.
- int bitWidthRatio = srcElemTyBitWidth / resElemTyBitWidth;
- if (layoutKind == xegpu::LayoutKind::Subgroup) {
- sgDataValue = sgData[innerMostDim];
- while ((sgDataValue <= resShape[innerMostDim]) &&
- (sgDataValue % bitWidthRatio) != 0)
- sgDataValue *= 2;
- } else if (layoutKind == xegpu::LayoutKind::InstData) {
- instDataValue = instData[innerMostDim];
- const int innermostDimLaneLayout = laneLayout.empty()
- ? uArch->getSubgroupSize()
- : laneLayout[innerMostDim];
- // Adjust instDataValue so it still fits within an instruction after
- // dividing by bitWidthRatio
- while ((instDataValue <= resShape[innerMostDim]) &&
- (instDataValue % (innermostDimLaneLayout * bitWidthRatio) != 0))
- instDataValue *= 2;
- assert((resShape[innerMostDim] % instDataValue) == 0 &&
- "resShape, instData, and lanelayout for innermost must be 2^n !");
- } else if (layoutKind == xegpu::LayoutKind::Lane) {
- laneDataValue = laneData[innerMostDim];
- while ((laneDataValue <= resShape[innerMostDim]) &&
- (laneDataValue % bitWidthRatio != 0))
- laneDataValue *= 2;
- }
- // Now set only instData and laneData, preserving sgData
- xegpu::DistributeLayoutAttr resLayout;
- resLayout = consumerLayout.setDimData(innerMostDim, sgDataValue,
- instDataValue, laneDataValue);
- return resLayout;
+/// Computes lane_layout and lane_data for scatter-style store anchor layouts
+/// (store scatter, store matrix). Lanes and the per-lane vector both live on
+/// the innermost dim:
+/// - laneLayout[innermost] = min(subgroupSize, srcShape[innermost])
+/// - laneData[innermost] = min(srcShape[innermost] / laneLayout[innermost],
+/// maxChunkSize)
+/// All other entries are 1.
+static std::pair<SmallVector<int64_t>, SmallVector<int64_t>>
+computeScatterIOLaneLayoutAndData(ArrayRef<int64_t> instShape,
+ int64_t subgroupSize, int64_t maxChunkSize) {
+ int64_t rank = instShape.size();
+ SmallVector<int64_t> laneLayout(rank, 1), laneData(rank, 1);
+ int64_t innermost = rank - 1;
+ laneLayout[innermost] = std::min(subgroupSize, instShape[innermost]);
+ laneData[innermost] =
+ std::min(instShape[innermost] / laneLayout[innermost], maxChunkSize);
+ return {laneLayout, laneData};
+}
+
+// Computes the per-lane layout and data for a 2D block load/store/prefetch:
+// lanes are spread across the subgroup along the last dim (or rank-2 if
+// transposed), and laneData packs sub-bitwidth elements along the packing dim.
+static std::pair<SmallVector<int64_t>, SmallVector<int64_t>>
+compute2DBlockIOLaneLayoutAndData(ArrayRef<int64_t> instShape,
+ int64_t subgroupSize, int64_t bitwidth,
+ int64_t packingSize, bool transform = false) {
+ int64_t rank = instShape.size();
+ SmallVector<int64_t> laneLayout(rank, 1), laneData(rank, 1);
+ int kDim = transform ? rank - 2 : rank - 1;
+ unsigned vnniFactor = packingSize / bitwidth;
+ laneData[kDim] = bitwidth < packingSize ? vnniFactor : 1;
+ laneLayout.back() =
+ std::min(subgroupSize, instShape.back() / laneData.back());
+
+ // assert that the lane layout and data fit in the inst shape
+ for (int64_t i = 0; i < rank; ++i) {
+ int64_t laneProduct = laneLayout[i] * laneData[i];
+ assert(instShape[i] % laneProduct == 0 &&
+ "lane_layout * lane_data must evenly divide the inst shape");
+ (void)laneProduct;
}
- return consumerLayout;
+ return {laneLayout, laneData};
}
-/// Sets up the result layout for an interleave operation to ensure the source
-/// layout can be safely derived. Interleave doubles the innermost dimension,
-/// so the result layout must ensure that laneData is a multiple
-/// of 2, and instData must be divisible by innermostDimLaneLayout * 2.
+/// Computes the (lane_layout, lane_data) for a multi-reduction's source layout.
+/// Only the innermost two dims are distributed; leading dims are assumed unit.
+/// `subgroupSize` lanes go on one dim; up to `maxReduceVectorSize` elements are
+/// packed into lane_data on the other. To minimize cross-lane reduction, lanes
+/// are spread across a non-reduction dim when possible so the reduction happens
+/// within a lane. inst_data is the element-wise product lane_layout *
+/// lane_data.
///
-/// Example:
-/// Interleave: vector<128x256xf4> -> vector<128x512xf4>
-/// Consumer layout: laneLayout=[1, 16], laneData=[1, 4], instData=[1, 64]
-/// Result layout adjustment to ensure source can be safely inferred:
-/// - laneData must be >= 2 and multiple of 2 (so source = laneData/2 is
-/// valid)
-/// - instData must be divisible by (16 * 2 = 32) (so source = instData/2 is
-/// valid)
-/// - Adjusted instData: ensure (instData % 32 == 0)
-///
-xegpu::DistributeLayoutAttr xegpu::setupInterleaveResultLayout(
- xegpu::LayoutKind layoutKind, VectorType srcVecTy, VectorType resVecTy,
- DistributeLayoutAttr consumerLayout, const xegpu::uArch::uArch *uArch) {
+/// e.g. with srcShape=[32, 128], subgroupSize=16, maxReduceVectorSize=2:
+/// - Switch: reductionDims=[1] and consumerReductionDims=[] -> lanes move
+/// to the non-reduction dim 0: lane_layout=[16, 1], lane_data=[1, 2].
+/// - Default: reductionDims=[0, 1] (both reduced) -> lanes stay on the
+/// innermost dim: lane_layout=[1, 16], lane_data=[2, 1].
+static std::pair<SmallVector<int64_t>, SmallVector<int64_t>>
+computeReductionLaneLayoutAndData(ArrayRef<int64_t> srcShape,
+ ArrayRef<int64_t> reductionDims,
+ int subgroupSize, int64_t maxReduceVectorSize,
+ bool verticalLaneLayout = false) {
+ int srcRank = srcShape.size();
+ SmallVector<int64_t> laneLayout(srcRank, 1), laneData(srcRank, 1);
- ArrayRef<int64_t> srcShape = srcVecTy.getShape();
- SmallVector<int64_t> sgData = consumerLayout.getEffectiveSgDataAsInt();
- SmallVector<int64_t> instData = consumerLayout.getEffectiveInstDataAsInt();
- SmallVector<int64_t> laneData = consumerLayout.getEffectiveLaneDataAsInt();
- SmallVector<int64_t> laneLayout =
- consumerLayout.getEffectiveLaneLayoutAsInt();
+ int innermost = srcRank - 1;
+ int secondInnermost = srcRank - 2;
- assert(consumerLayout.getRank() == static_cast<int64_t>(srcShape.size()) &&
- "consumer layout rank must match source shape rank");
- const size_t innerMostDim = srcShape.size() - 1;
- int64_t sgDataValue = -1;
- int64_t instDataValue = -1;
- int64_t laneDataValue = -1;
+ if (verticalLaneLayout && secondInnermost >= 0) {
+ std::swap(innermost, secondInnermost);
+ }
+ int laneDim = innermost;
+ int vectorDim = secondInnermost; // negative for rank 1
- // Interleave doubles the innermost dimension (ratio = 2)
- constexpr int ratio = 2;
+ laneLayout[laneDim] =
+ std::min(static_cast<int64_t>(subgroupSize), srcShape[laneDim]);
+ if (vectorDim >= 0)
+ laneData[vectorDim] = std::min(maxReduceVectorSize, srcShape[vectorDim]);
+
+ return {laneLayout, laneData};
+}
+
+//===----------------------------------------------------------------------===//
+// Result/anchor-layout setup. Each op category derives lane_layout/lane_data
+// (and inst_data / sgData)
diff erently. Two things vary across ops:
+//
+// * Consumer dependence: consumer-driven ops prefer the layout requested by
+// their downstream uses and fall back to uArch defaults only when it is
+// absent/invalid; sinks (StoreNd, PrefetchNd) have no consumer and always
+// pick their own layout from uArch.
+//
+// * Derivation direction between inst_data and lane_layout/lane_data. Both
+// obey the invariant inst_data = k * lane_layout * lane_data, where `k` is
+// a per-dim integer >= 1 giving how many times each lane repeats its
+// access to cover one instruction's data tile (k == 1 means one lane
+// position per element; k > 1 means the instruction loads/stores several
+// elements per lane along that dim). Ops solve this invariant from
+// opposite ends:
+// - Rigid-lane ops (Nd block IO, DPAS): hardware fixes lane_layout /
+// lane_data first, then inst_data is built as a multiple of their
+// product (using get2DBlockIOInstDataLayout / getDpasInstDataLayouts).
+// - inst_data-first ops (scatter load): take inst_data from the consumer
+// and derive lane_layout/lane_data underneath it.
+//
+// - DPAS (+DPAS_MX) : rigid lanes — inst_data from HW block dims; A/B/C/D
+// lanes/data follow each operand's matmul role; DPAS_MX
+// additionally lays out the scale operand.
+// - LoadNd : consumer-driven, rigid lanes — honors the consumer's
+// inst_data / lane / sg_layout (incl. transpose & VNNI
+// packing) when it satisfies uArch block constraints,
+// else falls back to the default 2D-block scheme (lanes
+// on the last dim, rank-2 if transposed). The fallback
+// picks the LARGEST uArch block that divides the data
+// shape, so the resulting inst_data block can be bigger
+// than what the consumer asked for (fewer, wider
+// loads).
+// - StoreNd/PrefetchNd: data sinks, no consumer, rigid lanes — pick the
+// 2D-block layout directly from uArch (no VNNI
+// packing).
+// - Load (scatter) : load_gather / load_matrix, consumer-driven,
+// inst_data-first — reuse the consumer's inst_data and
+// derive lane_layout/lane_data, else default to lanes +
+// per-lane chunk on the innermost dim (chunk capped by
+// maxChunkSize).
+// - Store (scatter) : store_scatter / store_matrix — same scatter scheme,
+// but always self-derived from the scatter default.
+// - Reduction : (multi_)reduction, consumer-driven — distribute the
+// inner two dims, with lanes on the innermost dim by
+// default (reducing across lanes) and switched to a
+// non-reduction dim only when that keeps the reduction
+// within a lane. Reuses the consumer's slice layout
+// when it slices exactly the reduction dims, otherwise
+// re-derives. See setupMultiReductionResultLayout for
+// the exact switch condition and worked examples.
+// - BitCast/Interleave: scale the innermost data field by the bitwidth /
+// interleave ratio so the source layout divides back
+// out.
+// - InsertStridedSlice: clamp lane_data per dim to fit the inserted slice
+// (Lane kind only; sg/inst layouts unsupported).
+//===----------------------------------------------------------------------===//
+
+/// Helper function to set up subgroup layouts for DPAS operands A, B, and
+/// C/D. Compute subgroup layout candidates based on wgtile and instData, and
+/// then pick the best one that satisfies all operands and the consumer (if
+/// specified).
+static std::optional<
+ std::tuple<xegpu::DistributeLayoutAttr, xegpu::DistributeLayoutAttr,
+ xegpu::DistributeLayoutAttr>>
+getDpasSubgroupLayouts(
+ mlir::MLIRContext *context, VectorType aTy, VectorType bTy, VectorType cdTy,
+ xegpu::DistributeLayoutAttr consumerLayout, int numSg,
+ std::tuple<SmallVector<int64_t>, SmallVector<int64_t>, SmallVector<int64_t>>
+ instDataVecs) {
+ auto [instDataA, instDataB, instDataCD] = instDataVecs;
+
+ std::optional<LayoutRepresentation> consumerSgLayout = std::nullopt;
+ if (consumerLayout && consumerLayout.isForWorkgroup()) {
+ consumerSgLayout = consumerLayout.getEffectiveSgLayoutAsInt();
+ }
+
+ // Get all valid layouts for A, B and C/D operands
+ auto layoutsA = getSgLayoutCandidates(aTy.getShape(), instDataA, numSg);
+ auto layoutsB = getSgLayoutCandidates(bTy.getShape(), instDataB, numSg);
+ auto layoutsCD = getSgLayoutCandidates(cdTy.getShape(), instDataCD, numSg);
+ if (layoutsA.empty() || layoutsB.empty() || layoutsCD.empty())
+ return std::nullopt;
+
+ // Pick the best subgroup layout
+ std::optional<LayoutRepresentation> bestPick;
+ auto checkAlignedSgDataAB = [&](const LayoutRepresentation &sgLayout) {
+ return aTy.getShape().back() / sgLayout[1] ==
+ bTy.getShape().front() / sgLayout[0];
+ };
+ for (auto &sgLayout : layoutsB) {
+ if (llvm::is_contained(layoutsA, sgLayout) &&
+ llvm::is_contained(layoutsCD, sgLayout)) {
+ if (!checkAlignedSgDataAB(sgLayout))
+ continue;
+ // Is in (A and B and CD) and matches consumer -> best pick
+ if (consumerSgLayout.has_value() && sgLayout == *consumerSgLayout) {
+ bestPick = sgLayout;
+ break;
+ }
+ // Is in (A and B and CD) layoutsB is ordered from most
+ // balanced to least. So the first one we see is the most balanced one,
+ // remember it and later only update if there is one that matches the
+ // consumer.
+ if (!bestPick)
+ bestPick = sgLayout;
+ }
+ }
+ if (!bestPick)
+ return std::nullopt;
+
+ const auto &picked = *bestPick;
+
+ auto dpasALayout = buildSgLayout(context, aTy.getShape(), picked,
+ /*dimK=*/aTy.getRank() - 1);
+ auto dpasBLayout = buildSgLayout(context, bTy.getShape(), picked,
+ /*dimK=*/bTy.getRank() - 2);
+ auto dpasCDLayout = buildSgLayout(context, cdTy.getShape(), picked);
+ return std::make_tuple(dpasALayout, dpasBLayout, dpasCDLayout);
+}
+
+/// Sets up the anchor layouts for dpas operands (A, B, and C/D).
+/// The numSg and consumerLayout (optional) are only used by sg layout
+/// creation.
+std::optional<
+ std::tuple<xegpu::DistributeLayoutAttr, xegpu::DistributeLayoutAttr,
+ xegpu::DistributeLayoutAttr>>
+xegpu::setupDpasLayout(xegpu::LayoutKind layoutKind, VectorType aTy,
+ VectorType bTy, VectorType cdTy,
+ xegpu::DistributeLayoutAttr consumerLayout, int numSg,
+ const xegpu::uArch::uArch *uArch) {
+ auto context = aTy.getContext();
+ const auto *uArchInstruction =
+ dyn_cast<xegpu::uArch::SubgroupMatrixMultiplyAcc>(uArch->getInstruction(
+ xegpu::uArch::InstructionKind::SubgroupMatrixMultiplyAcc));
+ if (!uArchInstruction)
+ return std::nullopt;
+ auto subgroupSize = uArch->getSubgroupSize();
+
+ auto [laneLayoutA, laneDataA] = compute2DBlockIOLaneLayoutAndData(
+ aTy.getShape(), subgroupSize,
+ aTy.getElementType().getIntOrFloatBitWidth(),
+ uArchInstruction->getPackedFormatBitSizeA());
+ auto [laneLayoutB, laneDataB] = compute2DBlockIOLaneLayoutAndData(
+ bTy.getShape(), subgroupSize,
+ bTy.getElementType().getIntOrFloatBitWidth(),
+ uArchInstruction->getPackedFormatBitSizeB(), /*vnni=*/true);
+ auto [laneLayoutCD, laneDataCD] = compute2DBlockIOLaneLayoutAndData(
+ cdTy.getShape(), subgroupSize,
+ cdTy.getElementType().getIntOrFloatBitWidth(),
+ cdTy.getElementType().getIntOrFloatBitWidth());
+
+ auto instDataVecs = getDpasInstDataLayouts(aTy, bTy, cdTy, uArchInstruction);
+ if (!instDataVecs)
+ return std::nullopt;
if (layoutKind == xegpu::LayoutKind::Subgroup) {
- sgDataValue = sgData[innerMostDim];
- // Ensure sgDataValue is divisible by ratio so source sgData can be inferred
- while ((sgDataValue <= srcShape[innerMostDim]) &&
- (sgDataValue % ratio != 0))
- sgDataValue *= ratio;
+ assert(numSg > 0 &&
+ "Number of subgroups must be provided for sg layout creation.");
+ return getDpasSubgroupLayouts(context, aTy, bTy, cdTy, consumerLayout,
+ numSg, *instDataVecs);
} else if (layoutKind == xegpu::LayoutKind::InstData) {
- instDataValue = instData[innerMostDim];
- const int innermostDimLaneLayout = laneLayout.empty()
- ? uArch->getSubgroupSize()
- : laneLayout[innerMostDim];
- // Adjust instDataValue so it can be divided by (innermostDimLaneLayout *
- // ratio) when inferring the source layout
- while ((instDataValue <= srcShape[innerMostDim]) &&
- (instDataValue % (innermostDimLaneLayout * ratio) != 0))
- instDataValue *= ratio;
- assert((srcShape[innerMostDim] % instDataValue) == 0 &&
- "srcShape, instData, and laneLayout for innermost must be 2^n!");
+ auto [instDataA, instDataB, instDataCD] = *instDataVecs;
+ return std::make_tuple(
+ buildInstDataLayoutWithLane(context, instDataA, laneLayoutA, laneDataA),
+ buildInstDataLayoutWithLane(context, instDataB, laneLayoutB, laneDataB),
+ buildInstDataLayoutWithLane(context, instDataCD, laneLayoutCD,
+ laneDataCD));
} else if (layoutKind == xegpu::LayoutKind::Lane) {
- laneDataValue = laneData[innerMostDim];
- // Ensure laneDataValue is at least 2 and divisible by ratio
- // so that source laneData = laneDataValue/2 is valid
- while ((laneDataValue <= srcShape[innerMostDim]) &&
- (laneDataValue % ratio != 0))
- laneDataValue *= ratio;
+ auto aLayout = buildLaneLayout(context, laneLayoutA, laneDataA);
+ auto bLayout = buildLaneLayout(context, laneLayoutB, laneDataB);
+ auto cdLayout = buildLaneLayout(context, laneLayoutCD, laneDataCD);
+ return std::make_tuple(aLayout, bLayout, cdLayout);
}
+ return std::nullopt;
+}
- return consumerLayout.setDimData(innerMostDim, sgDataValue, instDataValue,
- laneDataValue);
+/// Helper to create a scale layout derived from a matrix operand layout.
+/// The scale layout is computed by mapping each dimension of the matrix
+/// layout to the corresponding scale tensor dimension using the ratio
+/// between the matrix and scale shapes.
+static xegpu::DistributeLayoutAttr
+createScaleLayout(mlir::MLIRContext *context, VectorType matrixTy,
+ VectorType scaleTy, xegpu::DistributeLayoutAttr matrixLayout,
+ bool isBScale, const xegpu::uArch::uArch *uArch) {
+ if (!scaleTy || !matrixLayout)
+ return nullptr;
+
+ // Calculate scaling factor by dividing matrix shape by scale shape
+ ArrayRef<int64_t> matrixShape = matrixTy.getShape();
+ ArrayRef<int64_t> scaleShape = scaleTy.getShape();
+
+ // Scale shapes can be 1D or 2D, handle both cases
+ if (scaleShape.empty())
+ return nullptr;
+
+ auto uArchInstruction =
+ dyn_cast<xegpu::uArch::SubgroupScaledMatrixMultiplyAcc>(
+ uArch->getInstruction(
+ xegpu::uArch::InstructionKind::SubgroupScaledMatrixMultiplyAcc));
+
+ int64_t rank = matrixLayout.getRank();
+ assert(rank >= 2 && "dpas layouts must be at least two dimensions");
+
+ SmallVector<int64_t> sgLayout = matrixLayout.getEffectiveSgLayoutAsInt();
+ SmallVector<int64_t> sgData = matrixLayout.getEffectiveSgDataAsInt();
+ SmallVector<int64_t> instData = matrixLayout.getEffectiveInstDataAsInt();
+ SmallVector<int64_t> laneLayout = matrixLayout.getEffectiveLaneLayoutAsInt();
+ SmallVector<int64_t> laneData = matrixLayout.getEffectiveLaneDataAsInt();
+ auto order = matrixLayout.getOrder();
+
+ SmallVector<int64_t> scaleSgLayout;
+ SmallVector<int64_t> scaleSgData;
+ if (!sgLayout.empty() && !sgData.empty()) {
+ scaleSgLayout.assign(sgLayout.begin(), sgLayout.end());
+ scaleSgData.assign(sgData.begin(), sgData.end());
+ scaleSgData[rank - 2] = std::max<int64_t>(
+ scaleShape[rank - 2] / (matrixShape[rank - 2] / sgData[rank - 2]), 1);
+ scaleSgData[rank - 1] = std::max<int64_t>(
+ scaleShape[rank - 1] / (matrixShape[rank - 1] / sgData[rank - 1]), 1);
+ }
+
+ // For DPAS_MX scales: if matrix has inst_data, scale needs adjusted
+ // inst_data. Scale inst_data is derived from matrix inst_data divided by
+ // scale factor.
+ SmallVector<int64_t> scaleInstData;
+ if (!instData.empty()) {
+ scaleInstData.assign(instData.begin(), instData.end());
+ if (isBScale)
+ scaleInstData[rank - 2] = std::max<int64_t>(
+ scaleShape[rank - 2] / (matrixShape[rank - 2] / instData[rank - 2]),
+ 1);
+ else
+ scaleInstData[rank - 1] = std::max<int64_t>(
+ scaleShape[rank - 1] / (matrixShape[rank - 1] / instData[rank - 1]),
+ 1);
+ }
+
+ SmallVector<int64_t> scaleLaneLayout;
+ SmallVector<int64_t> scaleLaneData;
+ if (!laneLayout.empty() && !laneData.empty()) {
+ scaleLaneLayout.assign(laneLayout.begin(), laneLayout.end());
+ scaleLaneData.assign(laneData.size(), 1);
+
+ bool isRowMajor = uArchInstruction->isLaneLayoutRowMajorOrder();
+ if (isBScale ^ isRowMajor)
+ std::swap(scaleLaneLayout[rank - 2], scaleLaneLayout[rank - 1]);
+ // Cap lane_layout by the per-instruction tile (inst_data) on each dim.
+ // Then derive lane_data = inst_data / lane_layout so the Category A
+ // invariant inst_data = lane_layout * lane_data * k (with k = 1) holds
+ // for the scale operand's load_nd consumer.
+ auto layoutCap = scaleInstData.empty() ? scaleShape : scaleInstData;
+ for (int64_t d = rank - 2; d < rank; ++d)
+ scaleLaneLayout[d] = std::min<int64_t>(layoutCap[d], scaleLaneLayout[d]);
+ }
+ return buildLayout(context, scaleSgLayout, scaleSgData, scaleInstData,
+ scaleLaneLayout, scaleLaneData, order);
}
-/// Sets up the result layout for an insert strided slice operation.
-/// Creates a result layout based on the specified layout kind (InstData or
-/// Lane).
-xegpu::DistributeLayoutAttr xegpu::setupInsertStridedSliceResultLayout(
- xegpu::LayoutKind layoutKind, VectorType srcVectorTy,
- VectorType resVectorTy, xegpu::DistributeLayoutAttr consumerLayout,
- const xegpu::uArch::uArch *uArch) {
+/// Sets up the anchor layouts for dpas_mx operands (A, B, C/D, A_scale, and
+/// B_scale). The numSg and consumerLayout (optional) are only used by sg
+/// layout creation.
+std::optional<
+ std::tuple<xegpu::DistributeLayoutAttr, xegpu::DistributeLayoutAttr,
+ xegpu::DistributeLayoutAttr, xegpu::DistributeLayoutAttr,
+ xegpu::DistributeLayoutAttr>>
+xegpu::setupDpasMxLayout(xegpu::LayoutKind layoutKind, VectorType aTy,
+ VectorType bTy, VectorType cdTy, VectorType aScaleTy,
+ VectorType bScaleTy,
+ xegpu::DistributeLayoutAttr consumerLayout, int numSg,
+ const xegpu::uArch::uArch *uArch) {
+ auto context = aTy.getContext();
+ const auto *uArchInstruction =
+ dyn_cast<xegpu::uArch::SubgroupMatrixMultiplyAcc>(uArch->getInstruction(
+ xegpu::uArch::InstructionKind::SubgroupMatrixMultiplyAcc));
+ if (!uArchInstruction)
+ return std::nullopt;
+ auto subgroupSize = uArch->getSubgroupSize();
- xegpu::DistributeLayoutAttr requiredResLayout;
- SmallVector<int64_t> consumerInstData =
- consumerLayout.getEffectiveInstDataAsInt();
- SmallVector<int64_t> consumerLaneData =
- consumerLayout.getEffectiveLaneDataAsInt();
- SmallVector<int64_t> consumerLaneLayout =
- consumerLayout.getEffectiveLaneLayoutAsInt();
- ArrayRef<int64_t> srcShape = srcVectorTy.getShape();
- int64_t instDataValue = -1;
- int64_t laneDataValue = -1;
+ auto [laneLayoutA, laneDataA] = compute2DBlockIOLaneLayoutAndData(
+ aTy.getShape(), subgroupSize,
+ aTy.getElementType().getIntOrFloatBitWidth(),
+ uArchInstruction->getPackedFormatBitSizeA());
+ auto [laneLayoutB, laneDataB] = compute2DBlockIOLaneLayoutAndData(
+ bTy.getShape(), subgroupSize,
+ bTy.getElementType().getIntOrFloatBitWidth(),
+ uArchInstruction->getPackedFormatBitSizeB(), /*vnni=*/true);
+ auto [laneLayoutCD, laneDataCD] = compute2DBlockIOLaneLayoutAndData(
+ cdTy.getShape(), subgroupSize,
+ cdTy.getElementType().getIntOrFloatBitWidth(),
+ cdTy.getElementType().getIntOrFloatBitWidth());
+ auto instDataVecs = getDpasInstDataLayouts(aTy, bTy, cdTy, uArchInstruction);
+ if (!instDataVecs)
+ return std::nullopt;
- requiredResLayout = consumerLayout;
- int srcRank = srcShape.size();
+ if (layoutKind == xegpu::LayoutKind::Subgroup) {
+ assert(numSg > 0 &&
+ "Number of subgroups must be provided for sg layout creation.");
+ auto dpasLayouts = getDpasSubgroupLayouts(
+ context, aTy, bTy, cdTy, consumerLayout, numSg, *instDataVecs);
+ if (!dpasLayouts)
+ return std::nullopt;
+
+ auto [dpasALayout, dpasBLayout, dpasCDLayout] = *dpasLayouts;
+
+ // Create scale layouts
+ auto aScaleLayout =
+ createScaleLayout(context, aTy, aScaleTy, dpasALayout, false, uArch);
+
+ auto bScaleLayout =
+ createScaleLayout(context, bTy, bScaleTy, dpasBLayout, true, uArch);
+
+ return std::make_tuple(dpasALayout, dpasBLayout, dpasCDLayout, aScaleLayout,
+ bScaleLayout);
+ } else if (layoutKind == xegpu::LayoutKind::InstData) {
+
+ auto [instDataA, instDataB, instDataCD] = *instDataVecs;
+
+ auto dpasALayout =
+ buildInstDataLayoutWithLane(context, instDataA, laneLayoutA, laneDataA);
+ auto dpasBLayout =
+ buildInstDataLayoutWithLane(context, instDataB, laneLayoutB, laneDataB);
+ auto dpasCDLayout = buildInstDataLayoutWithLane(context, instDataCD,
+ laneLayoutCD, laneDataCD);
+
+ auto aScaleLayout =
+ createScaleLayout(context, aTy, aScaleTy, dpasALayout, false, uArch);
+ auto bScaleLayout =
+ createScaleLayout(context, bTy, bScaleTy, dpasBLayout, true, uArch);
+
+ return std::make_tuple(dpasALayout, dpasBLayout, dpasCDLayout, aScaleLayout,
+ bScaleLayout);
+ } else if (layoutKind == xegpu::LayoutKind::Lane) {
+ auto dpasALayout = buildLaneLayout(context, laneLayoutA, laneDataA);
+ auto dpasBLayout = buildLaneLayout(context, laneLayoutB, laneDataB);
+ auto dpasCDLayout = buildLaneLayout(context, laneLayoutCD, laneDataCD);
+
+ auto aScaleLayout =
+ createScaleLayout(context, aTy, aScaleTy, dpasALayout, false, uArch);
+ auto bScaleLayout =
+ createScaleLayout(context, bTy, bScaleTy, dpasBLayout, true, uArch);
+
+ return std::make_tuple(dpasALayout, dpasBLayout, dpasCDLayout, aScaleLayout,
+ bScaleLayout);
+ }
+ return std::nullopt;
+}
+
+/// Sets up the anchor layout for a store_nd operation. StoreNd picks its
+/// own layout based on uArch block parameters (it does not take a consumer
+/// layout, since it is a data sink).
+xegpu::DistributeLayoutAttr
+xegpu::setupStoreNdAnchorLayout(xegpu::LayoutKind layoutKind,
+ VectorType srcVecTy, int numSg,
+ const xegpu::uArch::uArch *uArch) {
+ const auto *uArchInstruction =
+ dyn_cast<xegpu::uArch::Subgroup2DBlockStoreInstruction>(
+ uArch->getInstruction(
+ xegpu::uArch::InstructionKind::Subgroup2DBlockStore));
+ if (!uArchInstruction)
+ return nullptr;
+
+ auto context = srcVecTy.getContext();
+ Type elemTy = srcVecTy.getElementType();
+ auto subgroupSize = uArch->getSubgroupSize();
+ auto dataShape = srcVecTy.getShape();
+ int rank = srcVecTy.getRank();
+ assert(rank >= 2 && "Expected at least 2D shape for ND op");
+
+ // Compute the default 2D block IO lane layout / lane data.
+ unsigned bitwidth = elemTy.getIntOrFloatBitWidth();
+ auto [laneLayout, laneData] = compute2DBlockIOLaneLayoutAndData(
+ dataShape, subgroupSize, bitwidth,
+ uArchInstruction->getPackedFormatBitSize());
+
+ if (layoutKind == xegpu::LayoutKind::Lane)
+ return buildLaneLayout(context, laneLayout, laneData);
+
+ auto instData =
+ get2DBlockIOInstDataLayout(dataShape, elemTy, uArchInstruction);
+
+ if (layoutKind == xegpu::LayoutKind::InstData) {
+ assert(instData && isValidLaneLayout(*instData, laneLayout, laneData) &&
+ "Expected the store layout to satisfy uArch block constraints");
+ return buildInstDataLayoutWithLane(context, *instData, laneLayout,
+ laneData);
+ }
+
+ if (layoutKind == xegpu::LayoutKind::Subgroup) {
+ assert(numSg > 0 &&
+ "Number of subgroups must be provided for sg layout creation.");
+ auto sgLayouts = getSgLayoutCandidates(dataShape, *instData, numSg);
+ if (sgLayouts.empty())
+ return nullptr;
+ return buildSgLayout(context, dataShape, sgLayouts.front(), /*dimK=*/-1);
+ }
+
+ return nullptr;
+}
+
+/// Sets up the anchor layout for a prefetch_nd operation. PrefetchNd has no
+/// consumer (it produces no value), so it picks its own layout from uArch
+/// block parameters.
+xegpu::DistributeLayoutAttr
+xegpu::setupPrefetchNdAnchorLayout(xegpu::LayoutKind layoutKind,
+ xegpu::TensorDescType tdescTy, int numSg,
+ const xegpu::uArch::uArch *uArch) {
+
+ const auto *uArchInstruction =
+ dyn_cast<xegpu::uArch::Subgroup2DBlockPrefetchInstruction>(
+ uArch->getInstruction(
+ xegpu::uArch::InstructionKind::Subgroup2DBlockPrefetch));
+ if (!uArchInstruction)
+ return nullptr;
+
+ auto context = tdescTy.getContext();
+ Type elemTy = tdescTy.getElementType();
+ auto subgroupSize = uArch->getSubgroupSize();
+ auto dataShape = tdescTy.getShape();
+ int rank = tdescTy.getRank();
+ assert(rank >= 2 && "Expected at least 2D shape for ND op");
+
+ // Compute the default 2D block IO lane layout / lane data.
+ unsigned bitwidth = elemTy.getIntOrFloatBitWidth();
+ auto [laneLayout, laneData] = compute2DBlockIOLaneLayoutAndData(
+ dataShape, subgroupSize, bitwidth,
+ uArchInstruction->getPackedFormatBitSize());
+
+ if (layoutKind == xegpu::LayoutKind::Lane)
+ return buildLaneLayout(context, laneLayout, laneData);
+
+ auto instData =
+ get2DBlockIOInstDataLayout(dataShape, elemTy, uArchInstruction);
+
+ if (layoutKind == xegpu::LayoutKind::InstData) {
+ assert(instData && isValidLaneLayout(*instData, laneLayout, laneData) &&
+ "Expected the prefetch layout to satisfy uArch block constraints");
+ return buildInstDataLayoutWithLane(context, *instData, laneLayout,
+ laneData);
+ }
+
+ if (layoutKind == xegpu::LayoutKind::Subgroup) {
+ assert(numSg > 0 &&
+ "Number of subgroups must be provided for sg layout creation.");
+ auto sgLayouts = getSgLayoutCandidates(dataShape, *instData, numSg);
+ if (sgLayouts.empty())
+ return nullptr;
+ return buildSgLayout(context, dataShape, sgLayouts.front(), /*dimK=*/-1);
+ }
+
+ return nullptr;
+}
+
+/// Sets up the anchor layout for a load_nd operation. LoadNd takes a
+/// consumer layout (from its result's downstream uses) and validates it
+/// against uArch constraints; if valid, the consumer's `inst_data` /
+/// `sg_layout` are honored. Otherwise the helper falls back to defaults
+/// derived from uArch block parameters.
+xegpu::DistributeLayoutAttr
+xegpu::setupLoadNdAnchorLayout(xegpu::LayoutKind layoutKind,
+ VectorType resVecTy,
+ xegpu::DistributeLayoutAttr consumerLayout,
+ int numSg, const xegpu::uArch::uArch *uArch) {
+
+ assert(consumerLayout && "Expected a valid consumer layout");
+ if (layoutKind == xegpu::LayoutKind::Subgroup) {
+ assert(consumerLayout.isForWorkgroup() &&
+ "Expected consumer layout to be a complete workgroup-level layout");
+ return consumerLayout;
+ }
+
+ auto context = resVecTy.getContext();
+ Type elemTy = resVecTy.getElementType();
+ auto subgroupSize = uArch->getSubgroupSize();
+ auto dataShape = resVecTy.getShape();
+ const auto *uArchInstruction =
+ dyn_cast<xegpu::uArch::Subgroup2DBlockLoadInstruction>(
+ uArch->getInstruction(
+ xegpu::uArch::InstructionKind::Subgroup2DBlockLoad));
+ if (!uArchInstruction)
+ return nullptr;
- if (layoutKind == xegpu::LayoutKind::Subgroup) {
- assert(true &&
- "subgroup layout assignment not supported for insertStridedSlice.");
- } else if (layoutKind == xegpu::LayoutKind::InstData) {
- for (int dim = 0; dim < srcRank; dim++) {
- instDataValue = std::min(srcShape[dim], consumerInstData[dim]);
- requiredResLayout =
- requiredResLayout.setDimData(dim, -1, instDataValue, -1);
+ int rank = resVecTy.getRank();
+ SmallVector<int64_t> consumerInstData =
+ consumerLayout.getEffectiveInstDataAsInt();
+ SmallVector<int64_t> consumerLaneLayout =
+ consumerLayout.getEffectiveLaneLayoutAsInt();
+ SmallVector<int64_t> consumerLaneData =
+ consumerLayout.getEffectiveLaneDataAsInt();
+ auto consumerOrderAttr = consumerLayout.getOrder();
+
+ assert(!consumerLaneLayout.empty() && !consumerLaneData.empty() &&
+ "Expected consumer layout to have lane_layout and lane_data");
+
+ // vertical lane layout means that the blockload must be transposed
+ // note scaleA on PVC has vertical lane layout even without transposed order
+ // attr
+ bool hasTranspose =
+ consumerLaneLayout[rank - 2] > 1 && consumerLaneLayout[rank - 1] == 1;
+ bool hasTransform = !hasTranspose && consumerLaneData[rank - 2] > 1 &&
+ consumerLaneData[rank - 1] == 1;
+ assert((consumerLaneData[rank - 2] == 1 || consumerLaneData[rank - 1] == 1) &&
+ "Expected consumer lane data to have at most one non-unit dim");
+
+ if (layoutKind == xegpu::LayoutKind::InstData) {
+ auto blockWHC = uArchInstruction->getBlockWidthHeightCount(
+ elemTy, hasTransform, hasTranspose,
+ /*upConv=*/false);
+ if (!blockWHC)
+ return nullptr;
+ auto [bWidths, bHeights, bCounts] = blockWHC.value();
+
+ SmallVector<int64_t> laneLayout;
+ // set the laneLayout to use consumer's LaneLayout as base, but adjust its
+ // size to match the subgroupsize in case its original value is larger than
+ // 1
+ for (int i = 0; i < rank; i++) {
+ if (consumerLaneLayout[i] > 1)
+ laneLayout.push_back(std::max(static_cast<int64_t>(subgroupSize),
+ consumerLaneLayout[i]));
+ else
+ laneLayout.push_back(1);
}
- } else if (layoutKind == xegpu::LayoutKind::Lane) {
- for (int dim = 0; dim < srcRank; dim++) {
- assert(srcShape[dim] % consumerLaneLayout[dim] == 0 &&
- "srcShape must be divisible by laneLayout for all dimensions");
- laneDataValue = std::min(srcShape[dim] / consumerLaneLayout[dim],
- consumerLaneData[dim]);
- requiredResLayout =
- requiredResLayout.setDimData(dim, -1, -1, laneDataValue);
+
+ // See whether the consumer's inst_data satisfies the block constraints.
+ int64_t height = consumerInstData[rank - 2];
+ int64_t width = consumerInstData[rank - 1];
+ auto maxBlockCount = *llvm::max_element(bCounts);
+ auto maxWidth = *llvm::max_element(bWidths);
+ if (llvm::is_contained(bWidths, static_cast<int>(width)) ||
+ (width % maxWidth == 0 && width / maxWidth < maxBlockCount)) {
+ if (llvm::is_contained(bHeights, static_cast<int>(height))) {
+ return buildInstDataLayoutWithLane(context, consumerInstData,
+ laneLayout, consumerLaneData,
+ consumerOrderAttr);
+ }
}
+
+ // if consumer instData size too small, try the larger one. like DPAS_MX's
+ // scale is smaller than block load
+ auto instData = get2DBlockIOInstDataLayout(
+ dataShape, elemTy, uArchInstruction, hasTransform, hasTranspose);
+ // assert instData is valid against consumer layout since
+ // transform/transpose attribute are derived from consumer layout
+ assert(instData &&
+ isValidLaneLayout(*instData, laneLayout, consumerLaneData) &&
+ "Expected the load layout to satisfy uArch block constraints");
+ return buildInstDataLayoutWithLane(context, *instData, laneLayout,
+ consumerLaneData, consumerOrderAttr);
}
- return requiredResLayout;
+ if (layoutKind == xegpu::LayoutKind::Lane) {
+ assert(isValidLaneLayout(dataShape, consumerLaneLayout, consumerLaneData) &&
+ "Expected the lane layout to satisfy uArch block constraints");
+ return consumerLayout;
+ }
+ return nullptr;
}
/// Sets up the anchor layout for load gather and load matrix operation.
/// load matrix lowers to load gather and 1d block load. All of them share the
/// same layout setup logic.
+///
/// For Subgroup layout, uses the consumer layout directly.
-/// non-chunked loads (1D or 2D):
-/// InstData = {1, ..., min(consumer, maxLaneLoadSize * subgroupSize)}
-/// LaneLayout = {1, ..., subgroupSize}
-/// lane_data = {1, ..., min(consumer, maxLaneLoadSize)}
-/// chunked loads (2D only):
-/// InstData = {subgroupSize, min(consumer, maxLaneLoadSize)}
-/// LaneLayout = {subgroupSize, 1}
-/// lane_data={1,min(consumer, maxLaneLoadSize)}
+///
+/// For InstData layout, takes consumer's inst_data as-is. lane_layout and
+/// lane_data are taken from the consumer when present; otherwise the helper
+/// derives the standard scatter-style default (subgroupSize lanes on the
+/// innermost dim, per-lane vector capped by maxChunkSize).
+///
+/// For Lane layout, lane_layout/lane_data are taken from the consumer when
+/// present; otherwise derived from the same default.
static xegpu::DistributeLayoutAttr setupGenericLoadAnchorLayout(
xegpu::LayoutKind layoutKind, mlir::MLIRContext *context,
- xegpu::DistributeLayoutAttr consumerLayout, bool isChunkedLoad,
- int maxChunkSize, ArrayRef<int64_t> resShape, int subgroupSize) {
+ xegpu::DistributeLayoutAttr consumerLayout, int maxChunkSize,
+ ArrayRef<int64_t> resShape, int subgroupSize) {
if (layoutKind == xegpu::LayoutKind::Subgroup)
return consumerLayout;
SmallVector<int64_t> consumerInstData =
consumerLayout.getEffectiveInstDataAsInt();
+ SmallVector<int64_t> consumerLaneLayout =
+ consumerLayout.getEffectiveLaneLayoutAsInt();
SmallVector<int64_t> consumerLaneData =
consumerLayout.getEffectiveLaneDataAsInt();
- SmallVector<int> instData(resShape.size(), 1);
- SmallVector<int> laneLayout(resShape.size(), 1);
- SmallVector<int> laneData(resShape.size(), 1);
-
- if (!isChunkedLoad) {
- if (layoutKind == xegpu::LayoutKind::InstData) {
- instData.back() = std::min(static_cast<int>(consumerInstData.back()),
- maxChunkSize * subgroupSize);
- return xegpu::LayoutAttr::get(context, instData);
- } else if (layoutKind == xegpu::LayoutKind::Lane) {
- laneData.back() =
- std::min(static_cast<int>(consumerLaneData.back()), maxChunkSize);
- laneLayout.back() = std::min(static_cast<int64_t>(subgroupSize),
- resShape.back() / laneData.back());
- return xegpu::LayoutAttr::get(context, laneLayout, laneData);
- }
- } else {
- assert(resShape.size() == 2 && "Chunked Store must access 2D tensor tile.");
- if (layoutKind == xegpu::LayoutKind::InstData) {
- instData[0] = subgroupSize;
- instData[1] =
- std::min(static_cast<int>(consumerInstData[1]), maxChunkSize);
- return xegpu::LayoutAttr::get(context, instData);
- } else if (layoutKind == xegpu::LayoutKind::Lane) {
- laneLayout[0] = subgroupSize;
- laneData[1] =
- std::min(static_cast<int>(consumerLaneData[1]), maxChunkSize);
- return xegpu::LayoutAttr::get(context, laneLayout, laneData);
+ // Pick lane_layout / lane_data: prefer consumer's, fall back to the
+ // scatter-store default (subgroupSize lanes on innermost dim, per-lane
+ // vector capped by maxChunkSize).
+ SmallVector<int64_t> laneLayout;
+ SmallVector<int64_t> laneData;
+ assert(!consumerLaneLayout.empty() && !consumerLaneData.empty() &&
+ "Expected consumer layout to have lane_layout and lane_data");
+ laneLayout.assign(consumerLaneLayout.begin(), consumerLaneLayout.end());
+ laneData.assign(consumerLaneData.begin(), consumerLaneData.end());
+
+ if (layoutKind == xegpu::LayoutKind::InstData) {
+ // Take consumer's inst_data as-is. If the consumer doesn't have one,
+ // fall back to lane_layout * lane_data per dim.
+ SmallVector<int64_t> instData;
+ if (!consumerInstData.empty()) {
+ instData.assign(consumerInstData.begin(), consumerInstData.end());
+ } else {
+ instData.resize(resShape.size());
+ for (size_t i = 0; i < resShape.size(); ++i)
+ instData[i] = laneLayout[i] * laneData[i];
}
+ return buildInstDataLayoutWithLane(context, instData, laneLayout, laneData);
}
+ if (layoutKind == xegpu::LayoutKind::Lane)
+ return buildLaneLayout(context, laneLayout, laneData);
return nullptr;
}
/// Sets up the anchor layout for a load gather operation.
xegpu::DistributeLayoutAttr xegpu::setupLoadGatherAnchorLayout(
- xegpu::LayoutKind layoutKind, VectorType resVecTy, int chunkSize,
+ xegpu::LayoutKind layoutKind, VectorType resVecTy, int contigChunkSize,
xegpu::DistributeLayoutAttr consumerLayout, const uArch::uArch *uArch) {
const int subgroupSize = uArch->getSubgroupSize();
@@ -1240,18 +1714,18 @@ xegpu::DistributeLayoutAttr xegpu::setupLoadGatherAnchorLayout(
const auto *uArchInstruction =
dyn_cast<xegpu::uArch::LoadGatherInstructionInterface>(
uArch->getInstruction(xegpu::uArch::InstructionKind::LoadGather));
- int maxChunkSize = uArchInstruction->getMaxLaneLoadSize(elemBitWidth);
+ int maxChunkSize = std::min(
+ uArchInstruction->getMaxLaneLoadSize(elemBitWidth), contigChunkSize);
return setupGenericLoadAnchorLayout(layoutKind, context, consumerLayout,
- (chunkSize > 1), maxChunkSize, resShape,
- subgroupSize);
+ maxChunkSize, resShape, subgroupSize);
}
/// Sets up the anchor layout for load matrix operation.
/// TODO: enhance load matrix to indicate lowering to chunked load or not.
xegpu::DistributeLayoutAttr
xegpu::setupLoadMatrixAnchorLayout(xegpu::LayoutKind layoutKind,
- VectorType resVecTy,
+ VectorType resVecTy, int contigChunkSize,
xegpu::DistributeLayoutAttr consumerLayout,
const xegpu::uArch::uArch *uArch) {
@@ -1263,61 +1737,41 @@ xegpu::setupLoadMatrixAnchorLayout(xegpu::LayoutKind layoutKind,
const auto *uArchInstruction =
dyn_cast<xegpu::uArch::LoadGatherInstructionInterface>(
uArch->getInstruction(xegpu::uArch::InstructionKind::LoadGather));
- int maxChunkSize = uArchInstruction->getMaxLaneLoadSize(elemBitWidth);
+ int maxChunkSize = std::min(
+ uArchInstruction->getMaxLaneLoadSize(elemBitWidth), contigChunkSize);
return setupGenericLoadAnchorLayout(layoutKind, context, consumerLayout,
- false, maxChunkSize, resShape,
- subgroupSize);
+ maxChunkSize, resShape, subgroupSize);
}
/// Sets up the anchor layout for store scatter and store matrix operation.
-/// store matrix lowers to store scatter and 1d block store. All of them share
-/// the same layout setup logic. For Subgroup layout, not supported yet.
-/// non-chunked stores (1D or 2D):
-/// InstData = {1, ..., subgroupSize}
-/// LaneLayout = {1, ..., subgroupSize}
-/// lane_data = {1, ..., 1}
-/// chunked stores (2D only):
-/// InstData = {subgroupSize, min(srcVec, maxLaneStoreSize)}
-/// LaneLayout = {subgroupSize, 1}
-/// lane_data={1,min(srcVec, maxLaneStoreSize)}
+/// store matrix lowers to store scatter and 1d block store. All of them
+/// share the same layout setup logic. For Subgroup layout, not supported
+/// yet.
+///
+/// Lane layout is derived first via `computeScatterIOLaneLayoutAndData`;
+/// inst_data is then the element-wise product lane_layout * lane_data.
static xegpu::DistributeLayoutAttr
setupGenericStoreAnchorLayout(xegpu::LayoutKind layoutKind,
- mlir::MLIRContext *context, bool isChunkedStore,
- int maxChunkSize, ArrayRef<int64_t> srcShape,
- int subgroupSize) {
-
- int srcShapeSize = srcShape.size();
- SmallVector<int> instData(srcShapeSize, 1);
- SmallVector<int> laneLayout(srcShapeSize, 1);
- SmallVector<int> laneData(srcShapeSize, 1);
+ mlir::MLIRContext *context, int maxChunkSize,
+ ArrayRef<int64_t> srcShape, int subgroupSize) {
if (layoutKind == xegpu::LayoutKind::Subgroup) {
- assert(true &&
+ assert(false &&
"subgroup layout assignment not supported for storeScatter.");
return nullptr;
}
- if (!isChunkedStore) {
- if (layoutKind == xegpu::LayoutKind::InstData) {
- instData[srcShapeSize - 1] =
- std::min(subgroupSize, static_cast<int>(srcShape.back()));
- return xegpu::LayoutAttr::get(context, instData);
- } else if (layoutKind == xegpu::LayoutKind::Lane) {
- laneLayout[srcShapeSize - 1] =
- std::min(subgroupSize, static_cast<int>(srcShape.back()));
- return xegpu::LayoutAttr::get(context, laneLayout, laneData);
- }
- } else {
- assert(srcShapeSize == 2 && "Chunked Store must access 2D tensor tile.");
- if (layoutKind == xegpu::LayoutKind::InstData) {
- instData[0] = subgroupSize;
- instData[1] = std::min(static_cast<int>(srcShape[1]), maxChunkSize);
- return xegpu::LayoutAttr::get(context, instData);
- } else if (layoutKind == xegpu::LayoutKind::Lane) {
- laneLayout[0] = subgroupSize;
- laneData[1] = std::min(static_cast<int>(srcShape[1]), maxChunkSize);
- return xegpu::LayoutAttr::get(context, laneLayout, laneData);
- }
+ auto [laneLayout, laneData] =
+ computeScatterIOLaneLayoutAndData(srcShape, subgroupSize, maxChunkSize);
+
+ if (layoutKind == xegpu::LayoutKind::InstData) {
+ SmallVector<int64_t> instData(srcShape.size());
+ for (size_t i = 0; i < srcShape.size(); ++i)
+ instData[i] = laneLayout[i] * laneData[i];
+ return buildInstDataLayoutWithLane(context, instData, laneLayout, laneData);
+ }
+ if (layoutKind == xegpu::LayoutKind::Lane) {
+ return buildLaneLayout(context, laneLayout, laneData);
}
return nullptr;
}
@@ -1325,7 +1779,7 @@ setupGenericStoreAnchorLayout(xegpu::LayoutKind layoutKind,
/// Sets up the anchor layout for a store scatter operation.
xegpu::DistributeLayoutAttr
xegpu::setupStoreScatterAnchorLayout(xegpu::LayoutKind layoutKind,
- VectorType srcVecTy, int chunkSize,
+ VectorType srcVecTy, int contigChunkSize,
const uArch::uArch *uArch) {
const int subgroupSize = uArch->getSubgroupSize();
@@ -1336,15 +1790,16 @@ xegpu::setupStoreScatterAnchorLayout(xegpu::LayoutKind layoutKind,
const auto *uArchInstruction =
dyn_cast<xegpu::uArch::StoreScatterInstructionInterface>(
uArch->getInstruction(xegpu::uArch::InstructionKind::StoreScatter));
- int maxChunkSize = uArchInstruction->getMaxLaneStoreSize(elemBitWidth);
- return setupGenericStoreAnchorLayout(layoutKind, context, (chunkSize > 1),
- maxChunkSize, srcShape, subgroupSize);
+ int maxChunkSize = std::min(
+ uArchInstruction->getMaxLaneStoreSize(elemBitWidth), contigChunkSize);
+ return setupGenericStoreAnchorLayout(layoutKind, context, maxChunkSize,
+ srcShape, subgroupSize);
}
/// Sets up the anchor layout for a store matrix operation.
xegpu::DistributeLayoutAttr
xegpu::setupStoreMatrixAnchorLayout(xegpu::LayoutKind layoutKind,
- VectorType srcVecTy,
+ VectorType srcVecTy, int contigChunkSize,
const xegpu::uArch::uArch *uArch) {
const int subgroupSize = uArch->getSubgroupSize();
@@ -1355,437 +1810,717 @@ xegpu::setupStoreMatrixAnchorLayout(xegpu::LayoutKind layoutKind,
const auto *uArchInstruction =
dyn_cast<xegpu::uArch::StoreScatterInstructionInterface>(
uArch->getInstruction(xegpu::uArch::InstructionKind::StoreScatter));
- int maxChunkSize = uArchInstruction->getMaxLaneStoreSize(elemBitWidth);
+ int maxChunkSize = std::min(
+ uArchInstruction->getMaxLaneStoreSize(elemBitWidth), contigChunkSize);
- return setupGenericStoreAnchorLayout(layoutKind, context, false, maxChunkSize,
+ return setupGenericStoreAnchorLayout(layoutKind, context, maxChunkSize,
srcShape, subgroupSize);
}
-// This function returns the default lane layout for a given vector type.
-// - `packingSize` means multiple consecutive elements can be accessed
-// together as a single unit.
-// - `vnni` means data packing is column-wise (i.e., 2x1xf16 with vnni vs.
-// 1x2xf16 w/o vnni).
-template <typename RankedTy>
-static xegpu::LayoutAttr getDefaultLaneLayout2DBlockIo(
- RankedTy ty, const xegpu::uArch::uArch *uArch,
- std::optional<unsigned> packingSize = std::nullopt, bool vnni = false) {
- // Expecting at least 1D vector. For rank > 2, leading dims are batch dims.
- assert(((ty.getRank() >= 1 && !vnni) || ty.getRank() >= 2) &&
- "Expected at least 1D non-vnni or 2D vector.");
- // Expecting int or float element type.
- assert(ty.getElementType().isIntOrFloat() &&
- "Expected int or float element type.");
-
- auto context = ty.getContext();
- auto rank = ty.getRank();
- SmallVector<int> laneLayout(rank, 1);
- SmallVector<int> laneData(rank, 1);
- if (packingSize.has_value()) {
- unsigned bitwidth = ty.getElementType().getIntOrFloatBitWidth();
- int &laneDataPos = vnni ? laneData[rank - 2] : laneData.back();
- laneDataPos = bitwidth < *packingSize ? *packingSize / bitwidth : 1;
+/// Completes a scatter IO layout by deriving lane_layout and lane_data from
+/// `specifiedLayout`'s inst_data when they are missing. The layout is returned
+/// unchanged if `specifiedLayout` is null, carries no inst_data, or already has
+/// both lane_layout and lane_data.
+///
+/// When lane info is absent, inst_data is treated as the effective shape and
+/// the lane factorization is filled in as follows:
+/// - If `consumerLayout` is present and its lane_layout / lane_data are a
+/// valid factorization of inst_data, that consumer lane info is reused so
+/// the completed layout matches the consumer (avoiding a relayout).
+/// - Otherwise a standard scatter-style factorization is computed via
+/// `computeScatterIOLaneLayoutAndData`, bounded by `maxChunkSize` — the
+/// per-lane load width reported by the uArch's LoadGather instruction
+/// (`getMaxLaneLoadSize`).
+///
+std::optional<xegpu::DistributeLayoutAttr>
+xegpu::completeScatterLoadLaneLayoutFromInstData(
+ xegpu::DistributeLayoutAttr specifiedLayout,
+ xegpu::DistributeLayoutAttr consumerLayout, Type elemTy,
+ const xegpu::uArch::LoadGatherInstructionInterface *uArchInstruction,
+ const int subgroupSize) {
+ if (!specifiedLayout)
+ return specifiedLayout;
+ SmallVector<int64_t> specifiedInstData =
+ specifiedLayout.getEffectiveInstDataAsInt();
+ if (specifiedInstData.empty())
+ return specifiedLayout;
+ if (!specifiedLayout.getEffectiveLaneLayoutAsInt().empty() &&
+ !specifiedLayout.getEffectiveLaneDataAsInt().empty())
+ return specifiedLayout;
+
+ // Reuse the load-side setup with inst_data as the destination shape.
+ auto *context = specifiedLayout.getContext();
+ auto elemBitWidth = elemTy.getIntOrFloatBitWidth();
+ int maxChunkSize = uArchInstruction->getMaxLaneLoadSize(elemBitWidth);
+ if (consumerLayout) {
+ auto consumerLaneLayout = consumerLayout.getEffectiveLaneLayoutAsInt();
+ auto consumerLaneData = consumerLayout.getEffectiveLaneDataAsInt();
+ if (!consumerLaneLayout.empty() && !consumerLaneData.empty() &&
+ isValidLaneLayout(specifiedInstData, consumerLaneLayout,
+ consumerLaneData))
+ return buildInstDataLayoutWithLane(context, specifiedInstData,
+ consumerLaneLayout, consumerLaneData);
}
- laneLayout.back() = uArch->getSubgroupSize();
- return xegpu::LayoutAttr::get(context, laneLayout, laneData);
+ auto [defLaneLayout, defLaneData] = computeScatterIOLaneLayoutAndData(
+ specifiedInstData, subgroupSize, maxChunkSize);
+ if (!isValidLaneLayout(specifiedInstData, defLaneLayout, defLaneData))
+ return std::nullopt;
+ return buildInstDataLayoutWithLane(context, specifiedInstData, defLaneLayout,
+ defLaneData);
}
-// This function returns all layouts for the given sgCount, whose sgData:
-// 1. Evenly divides the wgShape.
-// 2. Is a multiple of instData.
-// Example:
-// wgShape = [128, 64], instData = [8, 16], sgCount = 32
-// Returns layouts:
-// [(8,4), (16,2)], which correspond to sgData [16,16] and [8,32].
-using LayoutRepresentation = std::pair<int64_t, int64_t>;
-static SmallVector<LayoutRepresentation>
-getValidLayouts(ArrayRef<int64_t> wgShape, ArrayRef<int64_t> instData,
- int64_t sgCount) {
- SmallVector<LayoutRepresentation> candidates;
- for (int sgLayout0 = 1; sgLayout0 <= sgCount; ++sgLayout0) {
- if (sgCount % sgLayout0)
- continue;
- int64_t sgLayout1 = sgCount / sgLayout0;
- int64_t sgData0 = wgShape[0] / sgLayout0;
- int64_t sgData1 = wgShape[1] / sgLayout1;
- if ((wgShape[0] % sgLayout0 || wgShape[1] % sgLayout1) ||
- (sgData0 % instData[0] || sgData1 % instData[1]))
- continue;
- candidates.emplace_back(sgLayout0, sgLayout1);
- }
- // Sort primarily by how balanced they are
- // (i.e., minimize the absolute
diff erence between the two dimensions), and
- // secondarily by the first dimension in ascending order.
- llvm::sort(candidates, [](const LayoutRepresentation &lhs,
- const LayoutRepresentation &rhs) {
- int
diff Lhs = std::abs(lhs.first - lhs.second);
- int
diff Rhs = std::abs(rhs.first - rhs.second);
- if (
diff Lhs !=
diff Rhs)
- return
diff Lhs <
diff Rhs;
- return lhs.first < rhs.first;
- });
- return candidates;
+/// Like completeScatterLoadLaneLayoutFromInstData, but for scatter stores. A
+/// store is a data sink, so lane info is derived purely from inst_data (bounded
+/// by the uArch's per-lane store width); there is no consumer layout to reuse.
+std::optional<xegpu::DistributeLayoutAttr>
+xegpu::completeScatterStoreLaneLayoutFromInstData(
+ xegpu::DistributeLayoutAttr specifiedLayout, Type elemTy,
+ const xegpu::uArch::StoreScatterInstructionInterface *uArchInstruction,
+ const int subgroupSize) {
+ if (!specifiedLayout)
+ return specifiedLayout;
+ SmallVector<int64_t> specifiedInstData =
+ specifiedLayout.getEffectiveInstDataAsInt();
+ if (specifiedInstData.empty())
+ return specifiedLayout;
+ if (!specifiedLayout.getEffectiveLaneLayoutAsInt().empty() &&
+ !specifiedLayout.getEffectiveLaneDataAsInt().empty())
+ return specifiedLayout;
+
+ // Reuse the store-side setup with inst_data as the source shape.
+ auto *context = specifiedLayout.getContext();
+ auto elemBitWidth = elemTy.getIntOrFloatBitWidth();
+ int maxChunkSize = uArchInstruction->getMaxLaneStoreSize(elemBitWidth);
+ auto [defLaneLayout, defLaneData] = computeScatterIOLaneLayoutAndData(
+ specifiedInstData, subgroupSize, maxChunkSize);
+ if (!isValidLaneLayout(specifiedInstData, defLaneLayout, defLaneData))
+ return std::nullopt;
+ return buildInstDataLayoutWithLane(context, specifiedInstData, defLaneLayout,
+ defLaneData);
}
-/// Helper function to compute inst_data vectors for DPAS operands A, B, and
-/// C/D.
-static std::optional<std::tuple<SmallVector<int64_t>, SmallVector<int64_t>,
- SmallVector<int64_t>>>
-getDpasInstDataVectors(VectorType aTy, VectorType bTy, VectorType cdTy,
- const xegpu::uArch::uArch *uArch,
- bool isDpasMx = false) {
- const int subgroupSize = uArch->getSubgroupSize();
-
- const xegpu::uArch::MMAInstructionInterface *uArchInstruction;
- if (isDpasMx)
- uArchInstruction = dyn_cast<xegpu::uArch::SubgroupScaledMatrixMultiplyAcc>(
- uArch->getInstruction(
- xegpu::uArch::InstructionKind::SubgroupScaledMatrixMultiplyAcc));
- else
- uArchInstruction =
- dyn_cast<xegpu::uArch::SubgroupMatrixMultiplyAcc>(uArch->getInstruction(
- xegpu::uArch::InstructionKind::SubgroupMatrixMultiplyAcc));
-
- // M dimension is the second-to-last dim of A (handles batch dims).
- const unsigned dataALen = aTy.getShape()[aTy.getRank() - 2];
- auto supportedALen = uArchInstruction->getSupportedM(aTy.getElementType());
- const int maxALen =
- xegpu::getLargestDivisor(dataALen, ArrayRef<unsigned>(supportedALen));
-
- // N dimension is the last dim of B.
- const unsigned dataBLen = bTy.getShape().back();
- auto supportedBLen = uArchInstruction->getSupportedN(bTy.getElementType());
- const int maxBLen =
- xegpu::getLargestDivisor(dataBLen, ArrayRef<unsigned>(supportedBLen));
-
- auto supportedCLen = uArchInstruction->getSupportedN(cdTy.getElementType());
- const int maxCLen =
- xegpu::getLargestDivisor(dataBLen, ArrayRef<unsigned>(supportedCLen));
- if (maxALen == -1 || maxBLen == -1 || maxCLen == -1)
+/// Completes a 2D-block store/prefetch layout from its inst_data. store_nd and
+/// prefetch_nd are data sinks, so lane info is derived purely from inst_data
+/// (no consumer to reuse). One helper serves both via
+/// BlockIOInstructionInterface.
+std::optional<xegpu::DistributeLayoutAttr>
+xegpu::completeBlockStoreLaneLayoutFromInstData(
+ xegpu::DistributeLayoutAttr specifiedLayout, Type elemTy,
+ const xegpu::uArch::BlockIOInstructionInterface *uArchInstruction,
+ const int subgroupSize) {
+ if (!specifiedLayout)
+ return specifiedLayout;
+ SmallVector<int64_t> specifiedInstData =
+ specifiedLayout.getEffectiveInstDataAsInt();
+ if (specifiedInstData.empty())
+ return specifiedLayout;
+ if (!specifiedLayout.getEffectiveLaneLayoutAsInt().empty() &&
+ !specifiedLayout.getEffectiveLaneDataAsInt().empty())
+ return specifiedLayout;
+
+ auto *context = specifiedLayout.getContext();
+ auto [laneLayout, laneData] = compute2DBlockIOLaneLayoutAndData(
+ specifiedInstData, subgroupSize, elemTy.getIntOrFloatBitWidth(),
+ uArchInstruction->getPackedFormatBitSize());
+ if (!isValidLaneLayout(specifiedInstData, laneLayout, laneData))
return std::nullopt;
+ return buildInstDataLayoutWithLane(context, specifiedInstData, laneLayout,
+ laneData);
+}
- // For DPAS_MX, use getSupportedK to get the scaled K dimension.
- // assume single element in the returned vector.
- int kDimSize = subgroupSize;
- if (isDpasMx) {
- auto supportedKLen = uArchInstruction->getSupportedK(aTy.getElementType());
- if (supportedKLen.empty())
- return std::nullopt;
- kDimSize = supportedKLen[0];
+/// Like completeBlockStoreLaneLayoutFromInstData, but for load_nd. The
+/// consumer's lane_data and order are reused as-is; lane_layout is rebuilt from
+/// the consumer's lane_layout, bumping every non-unit dim up to the subgroup
+/// size. The user-provided inst_data is preserved.
+std::optional<xegpu::DistributeLayoutAttr>
+xegpu::completeBlockLoadLaneLayoutFromInstData(
+ xegpu::DistributeLayoutAttr specifiedLayout,
+ xegpu::DistributeLayoutAttr consumerLayout, Type elemTy,
+ const xegpu::uArch::BlockIOInstructionInterface *uArchInstruction,
+ const int subgroupSize) {
+ if (!specifiedLayout)
+ return specifiedLayout;
+ SmallVector<int64_t> specifiedInstData =
+ specifiedLayout.getEffectiveInstDataAsInt();
+ if (specifiedInstData.empty())
+ return specifiedLayout;
+ if (!specifiedLayout.getEffectiveLaneLayoutAsInt().empty() &&
+ !specifiedLayout.getEffectiveLaneDataAsInt().empty())
+ return specifiedLayout;
+ if (!consumerLayout)
+ return specifiedLayout;
+ SmallVector<int64_t> consumerLaneLayout =
+ consumerLayout.getEffectiveLaneLayoutAsInt();
+ SmallVector<int64_t> consumerLaneData =
+ consumerLayout.getEffectiveLaneDataAsInt();
+ if (consumerLaneLayout.empty() || consumerLaneData.empty())
+ return specifiedLayout;
+
+ auto *context = specifiedLayout.getContext();
+ int rank = specifiedInstData.size();
+
+ SmallVector<int64_t> laneLayout;
+ // set the laneLayout to use consumer's LaneLayout as base, but adjust its
+ // size to match the subgroupsize in case its original value is larger than 1
+ for (int i = 0; i < rank; i++) {
+ if (consumerLaneLayout[i] > 1) {
+ laneLayout.push_back(
+ std::max(static_cast<int64_t>(subgroupSize), consumerLaneLayout[i]));
+ } else {
+ laneLayout.push_back(1);
+ }
}
- SmallVector<int64_t> instDataA(aTy.getRank(), 1);
- instDataA[aTy.getRank() - 2] = maxALen;
- instDataA[aTy.getRank() - 1] = kDimSize;
- SmallVector<int64_t> instDataB(bTy.getRank(), 1);
- instDataB[bTy.getRank() - 2] = kDimSize;
- instDataB[bTy.getRank() - 1] = maxBLen;
- SmallVector<int64_t> instDataCD(cdTy.getRank(), 1);
- instDataCD[cdTy.getRank() - 2] = maxALen;
- instDataCD[cdTy.getRank() - 1] = maxCLen;
- return std::make_tuple(instDataA, instDataB, instDataCD);
+ if (!isValidLaneLayout(specifiedInstData, laneLayout, consumerLaneData))
+ return std::nullopt;
+ return buildInstDataLayoutWithLane(context, specifiedInstData, laneLayout,
+ consumerLaneData,
+ consumerLayout.getOrder());
}
-/// Helper function to set up subgroup layouts for DPAS operands A, B, and C/D.
-/// Returns the three layouts if successful, nullopt otherwise.
-static std::optional<
+/// Completes user-provided DPAS A/B/C-D anchors that carry only inst_data by
+/// filling in lane_layout / lane_data. The lane factorization mirrors the
+/// InstData branch of `setupDpasLayout` (derived from each operand's shape and
+/// matmul role, B using VNNI packing); the user's inst_data is preserved.
+std::optional<
std::tuple<xegpu::DistributeLayoutAttr, xegpu::DistributeLayoutAttr,
xegpu::DistributeLayoutAttr>>
-getupDpasSubgroupLayouts(mlir::MLIRContext *context, VectorType aTy,
- VectorType bTy, VectorType cdTy,
- xegpu::DistributeLayoutAttr consumerLayout, int numSg,
- const xegpu::uArch::uArch *uArch) {
- auto instDataVecs = getDpasInstDataVectors(aTy, bTy, cdTy, uArch);
- if (!instDataVecs)
+xegpu::completeDpasLaneLayoutFromInstData(xegpu::DistributeLayoutAttr aLayout,
+ xegpu::DistributeLayoutAttr bLayout,
+ xegpu::DistributeLayoutAttr cdLayout,
+ VectorType aTy, VectorType bTy,
+ VectorType cdTy,
+ const xegpu::uArch::uArch *uArch) {
+ auto context = aTy.getContext();
+ const auto *uArchInstruction =
+ dyn_cast<xegpu::uArch::SubgroupMatrixMultiplyAcc>(uArch->getInstruction(
+ xegpu::uArch::InstructionKind::SubgroupMatrixMultiplyAcc));
+ if (!uArchInstruction)
return std::nullopt;
- auto [instDataA, instDataB, instDataCD] = *instDataVecs;
- assert(instDataA.size() == 2 && instDataB.size() == 2 &&
- instDataCD.size() == 2 &&
- "Sg layout creation expects valid 2D inst data");
+ auto subgroupSize = uArch->getSubgroupSize();
- std::optional<LayoutRepresentation> consumerSgLayout = std::nullopt;
- if (consumerLayout && consumerLayout.isForWorkgroup()) {
- SmallVector<int64_t> sgLayoutD = consumerLayout.getEffectiveSgLayoutAsInt();
- consumerSgLayout = std::make_pair(sgLayoutD[0], sgLayoutD[1]);
- }
+ auto [laneLayoutA, laneDataA] = compute2DBlockIOLaneLayoutAndData(
+ aTy.getShape(), subgroupSize,
+ aTy.getElementType().getIntOrFloatBitWidth(),
+ uArchInstruction->getPackedFormatBitSizeA());
+ auto [laneLayoutB, laneDataB] = compute2DBlockIOLaneLayoutAndData(
+ bTy.getShape(), subgroupSize,
+ bTy.getElementType().getIntOrFloatBitWidth(),
+ uArchInstruction->getPackedFormatBitSizeB(), /*vnni=*/true);
+ auto [laneLayoutCD, laneDataCD] = compute2DBlockIOLaneLayoutAndData(
+ cdTy.getShape(), subgroupSize,
+ cdTy.getElementType().getIntOrFloatBitWidth(),
+ cdTy.getElementType().getIntOrFloatBitWidth());
+ SmallVector<int64_t> instDataA = aLayout.getEffectiveInstDataAsInt();
+ SmallVector<int64_t> instDataB = bLayout.getEffectiveInstDataAsInt();
+ SmallVector<int64_t> instDataCD = cdLayout.getEffectiveInstDataAsInt();
+ if (!isValidLaneLayout(instDataA, laneLayoutA, laneDataA) ||
+ !isValidLaneLayout(instDataB, laneLayoutB, laneDataB) ||
+ !isValidLaneLayout(instDataCD, laneLayoutCD, laneDataCD))
+ return std::nullopt;
+ return std::make_tuple(
+ buildInstDataLayoutWithLane(context, instDataA, laneLayoutA, laneDataA),
+ buildInstDataLayoutWithLane(context, instDataB, laneLayoutB, laneDataB),
+ buildInstDataLayoutWithLane(context, instDataCD, laneLayoutCD,
+ laneDataCD));
+}
- // Get all valid layouts for A, B and C/D operands
- auto layoutsA = getValidLayouts(aTy.getShape(), instDataA, numSg);
- auto layoutsB = getValidLayouts(bTy.getShape(), instDataB, numSg);
- auto layoutsCD = getValidLayouts(cdTy.getShape(), instDataCD, numSg);
- if (layoutsA.empty() || layoutsB.empty() || layoutsCD.empty())
+/// Like completeDpasLaneLayoutFromInstData, but for dpas_mx: also re-derives
+/// the A_scale / B_scale layouts from the completed A / B layouts via
+/// `createScaleLayout`, matching the default path of `setupDpasMxLayout`.
+std::optional<
+ std::tuple<xegpu::DistributeLayoutAttr, xegpu::DistributeLayoutAttr,
+ xegpu::DistributeLayoutAttr, xegpu::DistributeLayoutAttr,
+ xegpu::DistributeLayoutAttr>>
+xegpu::completeDpasMxLaneLayoutFromInstData(
+ xegpu::DistributeLayoutAttr aLayout, xegpu::DistributeLayoutAttr bLayout,
+ xegpu::DistributeLayoutAttr cdLayout, VectorType aTy, VectorType bTy,
+ VectorType cdTy, VectorType aScaleTy, VectorType bScaleTy,
+ const xegpu::uArch::uArch *uArch) {
+ auto completed = completeDpasLaneLayoutFromInstData(
+ aLayout, bLayout, cdLayout, aTy, bTy, cdTy, uArch);
+ if (!completed)
return std::nullopt;
+ auto context = aTy.getContext();
+ auto [completedA, completedB, completedCD] = *completed;
- // Pick the best subgroup layout
- llvm::DenseSet<LayoutRepresentation> setA(layoutsA.begin(), layoutsA.end());
- llvm::DenseSet<LayoutRepresentation> setCD(layoutsCD.begin(),
- layoutsCD.end());
- std::optional<LayoutRepresentation> bestPick;
- auto checkAlignedSgDataAB = [&](LayoutRepresentation sgLayout) {
- return aTy.getShape().back() / sgLayout.second ==
- bTy.getShape().front() / sgLayout.first;
- };
- for (auto &sgLayout : layoutsB) {
- if (setA.contains(sgLayout) && setCD.contains(sgLayout)) {
- if (!checkAlignedSgDataAB(sgLayout))
- continue;
- // Is in (A and B and CD) and matches consumer -> best pick
- if (consumerSgLayout.has_value() && sgLayout == *consumerSgLayout) {
- bestPick = sgLayout;
- break;
+ auto aScaleLayout =
+ createScaleLayout(context, aTy, aScaleTy, completedA, false, uArch);
+ auto bScaleLayout =
+ createScaleLayout(context, bTy, bScaleTy, completedB, true, uArch);
+
+ return std::make_tuple(completedA, completedB, completedCD, aScaleLayout,
+ bScaleLayout);
+}
+
+/// Sets up layout for reduction operations by creating a SliceAttr for the
+/// result.
+///
+/// Algorithm Overview:
+/// This function attempts to construct a source layout that, when sliced along
+/// reduction dimensions, produces a result layout compatible with the
+/// consumer layout.
+///
+/// For subgroup layouts, it first tries to align the source layout's subgroup
+/// layout and data with the consumer's layout on non-reduction dimensions.
+/// Then, it distributes remaining subgroups across reduction dimensions. This
+/// avoids subgroup data redistribution overhead between the reduced result and
+/// its consumer. When the consumer layout is a slice layout, it attempts to
+/// reuse the slice layout's parent layout for the source to further minimize
+/// potential data redistribution.
+///
+/// This is a best-effort alignment, not a hard constraint: the goal is only to
+/// pick a *legal* source layout that minimizes redistribution against the
+/// (single, first-arriving) consumer layout. There is no failure path - when
+/// the consumer's slice layout cannot be reused as-is (example 2 below), the
+/// function falls back to distributing all subgroups on the non-reduction
+/// dimensions first and the remainder on the reduction dimensions, which always
+/// yields a valid source layout. If the resulting source layout still
diff ers
+/// from what some consumer expects (e.g. a second, inconsistent consumer), that
+/// mismatch is reconciled later by the layout conflict resolution process
+/// (`ResolveLayoutConflicts`), which inserts a `convert_layout` op - this
+/// function never has to give up.
+///
+/// For the InstData and Lane layout kinds only the innermost two dimensions
+/// are distributed; all leading dimensions are assumed to be unit dimensions.
+/// This assumption is checked via `leadingDimsAreUnit`. The lane_layout and
+/// lane_data are computed by `computeReductionLaneLayoutAndData`, which picks
+/// a layout that minimizes cross-lane reduction (reducing within a lane when
+/// only one of the innermost two dims is a reduction dim). The inst_data is
+/// simply the element-wise product lane_layout * lane_data.
+///
+/// The function returns the *result* layout (the SliceAttr). The *source*
+/// layout it decides on is the parent of that slice; both are listed below so
+/// the relationship is explicit.
+///
+/// Examples:
+/// 1. Subgroup layout - Row reduction on 2D tensor:
+/// srcShape=[32, 128], reductionDims=[1], resShape=[32], subgroupSize=16,
+/// NumSg=32
+/// * Consumer Layout:
+/// #xegpu.slice<#xegpu.layout<sg_layout=[4, 8], sg_data=[8, 8]>, dims =
+/// [1]>}
+/// * Source Layout (decided by this function):
+/// #xegpu.layout<sg_layout=[4, 8], sg_data=[8, 16]>
+/// * Result Layout (returned):
+/// #xegpu.slice<#xegpu.layout<sg_layout=[4, 8], sg_data=[8, 16]>, dims =
+/// [1]>}
+/// The consumer slices exactly the reduction dim, so its parent layout is
+/// reused for the source: sg_layout is kept, but the source's sg_data on
+/// the reduction dim is grown from 8 to 16 (= srcShape[1] / sg_layout[1] =
+/// 128 / 8) so the source tile is evenly distributed over the reduction
+/// dim. Slicing that source over dim 1 reproduces the consumer.
+///
+/// 2. Subgroup layout - Same shapes as above but consumer doesn't have a
+/// reusable slice layout, so the algorithm distributes all subgroups on the
+/// non-reduction dims first and the remainder on the reduction dims.
+/// 2a. * Consumer Layout:
+/// #xegpu.layout<sg_layout=[32], sg_data=[1]>
+/// * Source Layout (decided by this function):
+/// #xegpu.layout<sg_layout=[32, 1], sg_data=[1, 128]>
+/// * Result Layout (returned):
+/// #xegpu.slice<#xegpu.layout<sg_layout=[32, 1], sg_data=[1, 128]>,
+/// dims = [1]>}
+/// All 32 subgroups land on the non-reduction dim 0; the reduction dim
+/// 1 gets the leftover (sg_layout=1, so the whole length 128 lives in
+/// one subgroup's sg_data).
+/// 2b. * Consumer Layout:
+/// #xegpu.slice<#xegpu.layout<sg_layout=[8, 2, 4], sg_data=[4, 64,
+/// 32]>, dims = [1, 2]>}
+/// * Source Layout (decided by this function):
+/// #xegpu.layout<sg_layout=[8, 4], sg_data=[4, 32]>
+/// * Result Layout (returned):
+/// #xegpu.slice<#xegpu.layout<sg_layout=[8, 4], sg_data=[4, 32]>,
+/// dims = [1]>}
+/// The consumer slices dims [1, 2] which do not match this op's
+/// reductionDims, so it can't be reused as-is; subgroups are
+/// re-distributed (non-reduction dim first, then reduction dim).
+///
+/// 3. Lane layout - Default (lanes on innermost dim):
+/// srcShape=[32, 64], reductionDims=[0], subgroupSize=16
+/// * Source Layout (decided by this function):
+/// laneLayout=[1, 16], laneData=[1, 1] (returned sliced over dim 0).
+/// The innermost dim is not reduced, so lanes stay on it.
+///
+/// 4. Lane layout - Switch (lanes moved off the reduction dim):
+/// srcShape=[32, 64], reductionDims=[1], subgroupSize=16
+/// * Source Layout (decided by this function):
+/// laneLayout=[16, 1], laneData=[1, 1] (returned sliced over dim 1).
+/// The innermost dim is the sole reduction dim, so lanes move to the
+/// non-reduction dim to reduce within a lane. This switch only happens
+/// when the consumer has no reduction dims to broadcast the result back
+/// along (i.e. the consumer layout is not a slice over this reduction);
+/// otherwise the default (example 3) is used.
+///
+/// 5. Lane layout - No switch when both inner dims are reduced (reduction to
+/// scalar):
+/// srcShape=[32, 64], reductionDims=[0, 1], subgroupSize=16
+/// * Source Layout (decided by this function):
+/// laneLayout=[1, 16], laneData=[1, 1] (returned sliced over dims
+/// [0,1]).
+/// Both dims are reduced, so this is not a *sole* innermost reduction; the
+/// switch condition (example 4) does not apply and lanes stay on the
+/// innermost dim. The cross-lane reduction here is unavoidable.
+///
+/// 6. Lane layout - No switch when the consumer slices the reduction dim:
+/// srcShape=[32, 64], reductionDims=[1], subgroupSize=16
+/// * Consumer Layout:
+/// #xegpu.slice<#xegpu.layout<laneLayout=[1, 16], laneData=[1, 1]>,
+/// dims = [1]>}
+/// * Source Layout (decided by this function):
+/// #xegpu.layout<laneLayout=[1, 16], laneData=[1, 1]> (the consumer
+/// slice's parent, reused directly; returned sliced over dim 1).
+/// Same shape/reductionDims as example 4, but here the consumer is a slice
+/// over the reduction dim, so it can broadcast the result back along that
+/// dim. The slice's parent layout is reused as the source (no switch, no
+/// re-derivation); the inst_data propagation step has already inserted a
+/// convert_layout if needed, so the lane-level layout can be reused as-is.
+
+xegpu::SliceAttr xegpu::setupMultiReductionResultLayout(
+ xegpu::LayoutKind layoutKind, VectorType srcVecTy,
+ DistributeLayoutAttr consumerLayout, SmallVector<int64_t> reductionDims,
+ int numSg, const xegpu::uArch::uArch *uArch) {
+
+ auto srcShape = srcVecTy.getShape();
+ int srcRank = srcShape.size();
+ auto context = srcVecTy.getContext();
+
+ const int subgroupSize = uArch->getSubgroupSize();
+ int64_t maxReduceVectorSize = 1; // could extend to spirv vector Size
+ xegpu::DistributeLayoutAttr srcLayout;
+ if (layoutKind == xegpu::LayoutKind::Subgroup) {
+ xegpu::SliceAttr consumerSliceLayout =
+ dyn_cast_if_present<xegpu::SliceAttr>(consumerLayout);
+ if (consumerSliceLayout &&
+ consumerSliceLayout.getDims().asArrayRef().equals(reductionDims)) {
+ srcLayout = consumerSliceLayout.getParent();
+ SmallVector<int64_t> sgLayoutFromConsumer =
+ srcLayout.getEffectiveSgLayoutAsInt();
+ auto srcSgData = computeShapeRatio(srcShape, sgLayoutFromConsumer);
+ if (srcSgData)
+ for (int dim = 0; dim < srcRank; dim++) {
+ if (llvm::is_contained(reductionDims, dim))
+ srcLayout =
+ srcLayout.setDimData(dim, srcSgData.value()[dim], -1, -1);
+ }
+ } else {
+ SmallVector<int64_t> consumerSgLayout =
+ consumerLayout ? consumerLayout.getEffectiveSgLayoutAsInt()
+ : SmallVector<int64_t>();
+ SmallVector<int64_t> consumerSgData =
+ consumerLayout ? consumerLayout.getEffectiveSgDataAsInt()
+ : SmallVector<int64_t>();
+ SmallVector<int64_t> consumerOrder =
+ consumerLayout ? consumerLayout.getEffectiveOrderAsInt()
+ : SmallVector<int64_t>();
+ DenseI32ArrayAttr orderAttr =
+ consumerLayout ? consumerLayout.getOrder() : nullptr;
+ SmallVector<int64_t> sgLayout(srcRank), sgData(srcRank), order(srcRank);
+ int remainingSgCount =
+ consumerLayout ? consumerLayout.getNumSubgroups() : numSg;
+ int consumerIdx = 0;
+
+ // First pass: Match consumer's layout on non-reduction dimensions
+ for (int i = 0; i < srcRank; i++) {
+ if (!llvm::is_contained(reductionDims, i) &&
+ consumerIdx < static_cast<int>(consumerSgLayout.size())) {
+ sgLayout[i] = consumerSgLayout[consumerIdx];
+ sgData[i] = consumerSgData[consumerIdx];
+ remainingSgCount /= sgLayout[i];
+ order[i] = consumerOrder[consumerIdx];
+ consumerIdx++;
+ }
}
- // Is in (A and B and CD) layoutsB is ordered from most
- // balanced to least. So the first one we see is the most balanced one,
- // remember it and later only update if there is one that matches the
- // consumer.
- if (!bestPick)
- bestPick = sgLayout;
+
+ // Second pass: Distribute remaining subgroups across reduction dimensions
+ // the reduction to scalar case is handled only by this loop
+ int64_t remainOrder = consumerSgLayout.size();
+ for (int i = 0; i < srcRank; i++) {
+ if (llvm::is_contained(reductionDims, i)) {
+ sgLayout[i] =
+ std::min(srcShape[i], static_cast<int64_t>(remainingSgCount));
+ assert((srcShape[i] % sgLayout[i] == 0) &&
+ "source shape not divisible by sg_layout");
+ sgData[i] = srcShape[i] / sgLayout[i];
+ remainingSgCount /= sgLayout[i];
+ order[i] = remainOrder++;
+ }
+ }
+ DenseI32ArrayAttr resOrderAttr = DenseI32ArrayAttr::get(
+ context, SmallVector<int32_t>(order.begin(), order.end()));
+ if (!orderAttr || orderAttr.empty())
+ resOrderAttr = nullptr;
+ assert(remainingSgCount == 1 && "not all subgroups distributed");
+ srcLayout = buildLayout(context, sgLayout, sgData,
+ /*instData=*/{}, /*laneLayout=*/{},
+ /*laneData=*/{}, resOrderAttr);
+ }
+ } else if (layoutKind == xegpu::LayoutKind::InstData) {
+ xegpu::SliceAttr consumerSliceLayout =
+ dyn_cast_if_present<xegpu::SliceAttr>(consumerLayout);
+ auto consumerReductionDims =
+ consumerSliceLayout
+ ? SmallVector<int64_t>(consumerSliceLayout.getDims().asArrayRef())
+ : SmallVector<int64_t>({});
+ // A[i] reduced from A[i, j] is stored out directly, use vertical Lane
+ // layout like [16, 1]
+ bool verticalLaneLayout = consumerReductionDims.empty() &&
+ reductionDims.size() == 1 &&
+ reductionDims[0] == (srcRank - 1);
+ auto [laneLayout, laneData] = computeReductionLaneLayoutAndData(
+ srcShape, reductionDims, subgroupSize, maxReduceVectorSize,
+ verticalLaneLayout);
+ // inst_data is the per-instruction data, i.e. the element-wise product of
+ // lane_layout and lane_data.
+ SmallVector<int64_t> instData(srcRank);
+ for (int i = 0; i < srcRank; i++)
+ instData[i] = laneLayout[i] * laneData[i];
+ srcLayout =
+ buildInstDataLayoutWithLane(context, instData, laneLayout, laneData);
+ } else if (layoutKind == xegpu::LayoutKind::Lane) {
+ // Only the innermost two dimensions are distributed; all leading dimensions
+ // are assumed to be unit dimensions.
+ assert(leadingDimsAreUnit(srcShape, /*numInnerDims=*/2) &&
+ "Lane reduction layout assumes all leading (non-innermost-two) "
+ "dimensions are unit dimensions");
+ xegpu::SliceAttr consumerSliceLayout =
+ dyn_cast_if_present<xegpu::SliceAttr>(consumerLayout);
+ auto consumerReductionDims =
+ consumerSliceLayout
+ ? SmallVector<int64_t>(consumerSliceLayout.getDims().asArrayRef())
+ : SmallVector<int64_t>({});
+ if (consumerSliceLayout &&
+ consumerSliceLayout.getDims().asArrayRef().equals(reductionDims)) {
+ // at the lane level, the consumerSliceLayout can be directly reused
+ // since the inst_data propagation already insert convert_layout if
+ // the layout is not consistent
+ srcLayout = consumerSliceLayout.getParent();
+ } else {
+ bool verticalLaneLayout = consumerReductionDims.empty() &&
+ reductionDims.size() == 1 &&
+ reductionDims[0] == (srcRank - 1);
+ auto [laneLayout, laneData] = computeReductionLaneLayoutAndData(
+ srcShape, reductionDims, subgroupSize, maxReduceVectorSize,
+ verticalLaneLayout);
+ srcLayout = buildLaneLayout(context, laneLayout, laneData);
}
}
- if (!bestPick)
- return std::nullopt;
-
- SmallVector<int> sgLayout = {static_cast<int>(bestPick->first),
- static_cast<int>(bestPick->second)};
- SmallVector<int> sgDataA = {static_cast<int>(aTy.getShape()[0] / sgLayout[0]),
- static_cast<int>(aTy.getShape()[1])};
- SmallVector<int> sgDataB = {
- static_cast<int>(bTy.getShape()[0]),
- static_cast<int>(bTy.getShape()[1] / sgLayout[1])};
- SmallVector<int> sgDataCD = {
- static_cast<int>(cdTy.getShape()[0] / sgLayout[0]),
- static_cast<int>(cdTy.getShape()[1] / sgLayout[1])};
-
- auto dpasALayout =
- xegpu::LayoutAttr::get(context, DenseI32ArrayAttr::get(context, sgLayout),
- DenseI32ArrayAttr::get(context, sgDataA), nullptr,
- nullptr, nullptr, nullptr);
- auto dpasBLayout =
- xegpu::LayoutAttr::get(context, DenseI32ArrayAttr::get(context, sgLayout),
- DenseI32ArrayAttr::get(context, sgDataB), nullptr,
- nullptr, nullptr, nullptr);
- auto dpasCDLayout =
- xegpu::LayoutAttr::get(context, DenseI32ArrayAttr::get(context, sgLayout),
- DenseI32ArrayAttr::get(context, sgDataCD), nullptr,
- nullptr, nullptr, nullptr);
- return std::make_tuple(dpasALayout, dpasBLayout, dpasCDLayout);
+ return xegpu::SliceAttr::get(context, srcLayout,
+ DenseI64ArrayAttr::get(context, reductionDims));
}
-/// Sets up the anchor layouts for dpas operands (A, B, and C/D).
-/// The numSg and consumerLayout (optional) are only used by sg layout
-/// creation.
-std::optional<
- std::tuple<xegpu::DistributeLayoutAttr, xegpu::DistributeLayoutAttr,
- xegpu::DistributeLayoutAttr>>
-xegpu::setupDpasLayout(xegpu::LayoutKind layoutKind, VectorType aTy,
- VectorType bTy, VectorType cdTy,
- xegpu::DistributeLayoutAttr consumerLayout, int numSg,
- const xegpu::uArch::uArch *uArch) {
- auto context = aTy.getContext();
- const auto *uArchInstruction =
- dyn_cast<xegpu::uArch::SubgroupMatrixMultiplyAcc>(uArch->getInstruction(
- xegpu::uArch::InstructionKind::SubgroupMatrixMultiplyAcc));
+/// Sets up layout for Reduction operations by creating a SliceAttr for the
+/// result.
+xegpu::SliceAttr
+xegpu::setupReductionResultLayout(xegpu::LayoutKind layoutKind,
+ VectorType srcVecTy,
+ const xegpu::uArch::uArch *uArch) {
+
+ auto srcShape = srcVecTy.getShape();
+ auto context = srcVecTy.getContext();
+ auto subgroupSize = uArch->getSubgroupSize();
+ xegpu::LayoutAttr srcLayout;
if (layoutKind == xegpu::LayoutKind::Subgroup) {
- assert(numSg > 0 &&
- "Number of subgroups must be provided for sg layout creation.");
- return getupDpasSubgroupLayouts(context, aTy, bTy, cdTy, consumerLayout,
- numSg, uArch);
+ assert(false &&
+ "subgroup layout assignment not supported for reduction (op "
+ "is not expected at this level).");
} else if (layoutKind == xegpu::LayoutKind::InstData) {
- auto instDataVecs = getDpasInstDataVectors(aTy, bTy, cdTy, uArch);
- if (!instDataVecs)
- return std::nullopt;
- auto [instDataA, instDataB, instDataCD] = *instDataVecs;
- return std::make_tuple(
- xegpu::LayoutAttr::get(
- context, SmallVector<int>(instDataA.begin(), instDataA.end())),
- xegpu::LayoutAttr::get(
- context, SmallVector<int>(instDataB.begin(), instDataB.end())),
- xegpu::LayoutAttr::get(
- context, SmallVector<int>(instDataCD.begin(), instDataCD.end())));
+ assert(false &&
+ "instData layout assignment not supported for reduction (op "
+ "is not expected at this level).");
} else if (layoutKind == xegpu::LayoutKind::Lane) {
- auto aLayout = getDefaultLaneLayout2DBlockIo(
- aTy, uArch, uArchInstruction->getPackedFormatBitSizeA());
- auto bLayout = getDefaultLaneLayout2DBlockIo(
- bTy, uArch, uArchInstruction->getPackedFormatBitSizeB(), true);
- auto cdLayout = getDefaultLaneLayout2DBlockIo(
- cdTy, uArch /*, packingSize = std::nullopt */);
- return std::make_tuple(aLayout, bLayout, cdLayout);
+ SmallVector<int64_t> laneLayout(1), laneData(1);
+ laneLayout[0] = std::min(static_cast<int64_t>(subgroupSize), srcShape[0]);
+ laneData[0] = 1;
+ srcLayout = buildLaneLayout(context, laneLayout, laneData);
}
- return std::nullopt;
+
+ auto result = xegpu::SliceAttr::get(context, srcLayout,
+ DenseI64ArrayAttr::get(context, 0));
+ return result;
}
-/// Helper to create a scale layout derived from a matrix operand layout.
-/// The scale layout is computed by mapping each dimension of the matrix layout
-/// to the corresponding scale tensor dimension using the ratio between the
-/// matrix and scale shapes.
+/// Adjusts `consumerLayout`'s innermost-dim data field selected by
+/// `layoutKind` so that the source layout can be safely inferred by dividing
+/// that value by `ratio`. Doubles the value until the divisibility constraint
+/// is met, bounded above by `bound` like result-shape.
+///
+/// Used by ops whose source relates to the result by a fixed factor along the
+/// innermost dim (e.g., bitcast: bitwidth ratio; interleave: 2x).
+///
+/// Divisibility constraints per LayoutKind:
+/// - Subgroup: sgData[innermost] % ratio == 0
+/// - InstData: instData[innermost] % (laneLayout[innermost] * ratio) == 0
+/// (laneLayout falls back to subgroupSize if absent)
+/// - Lane: laneData[innermost] % ratio == 0
static xegpu::DistributeLayoutAttr
-createScaleLayout(mlir::MLIRContext *context, VectorType matrixTy,
- VectorType scaleTy, xegpu::DistributeLayoutAttr matrixLayout,
- bool isBScale, const xegpu::uArch::uArch *uArch) {
- if (!scaleTy || !matrixLayout)
- return nullptr;
-
- // Calculate scaling factor by dividing matrix shape by scale shape
- ArrayRef<int64_t> matrixShape = matrixTy.getShape();
- ArrayRef<int64_t> scaleShape = scaleTy.getShape();
+adjustInnermostDimForDivisibility(xegpu::DistributeLayoutAttr consumerLayout,
+ xegpu::LayoutKind layoutKind,
+ size_t innerMostDim, int ratio, int64_t bound,
+ const xegpu::uArch::uArch *uArch) {
+ SmallVector<int64_t> sgData = consumerLayout.getEffectiveSgDataAsInt();
+ SmallVector<int64_t> instData = consumerLayout.getEffectiveInstDataAsInt();
+ SmallVector<int64_t> laneData = consumerLayout.getEffectiveLaneDataAsInt();
+ SmallVector<int64_t> laneLayout =
+ consumerLayout.getEffectiveLaneLayoutAsInt();
- // Scale shapes can be 1D or 2D, handle both cases
- if (scaleShape.empty())
- return nullptr;
+ int64_t sgDataValue = -1;
+ int64_t instDataValue = -1;
+ int64_t laneDataValue = -1;
- auto uArchInstruction =
- dyn_cast<xegpu::uArch::SubgroupScaledMatrixMultiplyAcc>(
- uArch->getInstruction(
- xegpu::uArch::InstructionKind::SubgroupScaledMatrixMultiplyAcc));
+ if (layoutKind == xegpu::LayoutKind::Subgroup) {
+ sgDataValue = sgData[innerMostDim];
+ while ((sgDataValue <= bound) && (sgDataValue % ratio) != 0)
+ sgDataValue *= 2;
+ } else if (layoutKind == xegpu::LayoutKind::InstData) {
+ instDataValue = instData[innerMostDim];
+ const int innermostDimLaneLayout = laneLayout.empty()
+ ? uArch->getSubgroupSize()
+ : laneLayout[innerMostDim];
+ while ((instDataValue <= bound) &&
+ (instDataValue % (innermostDimLaneLayout * ratio) != 0))
+ instDataValue *= 2;
+ assert((bound % instDataValue) == 0 &&
+ "bound, instData, and laneLayout for innermost must be 2^n!");
+ } else if (layoutKind == xegpu::LayoutKind::Lane) {
+ laneDataValue = laneData[innerMostDim];
+ while ((laneDataValue <= bound) && (laneDataValue % ratio) != 0)
+ laneDataValue *= 2;
+ }
- int64_t rank = matrixLayout.getRank();
- assert(rank >= 2 && "dpas layouts must be at least two dimensions");
+ return consumerLayout.setDimData(innerMostDim, sgDataValue, instDataValue,
+ laneDataValue);
+}
- SmallVector<int64_t> sgLayout = matrixLayout.getEffectiveSgLayoutAsInt();
- SmallVector<int64_t> sgData = matrixLayout.getEffectiveSgDataAsInt();
- SmallVector<int64_t> instData = matrixLayout.getEffectiveInstDataAsInt();
- SmallVector<int64_t> laneLayout = matrixLayout.getEffectiveLaneLayoutAsInt();
- SmallVector<int64_t> laneData = matrixLayout.getEffectiveLaneDataAsInt();
- auto order = matrixLayout.getOrder();
+/// Sets up the result layout for a bitcast operation.
+/// When casting to a smaller bitwidth, adjusts the layout dimensions (sgData,
+/// instData, or laneData) by multiplying by the bitwidth ratio to ensure the
+/// result layout can be correctly divided back to the source layout during
+/// inference.
+///
+/// Examples:
+/// 1. Casting f32 -> f16 (32-bit to 16-bit, bitWidthRatio = 2):
+/// Consumer layout: instData=[1, 16], subgroupSize=16
+/// Source shape: [8, 32]
+/// Result layout: instData=[1, 32] (16 * 2)
+/// The innermost dimension is multiplied by 2 to maintain consistency.
+///
+/// 2. Casting f32 -> i8 (32-bit to 8-bit, bitWidthRatio = 4):
+/// Consumer instData=[1, 16], subgroupSize=16
+/// Source shape: [4, 128]
+/// adjust the instData from [1, 16] to [1, 16 * 4 = 64]
+///
+/// 3. Casting i8 -> i32 (8-bit to 32-bit, bitWidthRatio = 1/4):
+/// Consumer layout: laneLayout=[1, 16], laneData=[1, 4]
+/// No adjustment needed - returns consumer layout directly.
+///
+xegpu::DistributeLayoutAttr xegpu::setupBitCastResultLayout(
+ xegpu::LayoutKind layoutKind, VectorType srcVecTy, VectorType resVecTy,
+ DistributeLayoutAttr consumerLayout, const xegpu::uArch::uArch *uArch) {
- SmallVector<int> scaleSgLayout;
- SmallVector<int> scaleSgData;
- if (!sgLayout.empty() && !sgData.empty()) {
- scaleSgLayout.assign(sgLayout.begin(), sgLayout.end());
- scaleSgData.assign(sgData.begin(), sgData.end());
- scaleSgData[rank - 2] = std::max<int64_t>(
- scaleShape[rank - 2] / (matrixShape[rank - 2] / sgData[rank - 2]), 1);
- scaleSgData[rank - 1] = std::max<int64_t>(
- scaleShape[rank - 1] / (matrixShape[rank - 1] / sgData[rank - 1]), 1);
- }
+ int srcElemTyBitWidth = srcVecTy.getElementType().getIntOrFloatBitWidth();
+ int resElemTyBitWidth = resVecTy.getElementType().getIntOrFloatBitWidth();
- // For DPAS_MX scales: if matrix has inst_data, scale needs adjusted
- // inst_data. Scale inst_data is derived from matrix inst_data divided by
- // scale factor.
- SmallVector<int> scaleInstData;
- if (!instData.empty()) {
- scaleInstData.assign(instData.begin(), instData.end());
- if (isBScale)
- scaleInstData[rank - 2] = std::max<int64_t>(
- scaleShape[rank - 2] / (matrixShape[rank - 2] / instData[rank - 2]),
- 1);
- else
- scaleInstData[rank - 1] = std::max<int64_t>(
- scaleShape[rank - 1] / (matrixShape[rank - 1] / instData[rank - 1]),
- 1);
- }
+ ArrayRef<int64_t> srcShape = srcVecTy.getShape();
+ ArrayRef<int64_t> resShape = resVecTy.getShape();
- SmallVector<int> scaleLaneLayout;
- SmallVector<int> scaleLaneData;
- if (!laneLayout.empty() && !laneData.empty()) {
- scaleLaneLayout.assign(laneLayout.begin(), laneLayout.end());
- scaleLaneData.assign(laneData.begin(), laneData.end());
- bool isRowMajor = uArchInstruction->isLaneLayoutRowMajorOrder();
- if (isBScale ^ isRowMajor) {
- std::swap(scaleLaneLayout[rank - 2], scaleLaneLayout[rank - 1]);
- scaleLaneLayout[rank - 2] =
- std::min<int64_t>(scaleShape[rank - 2], scaleLaneLayout[rank - 2]);
- }
- scaleLaneData[rank - 2] =
- std::max<int64_t>(scaleShape[rank - 2] / scaleLaneLayout[rank - 2], 1);
- scaleLaneData[rank - 1] =
- std::max<int64_t>(scaleShape[rank - 1] / scaleLaneLayout[rank - 1], 1);
- }
- return xegpu::LayoutAttr::get(
- context,
- scaleSgLayout.empty() ? nullptr
- : DenseI32ArrayAttr::get(context, scaleSgLayout),
- scaleSgData.empty() ? nullptr
- : DenseI32ArrayAttr::get(context, scaleSgData),
- scaleInstData.empty() ? nullptr
- : DenseI32ArrayAttr::get(context, scaleInstData),
- scaleLaneLayout.empty()
- ? nullptr
- : DenseI32ArrayAttr::get(context, scaleLaneLayout),
- scaleLaneData.empty() ? nullptr
- : DenseI32ArrayAttr::get(context, scaleLaneData),
- order);
-}
+ assert(consumerLayout.getRank() == static_cast<int64_t>(srcShape.size()) &&
+ "laneData must be available for all dimensions");
-/// Sets up the anchor layouts for dpas_mx operands (A, B, C/D, A_scale, and
-/// B_scale). The numSg and consumerLayout (optional) are only used by sg layout
-/// creation.
-std::optional<
- std::tuple<xegpu::DistributeLayoutAttr, xegpu::DistributeLayoutAttr,
- xegpu::DistributeLayoutAttr, xegpu::DistributeLayoutAttr,
- xegpu::DistributeLayoutAttr>>
-xegpu::setupDpasMxLayout(xegpu::LayoutKind layoutKind, VectorType aTy,
- VectorType bTy, VectorType cdTy, VectorType aScaleTy,
- VectorType bScaleTy,
- xegpu::DistributeLayoutAttr consumerLayout, int numSg,
- const xegpu::uArch::uArch *uArch) {
- auto context = aTy.getContext();
+ // Casting to same/larger element type: result has fewer (or equal) elements
+ // along the innermost dim, no adjustment needed.
+ if (srcElemTyBitWidth <= resElemTyBitWidth)
+ return consumerLayout;
- if (layoutKind == xegpu::LayoutKind::Subgroup) {
- assert(numSg > 0 &&
- "Number of subgroups must be provided for sg layout creation.");
- auto dpasLayouts = getupDpasSubgroupLayouts(context, aTy, bTy, cdTy,
- consumerLayout, numSg, uArch);
- if (!dpasLayouts)
- return std::nullopt;
+ // Casting to smaller element type: result has more elements along innermost
+ // dim. Adjust the innermost data field upward so the source layout can be
+ // recovered by dividing by bitWidthRatio.
+ size_t innerMostDim = srcShape.size() - 1;
+ int bitWidthRatio = srcElemTyBitWidth / resElemTyBitWidth;
+ return adjustInnermostDimForDivisibility(consumerLayout, layoutKind,
+ innerMostDim, bitWidthRatio,
+ resShape[innerMostDim], uArch);
+}
- auto [dpasALayout, dpasBLayout, dpasCDLayout] = *dpasLayouts;
+/// Sets up the result layout for an interleave operation to ensure the source
+/// layout can be safely derived. Interleave doubles the innermost dimension,
+/// so the result layout must ensure that laneData is a multiple
+/// of 2, and instData must be divisible by innermostDimLaneLayout * 2.
+///
+/// Example:
+/// Interleave: vector<128x256xf4> -> vector<128x512xf4>
+/// Consumer layout: laneLayout=[1, 16], laneData=[1, 4], instData=[1, 64]
+/// Result layout adjustment to ensure source can be safely inferred:
+/// - laneData must be >= 2 and multiple of 2 (so source = laneData/2 is
+/// valid)
+/// - instData must be divisible by (16 * 2 = 32) (so source = instData/2 is
+/// valid)
+/// - Adjusted instData: ensure (instData % 32 == 0)
+///
+xegpu::DistributeLayoutAttr xegpu::setupInterleaveResultLayout(
+ xegpu::LayoutKind layoutKind, VectorType srcVecTy, VectorType resVecTy,
+ DistributeLayoutAttr consumerLayout, const xegpu::uArch::uArch *uArch) {
- // Create scale layouts
- auto aScaleLayout =
- createScaleLayout(context, aTy, aScaleTy, dpasALayout, false, uArch);
+ ArrayRef<int64_t> resShape = resVecTy.getShape();
+ assert(consumerLayout.getRank() == static_cast<int64_t>(resShape.size()) &&
+ "consumer layout rank must match source shape rank");
- auto bScaleLayout =
- createScaleLayout(context, bTy, bScaleTy, dpasBLayout, true, uArch);
+ // Interleave doubles the innermost dimension (ratio = 2). Adjust the
+ // innermost data field so the source layout can be recovered by dividing
+ // by 2.
+ const size_t innerMostDim = resShape.size() - 1;
+ constexpr int ratio = 2;
+ return adjustInnermostDimForDivisibility(consumerLayout, layoutKind,
+ innerMostDim, ratio,
+ resShape[innerMostDim], uArch);
+}
- return std::make_tuple(dpasALayout, dpasBLayout, dpasCDLayout, aScaleLayout,
- bScaleLayout);
- } else if (layoutKind == xegpu::LayoutKind::InstData) {
- auto instDataVecs =
- getDpasInstDataVectors(aTy, bTy, cdTy, uArch, /*isDpasMx=*/true);
- if (!instDataVecs)
- return std::nullopt;
- auto [instDataA, instDataB, instDataCD] = *instDataVecs;
+/// Sets up the result layout for an insert strided slice operation.
+/// Creates a result layout based on the specified layout kind (InstData or
+/// Lane).
+xegpu::DistributeLayoutAttr xegpu::setupInsertStridedSliceResultLayout(
+ xegpu::LayoutKind layoutKind, VectorType srcVectorTy,
+ VectorType resVectorTy, xegpu::DistributeLayoutAttr consumerLayout,
+ const xegpu::uArch::uArch *uArch) {
- auto dpasALayout = xegpu::LayoutAttr::get(
- context, SmallVector<int>(instDataA.begin(), instDataA.end()));
- auto dpasBLayout = xegpu::LayoutAttr::get(
- context, SmallVector<int>(instDataB.begin(), instDataB.end()));
- auto dpasCDLayout = xegpu::LayoutAttr::get(
- context, SmallVector<int>(instDataCD.begin(), instDataCD.end()));
+ xegpu::DistributeLayoutAttr requiredResLayout;
+ SmallVector<int64_t> consumerInstData =
+ consumerLayout.getEffectiveInstDataAsInt();
+ SmallVector<int64_t> consumerLaneData =
+ consumerLayout.getEffectiveLaneDataAsInt();
+ SmallVector<int64_t> consumerLaneLayout =
+ consumerLayout.getEffectiveLaneLayoutAsInt();
+ ArrayRef<int64_t> srcShape = srcVectorTy.getShape();
+ int64_t laneDataValue = -1;
- // Create scale layouts
- auto aScaleLayout =
- createScaleLayout(context, aTy, aScaleTy, dpasALayout, false, uArch);
- auto bScaleLayout =
- createScaleLayout(context, bTy, bScaleTy, dpasBLayout, true, uArch);
+ requiredResLayout = consumerLayout;
+ int srcRank = srcShape.size();
- return std::make_tuple(dpasALayout, dpasBLayout, dpasCDLayout, aScaleLayout,
- bScaleLayout);
+ if (layoutKind == xegpu::LayoutKind::Subgroup ||
+ layoutKind == xegpu::LayoutKind::InstData) {
+ assert(false && "subgroup/instData layout assignment not supported for "
+ "insertStridedSlice.");
} else if (layoutKind == xegpu::LayoutKind::Lane) {
- const auto *uArchInstruction =
- dyn_cast<xegpu::uArch::SubgroupMatrixMultiplyAcc>(uArch->getInstruction(
- xegpu::uArch::InstructionKind::SubgroupMatrixMultiplyAcc));
- auto aLayout = getDefaultLaneLayout2DBlockIo(
- aTy, uArch, uArchInstruction->getPackedFormatBitSizeA());
- auto bLayout = getDefaultLaneLayout2DBlockIo(
- bTy, uArch, uArchInstruction->getPackedFormatBitSizeB(), true);
- auto cdLayout = getDefaultLaneLayout2DBlockIo(cdTy, uArch);
-
- // Create scale layouts
- auto aScaleLayout =
- createScaleLayout(context, aTy, aScaleTy, aLayout, false, uArch);
- auto bScaleLayout =
- createScaleLayout(context, bTy, bScaleTy, bLayout, true, uArch);
-
- return std::make_tuple(aLayout, bLayout, cdLayout, aScaleLayout,
- bScaleLayout);
+ for (int dim = 0; dim < srcRank; dim++) {
+ assert(srcShape[dim] % consumerLaneLayout[dim] == 0 &&
+ "srcShape must be divisible by laneLayout for all dimensions");
+ laneDataValue = std::min(srcShape[dim] / consumerLaneLayout[dim],
+ consumerLaneData[dim]);
+ requiredResLayout =
+ requiredResLayout.setDimData(dim, -1, -1, laneDataValue);
+ }
}
- return std::nullopt;
+ return requiredResLayout;
}
+/// Back-propagates a known result layout to the layout required on `operand`
+/// for a non-anchor (layout-propagating) vector op. Dispatches on the op kind —
+/// broadcast, (multi)reduction, bitcast, shape/transpose, insert/extract,
+/// interleave, etc. — applying the shape/permutation/bitwidth transform to
+/// derive the source layout; elementwise and pass-through ops reuse resLayout
+/// as-is. Returns nullptr for unknown ops or an absent result layout.
xegpu::DistributeLayoutAttr xegpu::inferSourceLayoutFromResultForNonAnchorOp(
OpOperand &operand, xegpu::DistributeLayoutAttr resLayout) {
if (!resLayout)
@@ -1837,8 +2572,8 @@ xegpu::DistributeLayoutAttr xegpu::inferSourceLayoutFromResultForNonAnchorOp(
shapeCast.getSourceVectorType().getShape());
}
- // For vector::InsertStridedSliceOp, infer source layout from result layout.
- // Dest vector must have the same layout as the result.
+ // For vector::InsertStridedSliceOp, infer source layout from result
+ // layout. Dest vector must have the same layout as the result.
if (auto insertSlice = dyn_cast<vector::InsertStridedSliceOp>(op)) {
if (idx == 0) {
return xegpu::inferInsertStridedSliceSourceLayout(
@@ -1907,14 +2642,17 @@ xegpu::DistributeLayoutAttr xegpu::inferSourceLayoutFromResultForNonAnchorOp(
if (dyn_cast<vector::ExtractStridedSliceOp>(op))
return resLayout;
- // For elementwise operations, all operands must have the same layout as the
- // result.
+ // For elementwise operations, all operands must have the same layout as
+ // the result.
if (OpTrait::hasElementwiseMappableTraits(op) && op->getNumResults() == 1)
return resLayout;
return nullptr;
}
+/// Returns the layout required on `operand`: anchor ops report their declared
+/// per-operand layout directly; non-anchor ops back-derive it from their result
+/// layout via inferSourceLayoutFromResultForNonAnchorOp.
xegpu::DistributeLayoutAttr xegpu::getConsumerLayoutAt(OpOperand &operand) {
Operation *op = operand.getOwner();
// Anchor ops declare the layout they
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
index 6a37ae6502b2d..8600492e4bf41 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
@@ -130,48 +130,6 @@ struct LayoutInfo {
void set(const xegpu::DistributeLayoutAttr &layout) { storage = layout; }
};
-SmallVector<int> LayoutInfo::getLaneLayout() const {
- if (!isAssigned())
- return {};
- return llvm::map_to_vector(storage.getEffectiveLaneLayoutAsInt(),
- [](int64_t val) { return static_cast<int>(val); });
-}
-
-SmallVector<int> LayoutInfo::getLaneData() const {
- if (!isAssigned())
- return {};
- return llvm::map_to_vector(storage.getEffectiveLaneDataAsInt(),
- [](int64_t val) { return static_cast<int>(val); });
-}
-
-SmallVector<int> LayoutInfo::getInstData() const {
- if (!isAssigned())
- return {};
- return llvm::map_to_vector(storage.getEffectiveInstDataAsInt(),
- [](int64_t val) { return static_cast<int>(val); });
-}
-
-SmallVector<int> LayoutInfo::getSgLayout() const {
- if (!isAssigned())
- return {};
- return llvm::map_to_vector(storage.getEffectiveSgLayoutAsInt(),
- [](int64_t val) { return static_cast<int>(val); });
-}
-
-SmallVector<int> LayoutInfo::getSgData() const {
- if (!isAssigned())
- return {};
- return llvm::map_to_vector(storage.getEffectiveSgDataAsInt(),
- [](int64_t val) { return static_cast<int>(val); });
-}
-
-SmallVector<int> LayoutInfo::getOrder() const {
- if (!isAssigned() || !storage.getOrder())
- return {};
- return llvm::map_to_vector(storage.getOrder().asArrayRef(),
- [](int64_t val) { return static_cast<int>(val); });
-}
-
void LayoutInfo::print(raw_ostream &os) const {
if (isAssigned()) {
os << storage;
@@ -191,64 +149,6 @@ LayoutInfo LayoutInfo::join(const LayoutInfo &lhs, const LayoutInfo &rhs) {
llvm_unreachable("Join should not be triggered by layout propagation.");
}
-/// Construct a new layout with the transposed inst_data or lane_layout,
-/// lane_data.
-LayoutInfo LayoutInfo::transpose(ArrayRef<int64_t> permutation) const {
- if (!isAssigned())
- return {};
- // Check if the permutation is valid.
- llvm::SmallSet<int64_t, 4> seen(permutation.begin(), permutation.end());
- bool hasDuplicates = seen.size() != permutation.size();
- bool withinRange = llvm::all_of(permutation, [&](int64_t idx) {
- return idx >= 0 && idx < static_cast<int64_t>(permutation.size());
- });
-
- if (!withinRange || hasDuplicates) {
- assert(false && "Invalid permutation for transpose.");
- return {};
- }
-
- SmallVector<int32_t> laneLayout;
- SmallVector<int32_t> laneData;
- SmallVector<int32_t> instData;
- SmallVector<int32_t> sgLayout;
- SmallVector<int32_t> sgData;
- SmallVector<int32_t> order;
-
- for (int64_t idx : permutation) {
- if (getLaneLayout().size()) {
- laneLayout.push_back(static_cast<int32_t>(getLaneLayout()[idx]));
- laneData.push_back(static_cast<int32_t>(getLaneData()[idx]));
- }
- if (getInstData().size())
- instData.push_back(static_cast<int32_t>(getInstData()[idx]));
- if (getSgData().size()) {
- sgLayout.push_back(static_cast<int32_t>(getSgLayout()[idx]));
- sgData.push_back(static_cast<int32_t>(getSgData()[idx]));
- }
- if (getOrder().size()) {
- order.push_back(static_cast<int32_t>(getOrder()[idx]));
- }
- }
- auto orderAttr = order.size()
- ? DenseI32ArrayAttr::get(storage.getContext(), order)
- : nullptr;
- xegpu::LayoutAttr layoutAttr;
- if (getLaneLayout().size())
- layoutAttr =
- xegpu::LayoutAttr::get(storage.getContext(), laneLayout, laneData);
- if (getInstData().size())
- layoutAttr = xegpu::LayoutAttr::get(storage.getContext(), instData);
- if (getSgData().size())
- layoutAttr = xegpu::LayoutAttr::get(
- storage.getContext(),
- DenseI32ArrayAttr::get(storage.getContext(), sgLayout),
- DenseI32ArrayAttr::get(storage.getContext(), sgData),
- /*inst_data =*/nullptr, /*lane_layout =*/nullptr,
- /*lane_data =*/nullptr, orderAttr);
- return LayoutInfo(layoutAttr);
-}
-
//===----------------------------------------------------------------------===//
// LayoutInfoLattice
//===----------------------------------------------------------------------===//
@@ -259,57 +159,6 @@ struct LayoutInfoLattice : public Lattice<LayoutInfo> {
using Lattice::Lattice;
};
-/// Helper Functions to get default layouts. A `default layout` is a layout that
-/// is assigned to a value when the layout is not fixed by some anchor operation
-/// (like DPAS).
-
-/// Helper Function to get the default layout for uniform values like constants.
-/// For 1D vector, lane_layout is [subgroupSize] and lane_data is [1].
-/// For 2D vector, lane_layout is [1, subgroupSize] and lane_data is [1, 1].
-/// For ND vector (N>2), leading dims get unit lane_layout and lane_data.
-static LayoutInfo getDefaultSIMTLayoutInfo(mlir::MLIRContext *ctx,
- unsigned rank,
- const xegpu::uArch::uArch *uArch) {
- assert(rank >= 1 && "Expected at least 1D vector.");
- if (rank == 1) {
- return LayoutInfo(
- xegpu::LayoutAttr::get(ctx, {uArch->getSubgroupSize()}, {1}));
- }
- // For rank >= 2, lane_layout is [1, ..., 1, subgroupSize] and
- // lane_data is [1, ..., 1, 1].
- SmallVector<int32_t> laneLayout(rank, 1);
- SmallVector<int32_t> laneData(rank, 1);
- laneLayout[rank - 1] = uArch->getSubgroupSize();
- return LayoutInfo(xegpu::LayoutAttr::get(ctx, laneLayout, laneData));
-}
-
-/// Helper to get the default layout for 2D block operations.
-/// For ND (N>2) types, leading dimensions get unit layout/data values.
-template <typename Ty>
-static LayoutInfo getSIMTLayoutInfoBlockIO(Ty ty,
- const xegpu::uArch::uArch *uArch,
- unsigned packingSize) {
- // Expecting at least 1D.
- assert(ty.getRank() >= 1 && "Expected at least 1D vector.");
- // Expecting int or float element type.
- assert(ty.getElementType().isIntOrFloat() &&
- "Expected int or float element type.");
- // If the rank is 1, then return default layout for 1D vector.
- if (ty.getRank() == 1)
- return getDefaultSIMTLayoutInfo(ty.getContext(), 1, uArch);
- // Packing factor is determined by the element type bitwidth.
- unsigned bitwidth = ty.getElementType().getIntOrFloatBitWidth();
- int packingFactor = bitwidth < packingSize ? packingSize / bitwidth : 1;
- // For rank >= 2, distribute along the last dimension with leading units.
- unsigned rank = ty.getRank();
- SmallVector<int32_t> laneLayout(rank, 1);
- SmallVector<int32_t> laneData(rank, 1);
- laneLayout[rank - 1] = uArch->getSubgroupSize();
- laneData[rank - 1] = packingFactor;
- return LayoutInfo(
- xegpu::LayoutAttr::get(ty.getContext(), laneLayout, laneData));
-}
-
//===----------------------------------------------------------------------===//
// LayoutInfoPropagation
//===----------------------------------------------------------------------===//
@@ -549,9 +398,9 @@ bool LayoutInfoPropagation::hasParamsOfLayoutKind(
// wgShape = [128, 64], instData = [8, 16], sgCount = 32
// Returns layouts:
// [(8,4), (16,2)], which correspond to sgData [16,16] and [8,32].
-SmallVector<std::pair<int, int>> getValidLayouts(ArrayRef<int64_t> wgShape,
- ArrayRef<int> instData,
- int64_t sgCount) {
+SmallVector<std::pair<int, int>>
+getSgLayoutCandidates(ArrayRef<int64_t> wgShape, ArrayRef<int> instData,
+ int64_t sgCount) {
SmallVector<std::pair<int, int>> candidates;
for (int sgLayout0 = 1; sgLayout0 <= sgCount; ++sgLayout0) {
if (sgCount % sgLayout0)
@@ -578,7 +427,15 @@ SmallVector<std::pair<int, int>> getValidLayouts(ArrayRef<int64_t> wgShape,
return candidates;
}
-FailureOr<int64_t> getNumSg(Operation *op, const int sgSize) {
+FailureOr<int64_t>
+getNumSg(Operation *op, const int sgSize,
+ xegpu::DistributeLayoutAttr consumerLayout = nullptr) {
+ // first look for the number of subgroups required by the consumer layout
+ if (consumerLayout) {
+ auto sgLayout = consumerLayout.getEffectiveSgLayoutAsInt();
+ if (!sgLayout.empty())
+ return llvm::product_of(sgLayout);
+ }
// Oblivious to workitem layout, the total count matters.
auto gpuFunc = op->getParentOfType<gpu::GPUFuncOp>();
if (!gpuFunc)
@@ -595,53 +452,48 @@ void LayoutInfoPropagation::visitPrefetchNdOp(
ArrayRef<const LayoutInfoLattice *> results) {
LayoutInfo prefetchLayout;
+ const uArch *uArch = getUArch(getChipStr(prefetch).value_or(""));
+ if (!uArch)
+ return;
xegpu::DistributeLayoutAttr anchorLayout = prefetch.getLayoutAttr();
if (hasParamsOfLayoutKind(anchorLayout)) {
prefetchLayout = LayoutInfo(anchorLayout);
+ if (layoutKind == xegpu::LayoutKind::InstData) {
+ const auto *uArchInstruction =
+ dyn_cast<xegpu::uArch::Subgroup2DBlockPrefetchInstruction>(
+ uArch->getInstruction(
+ xegpu::uArch::InstructionKind::Subgroup2DBlockPrefetch));
+ if (!uArchInstruction)
+ return;
+ auto completed = xegpu::completeBlockStoreLaneLayoutFromInstData(
+ anchorLayout, prefetch.getTensorDescType().getElementType(),
+ uArchInstruction, uArch->getSubgroupSize());
+ if (!completed) {
+ prefetch.emitWarning(
+ "Failed to identify lane layouts for the specified inst_data.");
+ return;
+ }
+ prefetch.setLayoutAttr(*completed);
+ prefetchLayout = LayoutInfo(*completed);
+ }
} else {
- // Here we assign the default layout to the tensor descriptor operand of
- // prefetch.
auto tdescTy = prefetch.getTensorDescType();
-
- const uArch *uArch = getUArch(getChipStr(prefetch).value_or(""));
- if (!uArch)
- return;
- const auto *uArchInstruction =
- dyn_cast<xegpu::uArch::Subgroup2DBlockPrefetchInstruction>(
- uArch->getInstruction(
- xegpu::uArch::InstructionKind::Subgroup2DBlockPrefetch));
-
- auto blockWHC =
- uArchInstruction->getBlockWidthHeightCount(tdescTy.getElementType());
- if (!blockWHC)
- prefetch.emitWarning("No known block params found for the element type.");
- auto [bWidth, bHeight, bCount] = blockWHC.value();
- SmallVector<int> instData;
- int instWidth = xegpu::getLargestDivisor(
- static_cast<int>(tdescTy.getDimSize(tdescTy.getRank() - 1)), bWidth);
- if (instWidth == -1)
+ auto numSgOrErr = getNumSg(prefetch, uArch->getSubgroupSize());
+ if (layoutKind == xegpu::LayoutKind::Subgroup && failed(numSgOrErr)) {
prefetch.emitWarning(
- "No suitable instruction multiple found for the given shape.");
- if (tdescTy.getRank() == 1)
- instData = {instWidth};
- else {
- int instHeight = xegpu::getLargestDivisor(
- static_cast<int>(tdescTy.getDimSize(tdescTy.getRank() - 2)), bHeight);
- if (instHeight == -1)
- prefetch.emitWarning(
- "No suitable instruction multiple found for the given shape.");
- instData = {instHeight, instWidth};
+ "Unable to determine the number of subgroups for the operation.");
+ return;
}
- if (layoutKind == xegpu::LayoutKind::InstData)
- prefetchLayout =
- LayoutInfo(xegpu::LayoutAttr::get(tdescTy.getContext(), instData));
- else
- prefetchLayout = getSIMTLayoutInfoBlockIO(
- tdescTy, uArch, uArchInstruction->getPackedFormatBitSize());
-
- prefetch.setLayoutAttr(
- dyn_cast<xegpu::DistributeLayoutAttr>(prefetchLayout.get()));
+ auto layoutAttr = xegpu::setupPrefetchNdAnchorLayout(
+ layoutKind, tdescTy, numSgOrErr.value_or(0), uArch);
+ if (!layoutAttr) {
+ prefetch.emitWarning(
+ "Failed to determine required layout for prefetch_nd.");
+ return;
+ }
+ prefetchLayout = LayoutInfo(layoutAttr);
+ prefetch.setLayoutAttr(layoutAttr);
}
// Propagate the layout to the source tensor descriptor.
propagateIfChanged(operands[0], operands[0]->meet(prefetchLayout));
@@ -669,11 +521,13 @@ void LayoutInfoPropagation::visitVectorMultiReductionOp(
const uArch *uArch = getUArch(xegpu::getChipStr(reduction).value_or(""));
if (!uArch)
return;
- int numSg = 0;
- if (layoutKind == xegpu::LayoutKind::Subgroup) {
- auto numSgOrErr = getNumSg(reduction, uArch->getSubgroupSize());
- if (succeeded(numSgOrErr))
- numSg = numSgOrErr.value();
+
+ auto numSgOrErr =
+ getNumSg(reduction, uArch->getSubgroupSize(), consumerLayoutAttr);
+ if (layoutKind == xegpu::LayoutKind::Subgroup && failed(numSgOrErr)) {
+ reduction.emitWarning(
+ "Unable to determine the number of subgroups for the operation.");
+ return;
}
// The result layout represents the layout requirements of the operation.
@@ -682,7 +536,8 @@ void LayoutInfoPropagation::visitVectorMultiReductionOp(
// propagated from consumer op, the conflict is resolved in later phase by
// converting the required result layout to the consumer layout
auto requiredResLayoutAttr = xegpu::setupMultiReductionResultLayout(
- layoutKind, sourceTy, consumerLayoutAttr, reductionDims, numSg, uArch);
+ layoutKind, sourceTy, consumerLayoutAttr, reductionDims,
+ numSgOrErr.value_or(0), uArch);
xegpu::setTemporaryLayout(reduction->getResult(0), requiredResLayoutAttr);
@@ -757,6 +612,13 @@ void LayoutInfoPropagation::visitShapeCastOp(
xegpu::DistributeLayoutAttr srcLayoutAttr =
xegpu::inferShapeCastSourceLayout(resultLayoutAttr, resShape, srcShape);
+ // TODO: turn this into a real pass failure once propagation failures are
+ // wired to signalPassFailure().
+ if (!srcLayoutAttr) {
+ shapeCast.emitWarning("Failed to infer source layout for shape_cast; "
+ "unsupported shape-cast pattern.");
+ return;
+ }
propagateIfChanged(operands[0], operands[0]->meet(LayoutInfo(srcLayoutAttr)));
}
@@ -769,6 +631,13 @@ void LayoutInfoPropagation::visitDpasOp(
LayoutInfo dpasBLayout;
LayoutInfo dpasCDLayout;
+ const uArch *uArch = getUArch(getChipStr(dpas).value_or(""));
+ if (!uArch)
+ return;
+ VectorType aTy = dpas.getLhsType();
+ VectorType bTy = dpas.getRhsType();
+ VectorType cdTy = dpas.getResultType();
+
xegpu::DistributeLayoutAttr anchorLayoutCD = dpas.getLayoutCdAttr();
if (hasParamsOfLayoutKind(anchorLayoutCD)) {
xegpu::DistributeLayoutAttr anchorLayoutA = dpas.getLayoutAAttr();
@@ -780,35 +649,45 @@ void LayoutInfoPropagation::visitDpasOp(
dpasALayout = LayoutInfo(anchorLayoutA);
dpasBLayout = LayoutInfo(anchorLayoutB);
dpasCDLayout = LayoutInfo(anchorLayoutCD);
+ if (layoutKind == xegpu::LayoutKind::InstData) {
+ auto completed = xegpu::completeDpasLaneLayoutFromInstData(
+ anchorLayoutA, anchorLayoutB, anchorLayoutCD, aTy, bTy, cdTy, uArch);
+ if (!completed) {
+ dpas.emitWarning(
+ "Failed to identify lane layouts for the specified inst_data.");
+ return;
+ }
+ auto [completedA, completedB, completedCD] = *completed;
+ dpas.setLayoutAAttr(completedA);
+ dpas.setLayoutBAttr(completedB);
+ dpas.setLayoutCdAttr(completedCD);
+ dpasALayout = LayoutInfo(completedA);
+ dpasBLayout = LayoutInfo(completedB);
+ dpasCDLayout = LayoutInfo(completedCD);
+ }
} else {
- const uArch *uArch = getUArch(getChipStr(dpas).value_or(""));
- if (!uArch)
- return;
- VectorType aTy = dpas.getLhsType();
- VectorType bTy = dpas.getRhsType();
- VectorType cdTy = dpas.getResultType();
xegpu::DistributeLayoutAttr consumerLayoutAttr = nullptr;
xegpu::DistributeLayoutAttr requiredCDLayoutAttr, requiredALayout,
requiredBLayout;
- int numSg = 0;
- if (layoutKind == xegpu::LayoutKind::Subgroup) {
- LayoutInfo consumerLayout = results[0]->getValue();
- if (!consumerLayout.isAssigned())
- return;
- consumerLayoutAttr =
- dyn_cast<xegpu::DistributeLayoutAttr>(consumerLayout.get());
- auto numSgOrErr = getNumSg(dpas, uArch->getSubgroupSize());
- if (failed(numSgOrErr)) {
- dpas.emitWarning(
- "Unable to determine the number of subgroups for the operation.");
- return;
- }
- numSg = numSgOrErr.value();
+ LayoutInfo consumerLayout = results[0]->getValue();
+ if (!consumerLayout.isAssigned())
+ return;
+ consumerLayoutAttr =
+ dyn_cast<xegpu::DistributeLayoutAttr>(consumerLayout.get());
+
+ auto numSgOrErr =
+ getNumSg(dpas, uArch->getSubgroupSize(), consumerLayoutAttr);
+ if (layoutKind == xegpu::LayoutKind::Subgroup && failed(numSgOrErr)) {
+ dpas.emitWarning(
+ "Unable to determine the number of subgroups for the operation.");
+ return;
}
- auto layouts = xegpu::setupDpasLayout(layoutKind, aTy, bTy, cdTy,
- consumerLayoutAttr, numSg, uArch);
+
+ auto layouts =
+ xegpu::setupDpasLayout(layoutKind, aTy, bTy, cdTy, consumerLayoutAttr,
+ numSgOrErr.value_or(0), uArch);
if (!layouts.has_value()) {
dpas.emitWarning(
"Failed to determine required layouts for DPAS operands.");
@@ -846,6 +725,24 @@ void LayoutInfoPropagation::visitDpasMxOp(
xegpu::DistributeLayoutAttr anchorLayoutB = dpasMx.getLayoutBAttr();
xegpu::DistributeLayoutAttr anchorLayoutCD = dpasMx.getLayoutCdAttr();
+ const uArch *uArch = getUArch(getChipStr(dpasMx).value_or(""));
+ if (!uArch)
+ return;
+
+ VectorType aTy = dpasMx.getAType();
+ VectorType bTy = dpasMx.getBType();
+ VectorType cdTy = dpasMx.getResultType();
+
+ // Get scale types if present
+ VectorType aScaleTy;
+ VectorType bScaleTy;
+ Value scaleA = dpasMx.getScaleA();
+ Value scaleB = dpasMx.getScaleB();
+ if (scaleA)
+ aScaleTy = dyn_cast<VectorType>(scaleA.getType());
+ if (scaleB)
+ bScaleTy = dyn_cast<VectorType>(scaleB.getType());
+
// Check if all layouts are already set
if (anchorLayoutA && anchorLayoutB && anchorLayoutCD &&
hasParamsOfLayoutKind(anchorLayoutA) &&
@@ -864,49 +761,55 @@ void LayoutInfoPropagation::visitDpasMxOp(
dpasMxAScaleLayout = LayoutInfo(anchorLayoutAScale);
if (anchorLayoutBScale)
dpasMxBScaleLayout = LayoutInfo(anchorLayoutBScale);
- } else {
- // Need to compute layouts
- const uArch *uArch = getUArch(getChipStr(dpasMx).value_or(""));
- if (!uArch)
- return;
-
- VectorType aTy = dpasMx.getAType();
- VectorType bTy = dpasMx.getBType();
- VectorType cdTy = dpasMx.getResultType();
-
- // Get scale types if present
- VectorType aScaleTy;
- VectorType bScaleTy;
- Value scaleA = dpasMx.getScaleA();
- Value scaleB = dpasMx.getScaleB();
- if (scaleA)
- aScaleTy = dyn_cast<VectorType>(scaleA.getType());
- if (scaleB)
- bScaleTy = dyn_cast<VectorType>(scaleB.getType());
+ if (layoutKind == xegpu::LayoutKind::InstData) {
+ auto completed = xegpu::completeDpasMxLaneLayoutFromInstData(
+ anchorLayoutA, anchorLayoutB, anchorLayoutCD, aTy, bTy, cdTy,
+ aScaleTy, bScaleTy, uArch);
+ if (!completed) {
+ dpasMx.emitWarning(
+ "Failed to identify lane layouts for the specified inst_data.");
+ return;
+ }
+ auto [completedA, completedB, completedCD, completedAScale,
+ completedBScale] = *completed;
+ dpasMx.setLayoutAAttr(completedA);
+ dpasMx.setLayoutBAttr(completedB);
+ dpasMx.setLayoutCdAttr(completedCD);
+ dpasMxALayout = LayoutInfo(completedA);
+ dpasMxBLayout = LayoutInfo(completedB);
+ dpasMxCDLayout = LayoutInfo(completedCD);
+ if (completedAScale) {
+ dpasMx.setLayoutAScaleAttr(completedAScale);
+ dpasMxAScaleLayout = LayoutInfo(completedAScale);
+ }
+ if (completedBScale) {
+ dpasMx.setLayoutBScaleAttr(completedBScale);
+ dpasMxBScaleLayout = LayoutInfo(completedBScale);
+ }
+ }
+ } else {
xegpu::DistributeLayoutAttr consumerLayoutAttr = nullptr;
xegpu::DistributeLayoutAttr requiredCDLayoutAttr, requiredALayout,
requiredBLayout, requiredAScaleLayout, requiredBScaleLayout;
- int numSg = 0;
- if (layoutKind == xegpu::LayoutKind::Subgroup) {
- LayoutInfo consumerLayout = results[0]->getValue();
- if (!consumerLayout.isAssigned())
- return;
- consumerLayoutAttr =
- dyn_cast<xegpu::DistributeLayoutAttr>(consumerLayout.get());
- auto numSgOrErr = getNumSg(dpasMx, uArch->getSubgroupSize());
- if (failed(numSgOrErr)) {
- dpasMx.emitWarning(
- "Unable to determine the number of subgroups for the operation.");
- return;
- }
- numSg = numSgOrErr.value();
+ LayoutInfo consumerLayout = results[0]->getValue();
+ if (!consumerLayout.isAssigned())
+ return;
+ consumerLayoutAttr =
+ dyn_cast<xegpu::DistributeLayoutAttr>(consumerLayout.get());
+
+ auto numSgOrErr =
+ getNumSg(dpasMx, uArch->getSubgroupSize(), consumerLayoutAttr);
+ if (layoutKind == xegpu::LayoutKind::Subgroup && failed(numSgOrErr)) {
+ dpasMx.emitWarning(
+ "Unable to determine the number of subgroups for the operation.");
+ return;
}
- auto layouts =
- xegpu::setupDpasMxLayout(layoutKind, aTy, bTy, cdTy, aScaleTy, bScaleTy,
- consumerLayoutAttr, numSg, uArch);
+ auto layouts = xegpu::setupDpasMxLayout(
+ layoutKind, aTy, bTy, cdTy, aScaleTy, bScaleTy, consumerLayoutAttr,
+ numSgOrErr.value_or(0), uArch);
if (!layouts.has_value()) {
dpasMx.emitWarning(
"Failed to determine required layouts for DPAS_MX operands.");
@@ -963,78 +866,47 @@ void LayoutInfoPropagation::visitStoreNdOp(
xegpu::StoreNdOp store, ArrayRef<LayoutInfoLattice *> operands,
ArrayRef<const LayoutInfoLattice *> results) {
LayoutInfo storeLayout;
+ const uArch *uArch = getUArch(getChipStr(store).value_or(""));
+ if (!uArch)
+ return;
xegpu::DistributeLayoutAttr anchorLayout = store.getLayoutAttr();
if (hasParamsOfLayoutKind(anchorLayout)) {
storeLayout = LayoutInfo(anchorLayout);
- } else {
- const uArch *uArch = getUArch(getChipStr(store).value_or(""));
- if (!uArch)
- return;
- const auto *uArchInstruction =
- dyn_cast<xegpu::uArch::Subgroup2DBlockStoreInstruction>(
- uArch->getInstruction(
- xegpu::uArch::InstructionKind::Subgroup2DBlockStore));
- VectorType dataTy = store.getValueType();
- auto blockWHC = uArchInstruction->getBlockWidthHeightCount(
- store.getValueType().getElementType());
- if (!blockWHC)
- store.emitWarning("No known block params found for the element type.");
- auto [bWidth, bHeight, bCount] = blockWHC.value();
- // Default to 1 for any leading batch dims; rank-1 and rank>=2 cases
- // overwrite the trailing entries below.
- SmallVector<int> instData(dataTy.getRank(), 1);
- int instWidth = xegpu::getLargestDivisor(
- static_cast<int>(dataTy.getDimSize(dataTy.getRank() - 1)), bWidth);
- if (instWidth == -1)
- store.emitWarning(
- "No suitable instruction multiple found for the given shape.");
- if (dataTy.getRank() == 1) {
- instData = {instWidth};
- } else {
- int instHeight = xegpu::getLargestDivisor(
- static_cast<int>(dataTy.getDimSize(dataTy.getRank() - 2)), bHeight);
- if (instHeight == -1)
- store.emitWarning(
- "No suitable instruction multiple found for the given shape.");
- instData[dataTy.getRank() - 2] = instHeight;
- instData[dataTy.getRank() - 1] = instWidth;
- }
+ if (layoutKind == xegpu::LayoutKind::InstData) {
- if (layoutKind == xegpu::LayoutKind::InstData)
- storeLayout =
- LayoutInfo(xegpu::LayoutAttr::get(dataTy.getContext(), instData));
- else if (layoutKind == xegpu::LayoutKind::Lane)
- storeLayout =
- getSIMTLayoutInfoBlockIO(store.getValueType(), uArch,
- uArchInstruction->getPackedFormatBitSize());
- else { // xegpu::LayoutKind::Subgroup
- auto sgSize = uArch->getSubgroupSize();
- auto numSgOrErr = getNumSg(store, sgSize);
- if (failed(numSgOrErr)) {
- store.emitWarning(
- "Unable to determine the number of subgroups for the operation.");
+ const auto *uArchInstruction =
+ dyn_cast<xegpu::uArch::Subgroup2DBlockStoreInstruction>(
+ uArch->getInstruction(
+ xegpu::uArch::InstructionKind::Subgroup2DBlockStore));
+ if (!uArchInstruction)
return;
- }
- auto sgLayouts = getValidLayouts(store.getValueType().getShape(),
- instData, numSgOrErr.value());
- if (sgLayouts.empty()) {
+ auto completed = xegpu::completeBlockStoreLaneLayoutFromInstData(
+ anchorLayout, store.getValueType().getElementType(), uArchInstruction,
+ uArch->getSubgroupSize());
+ if (!completed) {
store.emitWarning(
- "Unable to determine suitable subgroup layout for store value.");
+ "Failed to identify lane layouts for the specified inst_data.");
return;
}
- SmallVector<int> sgLayout = {sgLayouts[0].first, sgLayouts[0].second};
- SmallVector<int> sgData = {
- static_cast<int>(dataTy.getShape()[0]) / sgLayout[0],
- static_cast<int>(dataTy.getShape()[1]) / sgLayout[1]};
- storeLayout = LayoutInfo(xegpu::LayoutAttr::get(
- dataTy.getContext(),
- DenseI32ArrayAttr::get(dataTy.getContext(), sgLayout),
- DenseI32ArrayAttr::get(dataTy.getContext(), sgData),
- /*inst_data =*/nullptr, /*lane_layout =*/nullptr,
- /*lane_data =*/nullptr, /*order =*/nullptr));
+ store.setLayoutAttr(*completed);
+ storeLayout = LayoutInfo(*completed);
}
- store.setLayoutAttr(
- dyn_cast<xegpu::DistributeLayoutAttr>(storeLayout.get()));
+ } else {
+ auto numSgOrErr = getNumSg(store, uArch->getSubgroupSize());
+ if (layoutKind == xegpu::LayoutKind::Subgroup && failed(numSgOrErr)) {
+ store.emitWarning(
+ "Unable to determine the number of subgroups for the operation.");
+ return;
+ }
+
+ auto layoutAttr = xegpu::setupStoreNdAnchorLayout(
+ layoutKind, store.getValueType(), numSgOrErr.value_or(0), uArch);
+ if (!layoutAttr) {
+ store.emitWarning("Failed to determine required layout for store_nd.");
+ return;
+ }
+ storeLayout = LayoutInfo(layoutAttr);
+ store.setLayoutAttr(layoutAttr);
}
// Propagate the layout to the value operand.
// Both operands should have the same layout
@@ -1048,25 +920,54 @@ void LayoutInfoPropagation::visitLoadNdOp(
xegpu::LoadNdOp load, ArrayRef<LayoutInfoLattice *> operands,
ArrayRef<const LayoutInfoLattice *> results) {
LayoutInfo loadLayout;
+
+ const uArch *uArch = getUArch(getChipStr(load).value_or(""));
+ if (!uArch)
+ return;
+ LayoutInfo valueLayout = results[0]->getValue();
+ if (!valueLayout.isAssigned())
+ return;
+ auto consumerLayoutAttr =
+ dyn_cast<xegpu::DistributeLayoutAttr>(valueLayout.get());
xegpu::DistributeLayoutAttr anchorLayout = load.getLayoutAttr();
if (hasParamsOfLayoutKind(anchorLayout)) {
loadLayout = LayoutInfo(anchorLayout);
+ if (layoutKind == xegpu::LayoutKind::InstData &&
+ !consumerLayoutAttr.getEffectiveLaneLayoutAsInt().empty()) {
+ const auto *uArchInstruction =
+ dyn_cast<xegpu::uArch::Subgroup2DBlockLoadInstruction>(
+ uArch->getInstruction(
+ xegpu::uArch::InstructionKind::Subgroup2DBlockLoad));
+ if (!uArchInstruction)
+ return;
+ auto completed = xegpu::completeBlockLoadLaneLayoutFromInstData(
+ anchorLayout, consumerLayoutAttr, load.getType().getElementType(),
+ uArchInstruction, uArch->getSubgroupSize());
+ if (!completed) {
+ load.emitWarning(
+ "Failed to identify lane layouts for the specified inst_data.");
+ return;
+ }
+ load.setLayoutAttr(*completed);
+ loadLayout = LayoutInfo(*completed);
+ }
} else {
-
- LayoutInfo valueLayout = results[0]->getValue();
- // Need the layout of the value to propagate to the tensor descriptor.
- if (!valueLayout.isAssigned())
+ auto numSgOrErr =
+ getNumSg(load, uArch->getSubgroupSize(), consumerLayoutAttr);
+ if (layoutKind == xegpu::LayoutKind::Subgroup && failed(numSgOrErr)) {
+ load.emitWarning(
+ "Unable to determine the number of subgroups for the operation.");
return;
- loadLayout = valueLayout;
- // LoadNdOp has the transpose effect. However, at the stage of this analysis
- // this effect is not expected and should be abstracted away. Emit a
- // warning.
- if (auto transpose = load.getTranspose()) {
- load.emitWarning("Transpose effect is not expected for LoadNdOp at "
- "LayoutInfoPropagation stage.");
- loadLayout = valueLayout.transpose(transpose.value());
}
- load.setLayoutAttr(dyn_cast<xegpu::DistributeLayoutAttr>(loadLayout.get()));
+ auto layoutAttr = xegpu::setupLoadNdAnchorLayout(
+ layoutKind, load.getType(), consumerLayoutAttr, numSgOrErr.value_or(0),
+ uArch);
+ if (!layoutAttr) {
+ load.emitWarning("Failed to determine required layout for load_nd.");
+ return;
+ }
+ loadLayout = LayoutInfo(layoutAttr);
+ load.setLayoutAttr(layoutAttr);
}
// Propagate the new layout to the tensor descriptor operand.
propagateIfChanged(operands[0], operands[0]->meet(loadLayout));
@@ -1077,6 +978,49 @@ void LayoutInfoPropagation::visitLoadNdOp(
void LayoutInfoPropagation::visitConvertLayoutOp(
xegpu::ConvertLayoutOp convert, ArrayRef<LayoutInfoLattice *> operands,
ArrayRef<const LayoutInfoLattice *> results) {
+
+ LayoutInfo resultLayout = results[0]->getValue();
+
+ // TODO: fix if one of the layouts is a slice layout
+ auto targetLayoutAttr =
+ dyn_cast<xegpu::LayoutAttr>(convert.getTargetLayoutAttr());
+ auto inputLayoutAttr =
+ dyn_cast<xegpu::LayoutAttr>(convert.getInputLayoutAttr());
+
+ // The result's propagated layout is authoritative for the converted value.
+ // Fill the lane_layout / lane_data / order parameters the target_layout is
+ // missing from it (sg_layout / sg_data / inst_data are left as-is), so the
+ // target stays consistent with what is actually propagated downstream.
+ auto resultLayoutAttr = resultLayout.isAssigned()
+ ? dyn_cast<xegpu::LayoutAttr>(resultLayout.get())
+ : nullptr;
+ if (resultLayoutAttr && targetLayoutAttr) {
+ if (layoutKind == xegpu::LayoutKind::InstData &&
+ !targetLayoutAttr.getLaneLayout()) {
+ targetLayoutAttr = xegpu::LayoutAttr::get(
+ convert.getContext(), targetLayoutAttr.getSgLayout(),
+ targetLayoutAttr.getSgData(), targetLayoutAttr.getInstData(),
+ resultLayoutAttr.getLaneLayout(), resultLayoutAttr.getLaneData(),
+ resultLayoutAttr.getOrder());
+ convert.setTargetLayoutAttr(targetLayoutAttr);
+ }
+ }
+
+ // Fill only the lane_layout / lane_data / order parameters the input_layout
+ // is missing from the target_layout (sg_layout / sg_data / inst_data are left
+ // as-is), so the producer side receives a fully-populated lane layout.
+ if (inputLayoutAttr && targetLayoutAttr) {
+ if (layoutKind == xegpu::LayoutKind::InstData &&
+ !inputLayoutAttr.getLaneLayout()) {
+ auto merged = xegpu::LayoutAttr::get(
+ convert.getContext(), inputLayoutAttr.getSgLayout(),
+ inputLayoutAttr.getSgData(), inputLayoutAttr.getInstData(),
+ targetLayoutAttr.getLaneLayout(), targetLayoutAttr.getLaneData(),
+ targetLayoutAttr.getOrder());
+ convert.setInputLayoutAttr(merged);
+ }
+ }
+
xegpu::DistributeLayoutAttr anchorLayout = convert.getInputLayoutAttr();
LayoutInfo convertLayout(anchorLayout);
// Propagate the new layout to the tensor descriptor operand.
@@ -1135,9 +1079,9 @@ void LayoutInfoPropagation::visitVectorBitcastOp(
propagateIfChanged(operands[0], operands[0]->meet(LayoutInfo(srcLayoutAttr)));
}
-/// For vector::InterleaveOp, the result has double the innermost dimension size
-/// compared to each source operand. The layout is propagated from result to
-/// sources, adjusting for the 2x size increase.
+/// For vector::InterleaveOp, the result has double the innermost dimension
+/// size compared to each source operand. The layout is propagated from result
+/// to sources, adjusting for the 2x size increase.
void LayoutInfoPropagation::visitVectorInterleaveOp(
vector::InterleaveOp interleave, ArrayRef<LayoutInfoLattice *> operands,
ArrayRef<const LayoutInfoLattice *> results) {
@@ -1185,8 +1129,8 @@ void LayoutInfoPropagation::visitVectorDeinterleaveOp(
auto consumerLayoutAttr =
dyn_cast<xegpu::DistributeLayoutAttr>(resLayoutInfo.get());
- // Derive the source layout from the result layout (double the innermost dim)
- // No setup function needed - just infer directly
+ // Derive the source layout from the result layout (double the innermost
+ // dim) No setup function needed - just infer directly
auto srcLayoutAttr = xegpu::inferDeinterleaveSourceLayout(consumerLayoutAttr);
propagateIfChanged(operands[0], operands[0]->meet(LayoutInfo(srcLayoutAttr)));
@@ -1223,8 +1167,8 @@ void LayoutInfoPropagation::visitInsertStridedSliceOp(
operands[1]->meet(LayoutInfo(requiredResLayoutAttr)));
}
-/// Propagate the layout of the result to the tensor descriptor, mask and offset
-/// operands in LoadGatherOp.
+/// Propagate the layout of the result to the tensor descriptor, mask and
+/// offset operands in LoadGatherOp.
void LayoutInfoPropagation::visitLoadGatherOp(
xegpu::LoadGatherOp load, ArrayRef<LayoutInfoLattice *> operands,
ArrayRef<const LayoutInfoLattice *> results) {
@@ -1244,6 +1188,24 @@ void LayoutInfoPropagation::visitLoadGatherOp(
if (hasParamsOfLayoutKind(anchorLayoutAttr)) {
requiredAnchorLayoutAttr = anchorLayoutAttr;
+ if (layoutKind == xegpu::LayoutKind::InstData &&
+ !consumerLayoutAttr.getEffectiveLaneLayoutAsInt().empty()) {
+ const auto uArchInstruction =
+ dyn_cast<xegpu::uArch::LoadGatherInstructionInterface>(
+ uArch->getInstruction(xegpu::uArch::InstructionKind::LoadGather));
+ if (!uArchInstruction)
+ return;
+ auto completed = xegpu::completeScatterLoadLaneLayoutFromInstData(
+ anchorLayoutAttr, consumerLayoutAttr, resVecTy.getElementType(),
+ uArchInstruction, uArch->getSubgroupSize());
+ if (!completed) {
+ load.emitWarning(
+ "Failed to identify lane layouts for the specified inst_data.");
+ return;
+ }
+ requiredAnchorLayoutAttr = *completed;
+ load.setLayoutAttr(requiredAnchorLayoutAttr);
+ }
} else {
if (!resVecTy) {
load.emitWarning("Not propagating, non-vector payload supplied.");
@@ -1268,8 +1230,8 @@ void LayoutInfoPropagation::visitLoadGatherOp(
propagateIfChanged(operands[2], operands[2]->meet(maskLayoutInfo));
}
-/// Set the layout for the value, tensor descriptor, offset and mask operands in
-/// the StoreScatterOp.
+/// Set the layout for the value, tensor descriptor, offset and mask operands
+/// in the StoreScatterOp.
void LayoutInfoPropagation::visitStoreScatterOp(
xegpu::StoreScatterOp storeScatter, ArrayRef<LayoutInfoLattice *> operands,
ArrayRef<const LayoutInfoLattice *> results) {
@@ -1284,6 +1246,24 @@ void LayoutInfoPropagation::visitStoreScatterOp(
if (hasParamsOfLayoutKind(anchorLayoutAttr)) {
requiredAnchorLayoutAttr = anchorLayoutAttr;
+ if (layoutKind == xegpu::LayoutKind::InstData) {
+ const auto uArchInstruction =
+ dyn_cast<xegpu::uArch::StoreScatterInstructionInterface>(
+ uArch->getInstruction(
+ xegpu::uArch::InstructionKind::StoreScatter));
+ if (!uArchInstruction)
+ return;
+ auto completed = xegpu::completeScatterStoreLaneLayoutFromInstData(
+ anchorLayoutAttr, srcVecTy.getElementType(), uArchInstruction,
+ uArch->getSubgroupSize());
+ if (!completed) {
+ storeScatter.emitWarning(
+ "Failed to identify lane layouts for the specified inst_data.");
+ return;
+ }
+ requiredAnchorLayoutAttr = *completed;
+ storeScatter.setLayoutAttr(requiredAnchorLayoutAttr);
+ }
} else {
if (!srcVecTy) {
storeScatter.emitWarning("Not propagating, non-vector payload supplied.");
@@ -1331,8 +1311,10 @@ void LayoutInfoPropagation::visitLoadMatrixOp(
const uArch *uArch = getUArch(getChipStr(loadMatrixOp).value_or(""));
if (!uArch)
return;
+ int chunkSize =
+ 1; // placeHolder for future use when LoadMatrix supports coalescing
auto requiredAnchorLayoutAttr = xegpu::setupLoadMatrixAnchorLayout(
- layoutKind, resVecTy, consumerLayoutAttr, uArch);
+ layoutKind, resVecTy, chunkSize, consumerLayoutAttr, uArch);
loadMatrixOp.setLayoutAttr(requiredAnchorLayoutAttr);
}
}
@@ -1340,22 +1322,41 @@ void LayoutInfoPropagation::visitLoadMatrixOp(
void LayoutInfoPropagation::visitStoreMatrixOp(
xegpu::StoreMatrixOp storeMatrix, ArrayRef<LayoutInfoLattice *> operands,
ArrayRef<const LayoutInfoLattice *> results) {
- xegpu::DistributeLayoutAttr anchorLayout = storeMatrix.getLayoutAttr();
+ xegpu::DistributeLayoutAttr requiredAnchorLayoutAttr;
+ xegpu::DistributeLayoutAttr anchorLayoutAttr = storeMatrix.getLayoutAttr();
LayoutInfo layout;
- if (hasParamsOfLayoutKind(anchorLayout)) {
- layout = LayoutInfo(anchorLayout);
+ VectorType srcVecTy = llvm::cast<VectorType>(storeMatrix.getData().getType());
+ const uArch *uArch = getUArch(getChipStr(storeMatrix).value_or(""));
+ if (!uArch)
+ return;
+ if (hasParamsOfLayoutKind(anchorLayoutAttr)) {
+ requiredAnchorLayoutAttr = anchorLayoutAttr;
+ if (layoutKind == xegpu::LayoutKind::InstData) {
+ const auto uArchInstruction =
+ dyn_cast<xegpu::uArch::StoreScatterInstructionInterface>(
+ uArch->getInstruction(
+ xegpu::uArch::InstructionKind::StoreScatter));
+ if (!uArchInstruction)
+ return;
+ auto completed = xegpu::completeScatterStoreLaneLayoutFromInstData(
+ anchorLayoutAttr, srcVecTy.getElementType(), uArchInstruction,
+ uArch->getSubgroupSize());
+ if (!completed) {
+ storeMatrix.emitWarning(
+ "Failed to identify lane layouts for the specified inst_data.");
+ return;
+ }
+ requiredAnchorLayoutAttr = *completed;
+ storeMatrix.setLayoutAttr(requiredAnchorLayoutAttr);
+ }
} else {
- VectorType srcVecTy =
- llvm::cast<VectorType>(storeMatrix.getData().getType());
- const uArch *uArch = getUArch(getChipStr(storeMatrix).value_or(""));
- if (!uArch)
- return;
- auto requiredAnchorLayoutAttr =
- xegpu::setupStoreMatrixAnchorLayout(layoutKind, srcVecTy, uArch);
+ int chunkSize =
+ 1; // placeHolder for future use when StoreMatrix supports coalescing
+ requiredAnchorLayoutAttr = xegpu::setupStoreMatrixAnchorLayout(
+ layoutKind, srcVecTy, chunkSize, uArch);
storeMatrix.setLayoutAttr(requiredAnchorLayoutAttr);
- layout = LayoutInfo(requiredAnchorLayoutAttr);
}
-
+ layout = LayoutInfo(requiredAnchorLayoutAttr);
propagateIfChanged(operands[0], operands[0]->meet(layout));
}
@@ -1451,9 +1452,9 @@ namespace {
// ResolveLayoutConflicts
//===----------------------------------------------------------------------===//
-/// Helper to get the defining CreateNdDescOp of a tensor descriptor value. This
-/// function tries to find the defining CreateNdDescOp recursively accross
-/// control-flow boundaries.
+/// Helper to get the defining CreateNdDescOp of a tensor descriptor value.
+/// This function tries to find the defining CreateNdDescOp recursively
+/// accross control-flow boundaries.
static xegpu::CreateNdDescOp getDefiningCreateNdDescOp(Value tdescValue) {
// Try to get the defining CreateNdDescOp of the tensor descriptor.
auto definingOp = tdescValue.getDefiningOp<xegpu::CreateNdDescOp>();
@@ -1492,9 +1493,9 @@ LogicalResult ResolveLayoutConflicts::run() {
// Scan all operations in the parent op and resolve layout conflicts at
// tensor descriptor and vector use points.
auto r = parentOp->walk([&](Operation *op) -> WalkResult {
- // if the operation inputs vector and output scalar, like multi-reduction we
- // need to check if the result has layout and add a convert_layout to serve
- // as anchor op for the reduction op's layout.
+ // if the operation inputs vector and output scalar, like multi-reduction
+ // we need to check if the result has layout and add a convert_layout to
+ // serve as anchor op for the reduction op's layout.
if (isa<vector::MultiDimReductionOp>(op) || isa<vector::ReductionOp>(op)) {
for (OpResult result : op->getResults()) {
if (result.getType().isIntOrFloat()) {
@@ -1561,7 +1562,8 @@ ResolveLayoutConflicts::resolveVectorConsumer(OpOperand &operand) {
if (auto vectorTy = dyn_cast<VectorType>(vectorValue.getType());
vectorTy && vectorTy.getRank() > 1)
consumerOp->emitWarning("Expected layout for non-1D vectors.");
- return success(); // uniform non-tensor-data vector does not require layout
+ return success(); // uniform non-tensor-data vector does not require
+ // layout
}
// Region branch ops (e.g. scf.for) and their terminators (e.g. scf.yield)
// forward their operands to successor region inputs / parent op results;
@@ -1580,6 +1582,23 @@ ResolveLayoutConflicts::resolveVectorConsumer(OpOperand &operand) {
if (consumerLayout.isEqualTo(producerLayout))
return success();
+ // Consumer is a convert_layout: retarget its input_layout to the producer
+ // instead of chaining a second convert. Always safe (single source
+ // operand).
+ if (auto consumerConvert = dyn_cast<xegpu::ConvertLayoutOp>(consumerOp)) {
+ consumerConvert.setInputLayoutAttr(producerLayout);
+ return success();
+ }
+
+ // Producer is a convert_layout feeding only this use: retarget its
+ // target_layout to the consumer instead of appending another convert.
+ if (auto producerConvert =
+ vectorValue.getDefiningOp<xegpu::ConvertLayoutOp>();
+ producerConvert && vectorValue.hasOneUse()) {
+ producerConvert.setTargetLayoutAttr(consumerLayout);
+ return success();
+ }
+
// If the producer is trivially rematerializable (e.g. `vector.step`, splat
// `arith.constant`), clone it and stamp the consumer's expected layout on
// the clone instead of inserting a `xegpu.convert_layout`. The convert
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUUnroll.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUUnroll.cpp
index fc7c3b170dd3b..a9c73b3b84025 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUUnroll.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUUnroll.cpp
@@ -969,8 +969,6 @@ struct UnrollConvertLayoutOp : public UnrollPattern<xegpu::ConvertLayoutOp> {
if (valType.isIntOrFloat()) {
rewriter.replaceOp(op, op.getSource());
- assert(!inputLayout.dropInstData() && !targetLayout.dropInstData() &&
- "unexpected layout attributes for scalar type");
return success();
}
@@ -990,7 +988,7 @@ struct UnrollConvertLayoutOp : public UnrollPattern<xegpu::ConvertLayoutOp> {
Value newSource = op.getSource();
SmallVector<Value> newOps;
- if (inputLayout && targetLayout) {
+ if (inputLayout && targetLayout && !inputLayout.isEqualTo(targetLayout)) {
SmallVector<Type> convertedValTypes =
getUnrolledTypes(valueTy, *targetShape);
SmallVector<Value> convertedValues =
diff --git a/mlir/test/Dialect/XeGPU/propagate-layout-inst-data.mlir b/mlir/test/Dialect/XeGPU/propagate-layout-inst-data.mlir
index 5f493c8ca0df6..515c59db72819 100644
--- a/mlir/test/Dialect/XeGPU/propagate-layout-inst-data.mlir
+++ b/mlir/test/Dialect/XeGPU/propagate-layout-inst-data.mlir
@@ -4,13 +4,13 @@
// CHECK-LABEL: func.func @load_store_no_array_len(
// CHECK-SAME: %[[ARG0:[0-9a-zA-Z]+]]: memref<8x32xf32>, %[[ARG1:[0-9a-zA-Z]+]]: memref<8x32xf32>) {
// CHECK: %[[CST:.*]] = arith.constant dense<0.000000e+00> : vector<8x16xf32>
-// CHECK: %[[TDESC_SRC:.*]] = xegpu.create_nd_tdesc %[[ARG0]] : memref<8x32xf32> -> !xegpu.tensor_desc<8x32xf32, #xegpu.layout<inst_data = [8, 16]>>
-// CHECK: %[[TDESC_DST:.*]] = xegpu.create_nd_tdesc %[[ARG1]] : memref<8x32xf32> -> !xegpu.tensor_desc<8x32xf32, #xegpu.layout<inst_data = [8, 16]>>
-// CHECK: xegpu.prefetch_nd %[[TDESC_SRC]][0, 0] <{l1_hint = #xegpu.cache_hint<cached>, l2_hint = #xegpu.cache_hint<uncached>, layout = #xegpu.layout<inst_data = [8, 16]>}> :
-// CHECK-SAME: !xegpu.tensor_desc<8x32xf32, #xegpu.layout<inst_data = [8, 16]>>
-// CHECK: %[[LOADED:.*]] = xegpu.load_nd %0[0, 0] <{layout = #xegpu.layout<inst_data = [8, 16]>}>
-// CHECK-SAME: !xegpu.tensor_desc<8x32xf32, #xegpu.layout<inst_data = [8, 16]>> -> vector<8x32xf32>
-// CHECK: xegpu.store_nd %[[LOADED]], %[[TDESC_DST]][0, 0] <{layout = #xegpu.layout<inst_data = [8, 16]>}> : vector<8x32xf32>, !xegpu.tensor_desc<8x32xf32, #xegpu.layout<inst_data = [8, 16]>>
+// CHECK: %[[TDESC_SRC:.*]] = xegpu.create_nd_tdesc %[[ARG0]] : memref<8x32xf32> -> !xegpu.tensor_desc<8x32xf32, #xegpu.layout<inst_data = [8, 16], lane_layout = [1, 16], lane_data = [1, 1]>>
+// CHECK: %[[TDESC_DST:.*]] = xegpu.create_nd_tdesc %[[ARG1]] : memref<8x32xf32> -> !xegpu.tensor_desc<8x32xf32, #xegpu.layout<inst_data = [8, 16], lane_layout = [1, 16], lane_data = [1, 1]>>
+// CHECK: xegpu.prefetch_nd %[[TDESC_SRC]][0, 0] <{l1_hint = #xegpu.cache_hint<cached>, l2_hint = #xegpu.cache_hint<uncached>, layout = #xegpu.layout<inst_data = [8, 16], lane_layout = [1, 16], lane_data = [1, 1]>}> :
+// CHECK-SAME: !xegpu.tensor_desc<8x32xf32, #xegpu.layout<inst_data = [8, 16], lane_layout = [1, 16], lane_data = [1, 1]>>
+// CHECK: %[[LOADED:.*]] = xegpu.load_nd %0[0, 0] <{layout = #xegpu.layout<inst_data = [8, 16], lane_layout = [1, 16], lane_data = [1, 1]>}>
+// CHECK-SAME: !xegpu.tensor_desc<8x32xf32, #xegpu.layout<inst_data = [8, 16], lane_layout = [1, 16], lane_data = [1, 1]>> -> vector<8x32xf32>
+// CHECK: xegpu.store_nd %[[LOADED]], %[[TDESC_DST]][0, 0] <{layout = #xegpu.layout<inst_data = [8, 16], lane_layout = [1, 16], lane_data = [1, 1]>}> : vector<8x32xf32>, !xegpu.tensor_desc<8x32xf32, #xegpu.layout<inst_data = [8, 16], lane_layout = [1, 16], lane_data = [1, 1]>>
gpu.module @test {
// Although the uArch allows 8x32 inst data using block count (or array_len),
// it is up to optimization passes to decide on the block count usage.
@@ -29,17 +29,17 @@ func.func @load_store_no_array_len(%arg0: memref<8x32xf32>, %arg1: memref<8x32xf
// CHECK-LABEL: func.func @dpas_f16(
// CHECK-SAME: %[[ARG0:[0-9a-zA-Z]+]]: memref<8x16xf16>, %[[ARG1:[0-9a-zA-Z]+]]: memref<16x16xf16>, %[[ARG2:[0-9a-zA-Z]+]]: memref<8x16xf32>) {
-// CHECK: %[[CST:.*]] = arith.constant {layout_result_0 = #xegpu.layout<inst_data = [8, 16]>} dense<0.000000e+00> : vector<8x16xf32>
-// CHECK: %[[T0:.*]] = xegpu.create_nd_tdesc %[[ARG0]] : memref<8x16xf16> -> !xegpu.tensor_desc<8x16xf16, #xegpu.layout<inst_data = [8, 16]>
-// CHECK: %[[T1:.*]] = xegpu.create_nd_tdesc %[[ARG1]] : memref<16x16xf16> -> !xegpu.tensor_desc<16x16xf16, #xegpu.layout<inst_data = [16, 16]>>
-// CHECK: %[[T2:.*]] = xegpu.load_nd %[[T0]][0, 0] <{layout = #xegpu.layout<inst_data = [8, 16]>}> :
-// CHECK-SAME: !xegpu.tensor_desc<8x16xf16, #xegpu.layout<inst_data = [8, 16]>> -> vector<8x16xf16>
-// CHECK: %[[T3:.*]] = xegpu.load_nd %[[T1]][0, 0] <{layout = #xegpu.layout<inst_data = [16, 16]>}> :
-// CHECK-SAME: !xegpu.tensor_desc<16x16xf16, #xegpu.layout<inst_data = [16, 16]>> -> vector<16x16xf16>
-// CHECK: %[[T4:.*]] = xegpu.dpas %[[T2]], %[[T3]], %[[CST]] {layout_a = #xegpu.layout<inst_data = [8, 16]>, layout_b = #xegpu.layout<inst_data = [16, 16]>, layout_cd = #xegpu.layout<inst_data = [8, 16]>} :
+// CHECK: %[[CST:.*]] = arith.constant {layout_result_0 = #xegpu.layout<inst_data = [8, 16], lane_layout = [1, 16], lane_data = [1, 1]>} dense<0.000000e+00> : vector<8x16xf32>
+// CHECK: %[[T0:.*]] = xegpu.create_nd_tdesc %[[ARG0]] : memref<8x16xf16> -> !xegpu.tensor_desc<8x16xf16, #xegpu.layout<inst_data = [8, 16], lane_layout = [1, 16], lane_data = [1, 1]>
+// CHECK: %[[T1:.*]] = xegpu.create_nd_tdesc %[[ARG1]] : memref<16x16xf16> -> !xegpu.tensor_desc<16x16xf16, #xegpu.layout<inst_data = [16, 16], lane_layout = [1, 16], lane_data = [2, 1]>>
+// CHECK: %[[T2:.*]] = xegpu.load_nd %[[T0]][0, 0] <{layout = #xegpu.layout<inst_data = [8, 16], lane_layout = [1, 16], lane_data = [1, 1]>}> :
+// CHECK-SAME: !xegpu.tensor_desc<8x16xf16, #xegpu.layout<inst_data = [8, 16], lane_layout = [1, 16], lane_data = [1, 1]>> -> vector<8x16xf16>
+// CHECK: %[[T3:.*]] = xegpu.load_nd %[[T1]][0, 0] <{layout = #xegpu.layout<inst_data = [16, 16], lane_layout = [1, 16], lane_data = [2, 1]>}> :
+// CHECK-SAME: !xegpu.tensor_desc<16x16xf16, #xegpu.layout<inst_data = [16, 16], lane_layout = [1, 16], lane_data = [2, 1]>> -> vector<16x16xf16>
+// CHECK: %[[T4:.*]] = xegpu.dpas %[[T2]], %[[T3]], %[[CST]] {layout_a = #xegpu.layout<inst_data = [8, 16], lane_layout = [1, 16], lane_data = [1, 1]>, layout_b = #xegpu.layout<inst_data = [16, 16], lane_layout = [1, 16], lane_data = [2, 1]>, layout_cd = #xegpu.layout<inst_data = [8, 16], lane_layout = [1, 16], lane_data = [1, 1]>} :
// CHECK-SAME: vector<8x16xf16>, vector<16x16xf16>, vector<8x16xf32> -> vector<8x16xf32>
-// CHECK: %[[T5:.*]] = xegpu.create_nd_tdesc %[[ARG2]] : memref<8x16xf32> -> !xegpu.tensor_desc<8x16xf32, #xegpu.layout<inst_data = [8, 16]>
-// CHECK: xegpu.store_nd %[[T4]], %[[T5]][0, 0] <{layout = #xegpu.layout<inst_data = [8, 16]>}> : vector<8x16xf32>, !xegpu.tensor_desc<8x16xf32, #xegpu.layout<inst_data = [8, 16]>>
+// CHECK: %[[T5:.*]] = xegpu.create_nd_tdesc %[[ARG2]] : memref<8x16xf32> -> !xegpu.tensor_desc<8x16xf32, #xegpu.layout<inst_data = [8, 16], lane_layout = [1, 16], lane_data = [1, 1]>
+// CHECK: xegpu.store_nd %[[T4]], %[[T5]][0, 0] <{layout = #xegpu.layout<inst_data = [8, 16], lane_layout = [1, 16], lane_data = [1, 1]>}> : vector<8x16xf32>, !xegpu.tensor_desc<8x16xf32, #xegpu.layout<inst_data = [8, 16], lane_layout = [1, 16], lane_data = [1, 1]>>
gpu.module @test {
func.func @dpas_f16(%arg0: memref<8x16xf16>, %arg1: memref<16x16xf16>, %arg2: memref<8x16xf32>) {
%c0 = arith.constant 0 : index
@@ -70,15 +70,15 @@ gpu.module @test_kernel {
%c_tdesc = xegpu.create_nd_tdesc %C : memref<1024x1024xf16> -> !xegpu.tensor_desc<16x32xf16>
scf.for %k = %c0 to %c1024 step %c32 {
- //CHECK: xegpu.load_nd {{.*}} <{layout = #xegpu.layout<inst_data = [8, 16]>}> :
- //CHECK-SAME: !xegpu.tensor_desc<16x32xf16, #xegpu.layout<inst_data = [8, 16]>> -> vector<16x32xf16>
+ //CHECK: xegpu.load_nd {{.*}} <{layout = #xegpu.layout<inst_data = [8, 16], lane_layout = [1, 16], lane_data = [1, 1]>}> :
+ //CHECK-SAME: !xegpu.tensor_desc<16x32xf16, #xegpu.layout<inst_data = [8, 16], lane_layout = [1, 16], lane_data = [1, 1]>> -> vector<16x32xf16>
%a = xegpu.load_nd %a_tdesc[0, %k] : !xegpu.tensor_desc<16x32xf16> -> vector<16x32xf16>
%b = xegpu.load_nd %b_tdesc[0, %k] : !xegpu.tensor_desc<16x32xf16> -> vector<16x32xf16>
- //CHECK-COUNT: arith.addf {{.*}} {layout_result_0 = #xegpu.layout<inst_data = [8, 16]>} : vector<16x32xf16>
+ //CHECK-COUNT: arith.addf {{.*}} {layout_result_0 = #xegpu.layout<inst_data = [8, 16], lane_layout = [1, 16], lane_data = [1, 1]>} : vector<16x32xf16>
%c = arith.addf %a, %b : vector<16x32xf16>
- //CHECK-COUNT: xegpu.store_nd {{.*}} : vector<16x32xf16>, !xegpu.tensor_desc<16x32xf16, #xegpu.layout<inst_data = [8, 16]>>
+ //CHECK-COUNT: xegpu.store_nd {{.*}} : vector<16x32xf16>, !xegpu.tensor_desc<16x32xf16, #xegpu.layout<inst_data = [8, 16], lane_layout = [1, 16], lane_data = [1, 1]>>
xegpu.store_nd %c, %c_tdesc[0, %k] : vector<16x32xf16>, !xegpu.tensor_desc<16x32xf16>
}
gpu.return
@@ -100,46 +100,26 @@ gpu.module @test_kernel {
%c_tdesc = xegpu.create_nd_tdesc %C : memref<1024x1024xf16> -> !xegpu.tensor_desc<12x32xf16>
scf.for %k = %c0 to %c1024 step %c32 {
- //CHECK: xegpu.load_nd {{.*}} <{layout = #xegpu.layout<inst_data = [4, 16]>}> :
- //CHECK-SAME: !xegpu.tensor_desc<12x32xf16, #xegpu.layout<inst_data = [4, 16]>> -> vector<12x32xf16>
+ //CHECK: xegpu.load_nd {{.*}} <{layout = #xegpu.layout<inst_data = [4, 16], lane_layout = [1, 16], lane_data = [1, 1]>}> :
+ //CHECK-SAME: !xegpu.tensor_desc<12x32xf16, #xegpu.layout<inst_data = [4, 16], lane_layout = [1, 16], lane_data = [1, 1]>> -> vector<12x32xf16>
%a = xegpu.load_nd %a_tdesc[0, %k] : !xegpu.tensor_desc<12x32xf16> -> vector<12x32xf16>
%b = xegpu.load_nd %b_tdesc[0, %k] : !xegpu.tensor_desc<12x32xf16> -> vector<12x32xf16>
- //CHECK-COUNT: arith.addf {{.*}} {layout_result_0 = #xegpu.layout<inst_data = [4, 16]>} : vector<12x32xf16>
+ //CHECK-COUNT: arith.addf {{.*}} {layout_result_0 = #xegpu.layout<inst_data = [4, 16], lane_layout = [1, 16], lane_data = [1, 1]>} : vector<12x32xf16>
%c = arith.addf %a, %b : vector<12x32xf16>
- //CHECK-COUNT: xegpu.store_nd {{.*}} : vector<12x32xf16>, !xegpu.tensor_desc<12x32xf16, #xegpu.layout<inst_data = [4, 16]>>
+ //CHECK-COUNT: xegpu.store_nd {{.*}} : vector<12x32xf16>, !xegpu.tensor_desc<12x32xf16, #xegpu.layout<inst_data = [4, 16], lane_layout = [1, 16], lane_data = [1, 1]>>
xegpu.store_nd %c, %c_tdesc[0, %k] : vector<12x32xf16>, !xegpu.tensor_desc<12x32xf16>
}
gpu.return
}
}
-// -----
-gpu.module @test {
-// CHECK-LABEL: func.func @scatter_ops_chunksize(
-// CHECK-SAME: %[[ARG0:[0-9a-zA-Z]+]]: memref<256xf16>) {
-// CHECK: %{{.*}} = arith.constant {layout_result_0 = #xegpu.layout<inst_data = [16]>} dense<true> : vector<16xi1>
-// CHECK: %{{.*}} = arith.constant {layout_result_0 = #xegpu.layout<inst_data = [16]>} dense<12> : vector<16xindex>
-// CHECK: %{{.*}} = xegpu.load %[[ARG0]][%{{.*}}], %{{.*}} <{chunk_size = 8 : i64, layout = #xegpu.layout<inst_data = [16, 8]>}>
-// CHECK-SAME: memref<256xf16>, vector<16xindex>, vector<16xi1> -> vector<16x8xf16>
-// CHECK: xegpu.store %0, %[[ARG0]][%{{.*}}], %{{.*}} <{chunk_size = 8 : i64, layout = #xegpu.layout<inst_data = [16, 8]>}> : vector<16x8xf16>, memref<256xf16>, vector<16xindex>, vector<16xi1>
-func.func @scatter_ops_chunksize(%src: memref<256xf16>) {
- %1 = arith.constant dense<1>: vector<16xi1>
- %offset = arith.constant dense<12> : vector<16xindex>
- %3 = xegpu.load %src[%offset], %1 <{chunk_size=8}>
- : memref<256xf16>, vector<16xindex>, vector<16xi1> -> vector<16x8xf16>
- xegpu.store %3, %src[%offset], %1 <{chunk_size=8}>
- : vector<16x8xf16>, memref<256xf16>, vector<16xindex>, vector<16xi1>
- return
-}
-}
-
// -----
gpu.module @test {
// CHECK-LABEL: func.func @store_matrix(
// CHECK-SAME: %[[ARG0:[0-9a-zA-Z]+]]: !xegpu.mem_desc<16x64xf16>) {
-// CHECK: %{{.*}} = arith.constant {layout_result_0 = #xegpu.layout<inst_data = [1, 16]>} dense<0.000000e+00> : vector<16x16xf16>
+// CHECK: %{{.*}} = arith.constant {layout_result_0 = #xegpu.layout<inst_data = [1, 16], lane_layout = [1, 16], lane_data = [1, 1]>} dense<0.000000e+00> : vector<16x16xf16>
func.func @store_matrix(%arg0: !xegpu.mem_desc<16x64xf16>) {
%cst = arith.constant dense<0.0000> : vector<16x16xf16>
xegpu.store_matrix %cst, %arg0[8, 8]: vector<16x16xf16>, !xegpu.mem_desc<16x64xf16>
@@ -148,108 +128,49 @@ func.func @store_matrix(%arg0: !xegpu.mem_desc<16x64xf16>) {
}
}
-// -----
-gpu.module @test {
-// CHECK-LABEL: func.func @scatter_ops_chunksize_excessive(
-// CHECK-SAME: %[[ARG0:[0-9a-zA-Z]+]]: memref<1024xf32>) {
-// CHECK: %{{.*}} = arith.constant {layout_result_0 = #xegpu.layout<inst_data = [16]>} dense<true> : vector<16xi1>
-// CHECK: %{{.*}} = arith.constant {layout_result_0 = #xegpu.layout<inst_data = [16]>} dense<12> : vector<16xindex>
-// CHECK: %{{.*}} = xegpu.load %[[ARG0]][%{{.*}}], %{{.*}} <{chunk_size = 32 : i64, layout = #xegpu.layout<inst_data = [16, 16]>}> :
-// CHECK-SAME: memref<1024xf32>, vector<16xindex>, vector<16xi1> -> vector<16x32xf32>
-// CHECK: xegpu.store %0, %[[ARG0]][%{{.*}}], %{{.*}} <{chunk_size = 32 : i64, layout = #xegpu.layout<inst_data = [16, 16]>}> :
-// CHECK-SAME: vector<16x32xf32>, memref<1024xf32>, vector<16xindex>, vector<16xi1>
-func.func @scatter_ops_chunksize_excessive(%src: memref<1024xf32>) {
- %1 = arith.constant dense<1>: vector<16xi1>
- %offset = arith.constant dense<12> : vector<16xindex>
- %3 = xegpu.load %src[%offset], %1 <{chunk_size=32}>
- : memref<1024xf32>, vector<16xindex>, vector<16xi1> -> vector<16x32xf32>
- xegpu.store %3, %src[%offset], %1 <{chunk_size=32}>
- : vector<16x32xf32>, memref<1024xf32>, vector<16xindex>, vector<16xi1>
- return
-}
-}
-
-// -----
-
-gpu.module @test {
-// CHECK-LABEL: func.func @scatter_ops_chunksize_excessive_anchor(
-// CHECK-SAME: %[[ARG0:[0-9a-zA-Z]+]]: memref<1024xf32>) {
-// CHECK: %{{.*}} = arith.constant {layout_result_0 = #xegpu.layout<inst_data = [16]>} dense<true> : vector<16xi1>
-// CHECK: %{{.*}} = arith.constant {layout_result_0 = #xegpu.layout<inst_data = [16]>} dense<12> : vector<16xindex>
-// CHECK: %{{.*}} = xegpu.load %[[ARG0]][%{{.*}}], %{{.*}} <{chunk_size = 32 : i64, layout = #xegpu.layout<inst_data = [16, 16]>}> :
-// CHECK-SAME: memref<1024xf32>, vector<16xindex>, vector<16xi1> -> vector<16x32xf32>
-// CHECK: xegpu.store %0, %[[ARG0]][%{{.*}}], %{{.*}} <{chunk_size = 32 : i64, layout = #xegpu.layout<inst_data = [16, 16]>}> :
-// CHECK-SAME: vector<16x32xf32>, memref<1024xf32>, vector<16xindex>, vector<16xi1>
-func.func @scatter_ops_chunksize_excessive_anchor(%src: memref<1024xf32>) {
- %1 = arith.constant dense<1>: vector<16xi1>
- %offset = arith.constant dense<12> : vector<16xindex>
- %3 = xegpu.load %src[%offset], %1 <{chunk_size=32}>
- : memref<1024xf32>, vector<16xindex>, vector<16xi1> -> vector<16x32xf32>
- xegpu.store %3, %src[%offset], %1 <{chunk_size=32, layout = #xegpu.layout<inst_data = [16, 16]>}>
- : vector<16x32xf32>, memref<1024xf32>, vector<16xindex>, vector<16xi1>
- return
-}
-}
-
// -----
gpu.module @test {
-// CHECK-LABEL: func.func @scatter_ops_chunksize_slice(
-// CHECK-SAME: %[[ARG0:[0-9a-zA-Z]+]]: memref<1024xf32>) {
-// CHECK: %{{.*}} = arith.constant {layout_result_0 = #xegpu.layout<inst_data = [16]>} dense<true> : vector<16xi1>
-// CHECK: %{{.*}} = arith.constant {layout_result_0 = #xegpu.layout<inst_data = [16]>} dense<12> : vector<16xindex>
-// CHECK: %[[LOADED:.*]] = xegpu.load %[[ARG0]][%{{.*}}], %{{.*}} <{layout = #xegpu.layout<inst_data = [16]>}> :
-// CHECK-SAME: memref<1024xf32>, vector<16xindex>, vector<16xi1> -> vector<16xf32>
-// CHECK: %[[BCAST:.*]] = vector.broadcast %[[LOADED]] {layout_result_0 = #xegpu.layout<inst_data = [16, 16]>} : vector<16xf32> to vector<16x16xf32>
-// CHECK: xegpu.store %[[BCAST]], %[[ARG0]][%{{.*}}], %{{.*}} <{chunk_size = 16 : i64, layout = #xegpu.layout<inst_data = [16, 16]>}> :
-// CHECK-SAME: vector<16x16xf32>, memref<1024xf32>, vector<16xindex>, vector<16xi1>
-func.func @scatter_ops_chunksize_slice(%src: memref<1024xf32>) {
- %1 = arith.constant dense<1>: vector<16xi1>
- %offset = arith.constant dense<12> : vector<16xindex>
+// CHECK-LABEL: func.func @scatter_ops_coalesce_chunksize(
+// CHECK-SAME: %[[ARG0:[0-9a-zA-Z]+]]: memref<512xf32>) {
+// CHECK: %{{.*}} = arith.constant {layout_result_0 = #xegpu.layout<inst_data = [16, 16], lane_layout = [1, 16], lane_data = [1, 1]>} dense<true> : vector<16x32xi1>
+// CHECK: %{{.*}} = arith.constant {layout_result_0 = #xegpu.layout<inst_data = [16, 16], lane_layout = [1, 16], lane_data = [1, 1]>} dense<12> : vector<16x32xindex>
+// CHECK: %{{.*}} = xegpu.load %[[ARG0]][%{{.*}}], %{{.*}} <{layout = #xegpu.layout<inst_data = [16, 16], lane_layout = [1, 16], lane_data = [1, 1]>}> :
+// CHECK-SAME: memref<512xf32>, vector<16x32xindex>, vector<16x32xi1> -> vector<16x32xf32>
+// CHECK: xegpu.store %0, %[[ARG0]][%{{.*}}], %{{.*}} <{layout = #xegpu.layout<inst_data = [16, 16], lane_layout = [1, 16], lane_data = [1, 1]>}> :
+// CHECK-SAME: vector<16x32xf32>, memref<512xf32>, vector<16x32xindex>, vector<16x32xi1>
+func.func @scatter_ops_coalesce_chunksize(%src: memref<512xf32>) {
+ %1 = arith.constant dense<1>: vector<16x32xi1>
+ %offset = arith.constant dense<12> : vector<16x32xindex>
%3 = xegpu.load %src[%offset], %1
- : memref<1024xf32>, vector<16xindex>, vector<16xi1> -> vector<16xf32>
-
- %4 = vector.broadcast %3 : vector<16xf32> to vector<16x16xf32>
- xegpu.store %4, %src[%offset], %1 <{chunk_size=16, layout = #xegpu.layout<inst_data = [16, 16]>}>
- : vector<16x16xf32>, memref<1024xf32>, vector<16xindex>, vector<16xi1>
+ : memref<512xf32>, vector<16x32xindex>, vector<16x32xi1> -> vector<16x32xf32>
+ xegpu.store %3, %src[%offset], %1 <{layout = #xegpu.layout<inst_data = [16, 16]>}>
+ : vector<16x32xf32>, memref<512xf32>, vector<16x32xindex>, vector<16x32xi1>
return
}
}
// -----
-gpu.module @test {
-// CHECK-LABEL: func.func @insert_strided_slice_inst_data_no_packing(
-// CHECK-SAME: %[[ARG0:[0-9a-zA-Z]+]]: memref<8x32xf32>) {
-// CHECK: %[[CST_SMALL:.*]] = arith.constant {layout_result_0 = #xegpu.layout<inst_data = [4, 16]>} dense<1.000000e+00> : vector<4x16xf32>
-// CHECK: %[[CST_LARGE:.*]] = arith.constant {layout_result_0 = #xegpu.layout<inst_data = [4, 16]>} dense<0.000000e+00> : vector<8x32xf32>
-// CHECK: %[[INSERT:.*]] = vector.insert_strided_slice %[[CST_SMALL]], %[[CST_LARGE]] {layout_result_0 = #xegpu.layout<inst_data = [4, 16]>, offsets = [0, 0], strides = [1, 1]} : vector<4x16xf32> into vector<8x32xf32>
-// CHECK: %[[TDESC:.*]] = xegpu.create_nd_tdesc %[[ARG0]] : memref<8x32xf32> -> !xegpu.tensor_desc<8x32xf32, #xegpu.layout<inst_data = [8, 16]>>
-// CHECK: xegpu.store_nd %[[INSERT]], %[[TDESC]][0, 0] <{layout = #xegpu.layout<inst_data = [8, 16]>}> : vector<8x32xf32>, !xegpu.tensor_desc<8x32xf32, #xegpu.layout<inst_data = [8, 16]>>
-func.func @insert_strided_slice_inst_data_no_packing(%arg0: memref<8x32xf32>) {
- %c0 = arith.constant 0 : index
- %cst_small = arith.constant dense<1.0> : vector<4x16xf32>
- %cst_large = arith.constant dense<0.0> : vector<8x32xf32>
- %insert = vector.insert_strided_slice %cst_small, %cst_large {offsets = [0, 0], strides = [1, 1]} : vector<4x16xf32> into vector<8x32xf32>
- %tdesc = xegpu.create_nd_tdesc %arg0 : memref<8x32xf32> -> !xegpu.tensor_desc<8x32xf32>
- xegpu.store_nd %insert, %tdesc[0, 0] : vector<8x32xf32>, !xegpu.tensor_desc<8x32xf32>
- return
-}
-}
-// -----
gpu.module @test {
-// CHECK-LABEL: func.func @insert_strided_slice_inst_data_with_packing(
-// CHECK-SAME: %[[ARG0:[0-9a-zA-Z]+]]: memref<8x64xi8>) {
-// CHECK: %[[CST_SMALL:.*]] = arith.constant {layout_result_0 = #xegpu.layout<inst_data = [4, 64]>} dense<1> : vector<4x64xi8>
-// CHECK: %[[CST_LARGE:.*]] = arith.constant {layout_result_0 = #xegpu.layout<inst_data = [4, 64]>} dense<0> : vector<8x64xi8>
-// CHECK: %[[INSERT:.*]] = vector.insert_strided_slice %[[CST_SMALL]], %[[CST_LARGE]] {layout_result_0 = #xegpu.layout<inst_data = [4, 64]>, offsets = [0, 0], strides = [1, 1]} : vector<4x64xi8> into vector<8x64xi8>
-func.func @insert_strided_slice_inst_data_with_packing(%arg0: memref<8x64xi8>) {
+// CHECK-LABEL: func.func @load_gather_with_coalesce_chunksize(
+// CHECK-SAME: %[[ARG0:[0-9a-zA-Z]+]]: memref<8x16xf16>, %[[ARG1:[0-9a-zA-Z]+]]: memref<256xf16>, %[[ARG2:[0-9a-zA-Z]+]]: memref<8x16xf32>) {
+// CHECK: %[[OFFSET:.*]] = arith.constant {layout_result_0 = #xegpu.layout<inst_data = [16, 16], lane_layout = [16, 1], lane_data = [1, 2]>}
+// CHECK-SAME: dense<0> : vector<16x16xindex>
+// CHECK-NEXT: %[[MASK:.*]] = arith.constant {layout_result_0 = #xegpu.layout<inst_data = [16, 16], lane_layout = [16, 1], lane_data = [1, 2]>} dense<true> : vector<16x16xi1>
+// CHECK-NEXT: %{{.*}} = xegpu.load %arg1[%[[OFFSET]]], %[[MASK]] <{layout = #xegpu.layout<inst_data = [16, 16], lane_layout = [16, 1], lane_data = [1, 2]>}> : memref<256xf16>, vector<16x16xindex>, vector<16x16xi1> -> vector<16x16xf16>
+func.func @load_gather_with_coalesce_chunksize(%arg0: memref<8x16xf16>, %arg1: memref<256xf16>, %arg2: memref<8x16xf32>) {
%c0 = arith.constant 0 : index
- %cst_small = arith.constant dense<1> : vector<4x64xi8>
- %cst_large = arith.constant dense<0> : vector<8x64xi8>
- %insert = vector.insert_strided_slice %cst_small, %cst_large {offsets = [0, 0], strides = [1, 1]} : vector<4x64xi8> into vector<8x64xi8>
- %tdesc = xegpu.create_nd_tdesc %arg0 : memref<8x64xi8> -> !xegpu.tensor_desc<8x64xi8>
- xegpu.store_nd %insert, %tdesc[0, 0] <{layout = #xegpu.layout<inst_data = [8, 64]>}>: vector<8x64xi8>, !xegpu.tensor_desc<8x64xi8>
+ %0 = xegpu.create_nd_tdesc %arg0 : memref<8x16xf16> -> !xegpu.tensor_desc<8x16xf16>
+ %1 = xegpu.load_nd %0[0, 0] : !xegpu.tensor_desc<8x16xf16> -> vector<8x16xf16>
+ %offset = arith.constant dense<0> : vector<16x16xindex>
+ %mask = arith.constant dense<true> : vector<16x16xi1>
+ %3 = xegpu.load %arg1[%offset], %mask
+ : memref<256xf16>, vector<16x16xindex>, vector<16x16xi1> -> vector<16x16xf16>
+ %4 = vector.transpose %3, [1, 0] : vector<16x16xf16> to vector<16x16xf16>
+ %5 = xegpu.dpas %1, %4 : vector<8x16xf16>, vector<16x16xf16> -> vector<8x16xf32>
+ %6 = xegpu.create_nd_tdesc %arg2 : memref<8x16xf32> -> !xegpu.tensor_desc<8x16xf32>
+ xegpu.store_nd %5, %6[0, 0] : vector<8x16xf32>, !xegpu.tensor_desc<8x16xf32>
return
}
}
@@ -257,12 +178,12 @@ func.func @insert_strided_slice_inst_data_with_packing(%arg0: memref<8x64xi8>) {
// -----
gpu.module @test {
// CHECK-LABEL: func.func @vector_shape_cast_expand_non_unit_dims(
-// CHECK: %[[LOAD:.*]] = xegpu.load %arg0[%[[STEP:.*]]], %[[CST:.*]] <{layout = #xegpu.layout<inst_data = [16]>}> : memref<1024xf16>, vector<1024xindex>, vector<1024xi1> -> vector<1024xf16>
-// CHECK: %[[CAST:.*]] = vector.shape_cast %[[LOAD]] {layout_result_0 = #xegpu.layout<inst_data = [1, 1, 16]>} : vector<1024xf16> to vector<8x8x16xf16>
-// CHECK: %[[CST_0:.*]] = arith.constant {layout_result_0 = #xegpu.slice<#xegpu.layout<inst_data = [1, 1, 16]>, dims = [0]>} dense<0.000000e+00> : vector<8x16xf16>
-// CHECK: %[[CST_1:.*]] = arith.constant {layout_result_0 = #xegpu.slice<#xegpu.layout<inst_data = [1, 16]>, dims = [0]>} dense<0.000000e+00> : vector<16xf16>
-// CHECK: %[[REDUCE_0:.*]] = vector.multi_reduction <add>, %[[CAST]], %[[CST_0]] {layout_result_0 = #xegpu.slice<#xegpu.layout<inst_data = [1, 1, 16]>, dims = [0]>} [0] : vector<8x8x16xf16> to vector<8x16xf16>
-// CHECK: %[[REDUCE_1:.*]] = vector.multi_reduction <add>, %[[REDUCE_0]], %[[CST_1]] {layout_result_0 = #xegpu.slice<#xegpu.layout<inst_data = [1, 16]>, dims = [0]>} [0] : vector<8x16xf16> to vector<16xf16>
+// CHECK: %[[LOAD:.*]] = xegpu.load %arg0[%[[STEP:.*]]], %[[CST:.*]] <{layout = #xegpu.layout<inst_data = [16], lane_layout = [16], lane_data = [1]>}> : memref<1024xf16>, vector<1024xindex>, vector<1024xi1> -> vector<1024xf16>
+// CHECK: %[[CAST:.*]] = vector.shape_cast %[[LOAD]] {layout_result_0 = #xegpu.layout<inst_data = [1, 1, 16], lane_layout = [1, 1, 16], lane_data = [1, 1, 1]>} : vector<1024xf16> to vector<8x8x16xf16>
+// CHECK: %[[CST_0:.*]] = arith.constant {layout_result_0 = #xegpu.slice<#xegpu.layout<inst_data = [1, 1, 16], lane_layout = [1, 1, 16], lane_data = [1, 1, 1]>, dims = [0]>} dense<0.000000e+00> : vector<8x16xf16>
+// CHECK: %[[CST_1:.*]] = arith.constant {layout_result_0 = #xegpu.slice<#xegpu.layout<inst_data = [1, 16], lane_layout = [1, 16], lane_data = [1, 1]>, dims = [0]>} dense<0.000000e+00> : vector<16xf16>
+// CHECK: %[[REDUCE_0:.*]] = vector.multi_reduction <add>, %[[CAST]], %[[CST_0]] {layout_result_0 = #xegpu.slice<#xegpu.layout<inst_data = [1, 1, 16], lane_layout = [1, 1, 16], lane_data = [1, 1, 1]>, dims = [0]>} [0] : vector<8x8x16xf16> to vector<8x16xf16>
+// CHECK: %[[REDUCE_1:.*]] = vector.multi_reduction <add>, %[[REDUCE_0]], %[[CST_1]] {layout_result_0 = #xegpu.slice<#xegpu.layout<inst_data = [1, 16], lane_layout = [1, 16], lane_data = [1, 1]>, dims = [0]>} [0] : vector<8x16xf16> to vector<16xf16>
func.func @vector_shape_cast_expand_non_unit_dims(%arg0: memref<1024xf16>, %arg1: memref<16xf16>) {
%cst = arith.constant dense<true> : vector<1024xi1>
%0 = vector.step : vector<1024xindex>
@@ -282,7 +203,7 @@ func.func @vector_shape_cast_expand_non_unit_dims(%arg0: memref<1024xf16>, %arg1
// -----
gpu.module @test {
// CHECK-LABEL: func.func @vector_2d_reduction_with_fractional_subgroup_size(
-// CHECK: %[[ReduceVal:.*]] = vector.multi_reduction <add>, %[[Val:.*]], %[[CST:.*]] {layout_result_0 = #xegpu.slice<#xegpu.layout<inst_data = [1, 1, 1]>, dims = [1, 2]>} [1, 2] : vector<1x16x1xf16> to vector<1xf16>
+// CHECK: %[[ReduceVal:.*]] = vector.multi_reduction <add>, %[[Val:.*]], %[[CST:.*]] {layout_result_0 = #xegpu.slice<#xegpu.layout<inst_data = [1, 1, 1], lane_layout = [1, 1, 1], lane_data = [1, 1, 1]>, dims = [1, 2]>} [1, 2] : vector<1x16x1xf16> to vector<1xf16>
func.func @vector_2d_reduction_with_fractional_subgroup_size(%arg0: memref<1024xf16>, %arg1: memref<16xf16>) {
%cst = arith.constant dense<true> : vector<16xi1>
%0 = vector.step : vector<16xindex>
@@ -300,7 +221,7 @@ func.func @vector_2d_reduction_with_fractional_subgroup_size(%arg0: memref<1024x
// -----
gpu.module @test {
// CHECK-LABEL: func.func @vector_2d_reduction_with_fractional_subgroup_size_1x4x1(
-// CHECK: %[[ReduceVal:.*]] = vector.multi_reduction <add>, %[[Val:.*]], %[[CST:.*]] {layout_result_0 = #xegpu.slice<#xegpu.layout<inst_data = [1, 1, 4]>, dims = [1, 2]>} [1, 2] : vector<1x16x4xf16> to vector<1xf16>
+// CHECK: %[[ReduceVal:.*]] = vector.multi_reduction <add>, %[[Val:.*]], %[[CST:.*]] {layout_result_0 = #xegpu.slice<#xegpu.layout<inst_data = [1, 1, 4], lane_layout = [1, 1, 4], lane_data = [1, 1, 1]>, dims = [1, 2]>} [1, 2] : vector<1x16x4xf16> to vector<1xf16>
func.func @vector_2d_reduction_with_fractional_subgroup_size_1x4x1(%arg0: memref<1024xf16>, %arg1: memref<16xf16>) {
%cst = arith.constant dense<true> : vector<64xi1>
%0 = vector.step : vector<64xindex>
@@ -318,13 +239,13 @@ func.func @vector_2d_reduction_with_fractional_subgroup_size_1x4x1(%arg0: memref
// -----
gpu.module @test {
// CHECK-LABEL: func.func @vector_shape_cast_expand_and_merge(
-// CHECK: %[[CST:.*]] = arith.constant {layout_result_0 = #xegpu.layout<inst_data = [32]>} dense<true> : vector<256xi1>
-// CHECK: %[[STEP:.*]] = vector.step {layout_result_0 = #xegpu.layout<inst_data = [32]>} : vector<256xindex>
-// CHECK: %[[LOAD:.*]] = xegpu.load %arg0[%[[STEP]]], %[[CST]] <{layout = #xegpu.layout<inst_data = [32]>}> : memref<256xf16>, vector<256xindex>, vector<256xi1> -> vector<256xf16>
-// CHECK: %[[CAST_0:.*]] = vector.shape_cast %[[LOAD]] {layout_result_0 = #xegpu.layout<inst_data = [1, 1, 32]>} : vector<256xf16> to vector<2x4x32xf16>
-// CHECK: %[[CAST_1:.*]] = vector.shape_cast %[[CAST_0]] {layout_result_0 = #xegpu.layout<inst_data = [1, 32]>} : vector<2x4x32xf16> to vector<1x256xf16>
-// CHECK: %[[CAST_2:.*]] = vector.shape_cast %[[CAST_1]] {layout_result_0 = #xegpu.layout<inst_data = [32]>} : vector<1x256xf16> to vector<256xf16>
-// CHECK: xegpu.store %[[CAST_2]], %arg1[%[[STEP]]], %[[CST]] <{layout = #xegpu.layout<inst_data = [32]>}> : vector<256xf16>, memref<256xf16>, vector<256xindex>, vector<256xi1>
+// CHECK: %[[CST:.*]] = arith.constant {layout_result_0 = #xegpu.layout<inst_data = [32], lane_layout = [16], lane_data = [2]>} dense<true> : vector<256xi1>
+// CHECK: %[[STEP:.*]] = vector.step {layout_result_0 = #xegpu.layout<inst_data = [32], lane_layout = [16], lane_data = [2]>} : vector<256xindex>
+// CHECK: %[[LOAD:.*]] = xegpu.load %arg0[%[[STEP]]], %[[CST]] <{layout = #xegpu.layout<inst_data = [32], lane_layout = [16], lane_data = [2]>}> : memref<256xf16>, vector<256xindex>, vector<256xi1> -> vector<256xf16>
+// CHECK: %[[CAST_0:.*]] = vector.shape_cast %[[LOAD]] {layout_result_0 = #xegpu.layout<inst_data = [2, 4, 4], lane_layout = [2, 4, 2], lane_data = [1, 1, 2]>} : vector<256xf16> to vector<2x4x32xf16>
+// CHECK: %[[CAST_1:.*]] = vector.shape_cast %[[CAST_0]] {layout_result_0 = #xegpu.layout<inst_data = [1, 32], lane_layout = [1, 16], lane_data = [1, 2]>} : vector<2x4x32xf16> to vector<1x256xf16>
+// CHECK: %[[CAST_2:.*]] = vector.shape_cast %[[CAST_1]] {layout_result_0 = #xegpu.layout<inst_data = [32], lane_layout = [16], lane_data = [2]>} : vector<1x256xf16> to vector<256xf16>
+// CHECK: xegpu.store %[[CAST_2]], %arg1[%[[STEP]]], %[[CST]] <{layout = #xegpu.layout<inst_data = [32], lane_layout = [16], lane_data = [2]>}> : vector<256xf16>, memref<256xf16>, vector<256xindex>, vector<256xi1>
func.func @vector_shape_cast_expand_and_merge(%arg0: memref<256xf16>, %arg1: memref<256xf16>) {
%cst = arith.constant dense<true> : vector<256xi1>
%0 = vector.step : vector<256xindex>
@@ -341,8 +262,8 @@ func.func @vector_shape_cast_expand_and_merge(%arg0: memref<256xf16>, %arg1: mem
// -----
gpu.module @test{
// CHECK-LABEL: load_store_matrix
- // CHECK: xegpu.load_matrix %{{.*}} <{layout = #xegpu.layout<inst_data = [1, 1]>}>
- // CHECK: xegpu.store_matrix %{{.*}} <{layout = #xegpu.layout<inst_data = [1, 1]>}>
+ // CHECK: xegpu.load_matrix %{{.*}} <{layout = #xegpu.layout<inst_data = [1, 1], lane_layout = [1, 1], lane_data = [1, 1]>}>
+ // CHECK: xegpu.store_matrix %{{.*}} <{layout = #xegpu.layout<inst_data = [1, 1], lane_layout = [1, 1], lane_data = [1, 1]>}>
func.func @load_store_matrix(%arg0: !xegpu.mem_desc<64x128xf32>, %arg1: i1) {
%c0 = arith.constant 0 : index
scf.if %arg1 {
@@ -356,11 +277,11 @@ gpu.module @test{
// -----
gpu.module @test{
// CHECK-LABEL: broadcast_both_leadingdims_innerdims
- // CHECK: arith.constant {layout_result_0 = #xegpu.layout<inst_data = [1, 1, 1, 16]>} dense<true> : vector<2x2x6x32xi1>
- // CHECK: arith.constant {layout_result_0 = #xegpu.layout<inst_data = [1, 1, 1, 16]>} dense<1.000000e+00> : vector<2x2x6x32xf32>
- // CHECK: vector.step {layout_result_0 = #xegpu.slice<#xegpu.slice<#xegpu.layout<inst_data = [1, 1, 1, 1]>, dims = [0, 1]>, dims = [1]>} : vector<6xindex>
- // CHECK: vector.shape_cast {{.*}} {layout_result_0 = #xegpu.slice<#xegpu.layout<inst_data = [1, 1, 1, 1]>, dims = [0, 1]>} : vector<6xindex> to vector<6x1xindex>
- // CHECK: vector.broadcast {{.*}} {layout_result_0 = #xegpu.layout<inst_data = [1, 1, 1, 16]>} : vector<6x1xindex> to vector<2x2x6x32xindex>
+ // CHECK: arith.constant {layout_result_0 = #xegpu.layout<inst_data = [1, 1, 1, 16], lane_layout = [1, 1, 1, 16], lane_data = [1, 1, 1, 1]>} dense<true> : vector<2x2x6x32xi1>
+ // CHECK: arith.constant {layout_result_0 = #xegpu.layout<inst_data = [1, 1, 1, 16], lane_layout = [1, 1, 1, 16], lane_data = [1, 1, 1, 1]>} dense<1.000000e+00> : vector<2x2x6x32xf32>
+ // CHECK: vector.step {layout_result_0 = #xegpu.slice<#xegpu.slice<#xegpu.layout<inst_data = [1, 1, 1, 1], lane_layout = [1, 1, 1, 16], lane_data = [1, 1, 1, 1]>, dims = [0, 1]>, dims = [1]>} : vector<6xindex>
+ // CHECK: vector.shape_cast {{.*}} {layout_result_0 = #xegpu.slice<#xegpu.layout<inst_data = [1, 1, 1, 1], lane_layout = [1, 1, 1, 16], lane_data = [1, 1, 1, 1]>, dims = [0, 1]>} : vector<6xindex> to vector<6x1xindex>
+ // CHECK: vector.broadcast {{.*}} {layout_result_0 = #xegpu.layout<inst_data = [1, 1, 1, 16], lane_layout = [1, 1, 1, 16], lane_data = [1, 1, 1, 1]>} : vector<6x1xindex> to vector<2x2x6x32xindex>
gpu.func @broadcast_both_leadingdims_innerdims(%arg0: memref<32x2x192xf32>, %arg1: memref<32x2x192xf32>, %arg2: memref<32x2x192xf32>) kernel attributes {known_block_size = array<i32: 768, 1, 1>, known_grid_size = array<i32: 16, 1, 1>} {
%cst = arith.constant dense<true> : vector<2x2x6x32xi1>
%cst_0 = arith.constant dense<1.000000e+00> : vector<2x2x6x32xf32>
@@ -391,7 +312,7 @@ gpu.module @test_collapse_dims [#xevm.target<O = 3, chip = "pvc">] {
%mask = arith.constant dense<true> : vector<32x32xi1>
%data = arith.constant dense<0.0> : vector<32x32xf32>
- // CHECK: xegpu.store {{.*}} <{{{.*}}layout = #xegpu.layout<inst_data = [32, 32]>{{.*}}}> :
+ // CHECK: xegpu.store {{.*}} <{{{.*}}layout = #xegpu.layout<inst_data = [32, 32], lane_layout = [1, 16], lane_data = [1, 2]>{{.*}}}> :
xegpu.store %data, %ptr_i64[%1], %mask {
layout = #xegpu.layout<inst_data = [32, 32]>
} : vector<32x32xf32>, i64, vector<32x32xindex>, vector<32x32xi1>
@@ -403,22 +324,22 @@ gpu.module @test_collapse_dims [#xevm.target<O = 3, chip = "pvc">] {
// -----
gpu.module @test {
// CHECK-LABEL: func.func @bitcast_ui8_to_f4(
-// CHECK-SAME: %[[ARG0:[0-9a-zA-Z]+]]: memref<256x16xui8>) {
-// CHECK: %[[TDESC:.*]] = xegpu.create_nd_tdesc %[[ARG0]] : memref<256x16xui8> -> !xegpu.tensor_desc<256x16xui8, #xegpu.layout<inst_data = [32, 16]>>
-// CHECK: %[[LOAD:.*]] = xegpu.load_nd %[[TDESC]][0, 0] <{layout = #xegpu.layout<inst_data = [32, 16]>}> :
-// CHECK-SAME: !xegpu.tensor_desc<256x16xui8, #xegpu.layout<inst_data = [32, 16]>> -> vector<256x16xui8>
-// CHECK: %[[BC:.*]] = vector.bitcast %[[LOAD]] {layout_result_0 = #xegpu.layout<inst_data = [32, 32]>} : vector<256x16xui8> to vector<256x32xf4E2M1FN>
+// CHECK-SAME: %[[ARG0:[0-9a-zA-Z]+]]: memref<256x32xui8>) {
+// CHECK: %[[TDESC:.*]] = xegpu.create_nd_tdesc %[[ARG0]] : memref<256x32xui8> -> !xegpu.tensor_desc<256x32xui8, #xegpu.layout<inst_data = [32, 16], lane_layout = [1, 16], lane_data = [1, 1]>>
+// CHECK: %[[LOAD:.*]] = xegpu.load_nd %[[TDESC]][0, 0] <{layout = #xegpu.layout<inst_data = [32, 16], lane_layout = [1, 16], lane_data = [1, 1]>}> :
+// CHECK-SAME: !xegpu.tensor_desc<256x32xui8, #xegpu.layout<inst_data = [32, 16], lane_layout = [1, 16], lane_data = [1, 1]>> -> vector<256x32xui8>
+// CHECK: %[[BC:.*]] = vector.bitcast %[[LOAD]] {layout_result_0 = #xegpu.layout<inst_data = [32, 32], lane_layout = [1, 16], lane_data = [1, 2]>} : vector<256x32xui8> to vector<256x64xf4E2M1FN>
// CHECK: xegpu.convert_layout %[[BC]]
-// CHECK-SAME: <{input_layout = #xegpu.layout<inst_data = [32, 32]>, target_layout = #xegpu.layout<inst_data = [32, 32]>}>
-// CHECK-SAME: : vector<256x32xf4E2M1FN>
-func.func @bitcast_ui8_to_f4(%arg0: memref<256x16xui8>) {
- %0 = xegpu.create_nd_tdesc %arg0 : memref<256x16xui8> -> !xegpu.tensor_desc<256x16xui8>
- %1 = xegpu.load_nd %0[0, 0] : !xegpu.tensor_desc<256x16xui8> -> vector<256x16xui8>
- %2 = vector.bitcast %1 : vector<256x16xui8> to vector<256x32xf4E2M1FN>
+// CHECK-SAME: <{input_layout = #xegpu.layout<inst_data = [32, 32], lane_layout = [1, 16], lane_data = [1, 2]>, target_layout = #xegpu.layout<inst_data = [32, 32], lane_layout = [1, 16], lane_data = [1, 2]>}>
+// CHECK-SAME: : vector<256x64xf4E2M1FN>
+func.func @bitcast_ui8_to_f4(%arg0: memref<256x32xui8>) {
+ %0 = xegpu.create_nd_tdesc %arg0 : memref<256x32xui8> -> !xegpu.tensor_desc<256x32xui8>
+ %1 = xegpu.load_nd %0[0, 0] : !xegpu.tensor_desc<256x32xui8> -> vector<256x32xui8>
+ %2 = vector.bitcast %1 : vector<256x32xui8> to vector<256x64xf4E2M1FN>
%3 = xegpu.convert_layout %2
- <{input_layout = #xegpu.layout<inst_data = [32, 32]>,
- target_layout = #xegpu.layout<inst_data = [32, 32]>}>
- : vector<256x32xf4E2M1FN>
+ <{input_layout = #xegpu.layout<inst_data = [32, 32], lane_layout = [1, 16], lane_data = [1, 2]>,
+ target_layout = #xegpu.layout<inst_data = [32, 32], lane_layout = [1, 16], lane_data = [1, 2]>}>
+ : vector<256x64xf4E2M1FN>
return
}
}
@@ -427,20 +348,20 @@ func.func @bitcast_ui8_to_f4(%arg0: memref<256x16xui8>) {
gpu.module @test {
// CHECK-LABEL: func.func @bitcast_ui16_to_f4(
// CHECK-SAME: %[[ARG0:[0-9a-zA-Z]+]]: memref<256x16xui16>) {
-// CHECK: %[[TDESC:.*]] = xegpu.create_nd_tdesc %[[ARG0]] : memref<256x16xui16> -> !xegpu.tensor_desc<256x16xui16, #xegpu.layout<inst_data = [32, 16]>>
-// CHECK: %[[LOAD:.*]] = xegpu.load_nd %[[TDESC]][0, 0] <{layout = #xegpu.layout<inst_data = [32, 16]>}> :
-// CHECK-SAME: !xegpu.tensor_desc<256x16xui16, #xegpu.layout<inst_data = [32, 16]>> -> vector<256x16xui16>
-// CHECK: %[[BC:.*]] = vector.bitcast %[[LOAD]] {layout_result_0 = #xegpu.layout<inst_data = [32, 64]>} : vector<256x16xui16> to vector<256x64xf4E2M1FN>
+// CHECK: %[[TDESC:.*]] = xegpu.create_nd_tdesc %[[ARG0]] : memref<256x16xui16> -> !xegpu.tensor_desc<256x16xui16, #xegpu.layout<inst_data = [32, 16], lane_layout = [1, 16], lane_data = [1, 1]>>
+// CHECK: %[[LOAD:.*]] = xegpu.load_nd %[[TDESC]][0, 0] <{layout = #xegpu.layout<inst_data = [32, 16], lane_layout = [1, 16], lane_data = [1, 1]>}> :
+// CHECK-SAME: !xegpu.tensor_desc<256x16xui16, #xegpu.layout<inst_data = [32, 16], lane_layout = [1, 16], lane_data = [1, 1]>> -> vector<256x16xui16>
+// CHECK: %[[BC:.*]] = vector.bitcast %[[LOAD]] {layout_result_0 = #xegpu.layout<inst_data = [32, 64], lane_layout = [1, 16], lane_data = [1, 4]>} : vector<256x16xui16> to vector<256x64xf4E2M1FN>
// CHECK: xegpu.convert_layout %[[BC]]
-// CHECK-SAME: <{input_layout = #xegpu.layout<inst_data = [32, 32]>, target_layout = #xegpu.layout<inst_data = [32, 32]>}>
+// CHECK-SAME: <{input_layout = #xegpu.layout<inst_data = [32, 64], lane_layout = [1, 16], lane_data = [1, 4]>, target_layout = #xegpu.layout<inst_data = [32, 64], lane_layout = [1, 16], lane_data = [1, 4]>}>
// CHECK-SAME: : vector<256x64xf4E2M1FN>
func.func @bitcast_ui16_to_f4(%arg0: memref<256x16xui16>) {
%0 = xegpu.create_nd_tdesc %arg0 : memref<256x16xui16> -> !xegpu.tensor_desc<256x16xui16>
%1 = xegpu.load_nd %0[0, 0] : !xegpu.tensor_desc<256x16xui16> -> vector<256x16xui16>
%2 = vector.bitcast %1 : vector<256x16xui16> to vector<256x64xf4E2M1FN>
%3 = xegpu.convert_layout %2
- <{input_layout = #xegpu.layout<inst_data = [32, 32]>,
- target_layout = #xegpu.layout<inst_data = [32, 32]>}>
+ <{input_layout = #xegpu.layout<inst_data = [32, 64], lane_layout = [1, 16], lane_data = [1, 4]>,
+ target_layout = #xegpu.layout<inst_data = [32, 64], lane_layout = [1, 16], lane_data = [1, 4]>}>
: vector<256x64xf4E2M1FN>
return
}
@@ -449,40 +370,40 @@ func.func @bitcast_ui16_to_f4(%arg0: memref<256x16xui16>) {
// -----
// CHECK-LABEL: func.func @dpas_mx_f8e5m2
-// CHECK-SAME: %[[ARG0:[0-9a-zA-Z]+]]: memref<16x64xf8E5M2>, %[[ARG1:[0-9a-zA-Z]+]]: memref<64x32xf8E5M2>, %[[ARG2:[0-9a-zA-Z]+]]: memref<16x32xbf16>
-// CHECK-SAME: %[[ARG3:[0-9a-zA-Z]+]]: memref<16x2xf8E8M0FNU>, %[[ARG4:[0-9a-zA-Z]+]]: memref<2x32xf8E8M0FNU>
-// CHECK: %[[CST:.*]] = arith.constant {layout_result_0 = #xegpu.layout<inst_data = [8, 16]>} dense<0.000000e+00> : vector<16x32xbf16>
-// CHECK: %[[T0:.*]] = xegpu.create_nd_tdesc %[[ARG0]] : memref<16x64xf8E5M2> -> !xegpu.tensor_desc<16x64xf8E5M2, #xegpu.layout<inst_data = [8, 32]>>
-// CHECK: %[[T1:.*]] = xegpu.create_nd_tdesc %[[ARG1]] : memref<64x32xf8E5M2> -> !xegpu.tensor_desc<64x32xf8E5M2, #xegpu.layout<inst_data = [32, 16]>>
-// CHECK: %[[T2:.*]] = xegpu.load_nd %[[T0]][0, 0] <{layout = #xegpu.layout<inst_data = [8, 32]>}> :
-// CHECK-SAME: !xegpu.tensor_desc<16x64xf8E5M2, #xegpu.layout<inst_data = [8, 32]>> -> vector<16x64xf8E5M2>
-// CHECK: %[[T3:.*]] = xegpu.load_nd %[[T1]][0, 0] <{layout = #xegpu.layout<inst_data = [32, 16]>}> :
-// CHECK-SAME: !xegpu.tensor_desc<64x32xf8E5M2, #xegpu.layout<inst_data = [32, 16]>> -> vector<64x32xf8E5M2>
-// CHECK: %[[T4:.*]] = xegpu.create_nd_tdesc %[[ARG3]] : memref<16x2xf8E8M0FNU> -> !xegpu.tensor_desc<16x2xf8E8M0FNU, #xegpu.layout<inst_data = [8, 1]>>
-// CHECK: %[[T5:.*]] = xegpu.load_nd %[[T4]][0, 0] <{layout = #xegpu.layout<inst_data = [8, 1]>}> :
-// CHECK-SAME: !xegpu.tensor_desc<16x2xf8E8M0FNU, #xegpu.layout<inst_data = [8, 1]>> -> vector<16x2xf8E8M0FNU>
-// CHECK: %[[T6:.*]] = xegpu.create_nd_tdesc %[[ARG4]] : memref<2x32xf8E8M0FNU> -> !xegpu.tensor_desc<2x32xf8E8M0FNU, #xegpu.layout<inst_data = [1, 16]>>
-// CHECK: %[[T7:.*]] = xegpu.load_nd %[[T6]][0, 0] <{layout = #xegpu.layout<inst_data = [1, 16]>}> :
-// CHECK-SAME: !xegpu.tensor_desc<2x32xf8E8M0FNU, #xegpu.layout<inst_data = [1, 16]>> -> vector<2x32xf8E8M0FNU>
+// CHECK-SAME: %[[ARG0:[0-9a-zA-Z]+]]: memref<16x1024xf8E5M2>, %[[ARG1:[0-9a-zA-Z]+]]: memref<1024x32xf8E5M2>, %[[ARG2:[0-9a-zA-Z]+]]: memref<16x32xbf16>
+// CHECK-SAME: %[[ARG3:[0-9a-zA-Z]+]]: memref<16x32xf8E8M0FNU>, %[[ARG4:[0-9a-zA-Z]+]]: memref<32x32xf8E8M0FNU>
+// CHECK: %[[CST:.*]] = arith.constant {layout_result_0 = #xegpu.layout<inst_data = [8, 16], lane_layout = [1, 16], lane_data = [1, 1]>} dense<0.000000e+00> : vector<16x32xbf16>
+// CHECK: %[[T0:.*]] = xegpu.create_nd_tdesc %[[ARG0]] : memref<16x1024xf8E5M2> -> !xegpu.tensor_desc<16x1024xf8E5M2, #xegpu.layout<inst_data = [8, 32], lane_layout = [1, 16], lane_data = [1, 2]>>
+// CHECK: %[[T1:.*]] = xegpu.create_nd_tdesc %[[ARG1]] : memref<1024x32xf8E5M2> -> !xegpu.tensor_desc<1024x32xf8E5M2, #xegpu.layout<inst_data = [32, 16], lane_layout = [1, 16], lane_data = [4, 1]>>
+// CHECK: %[[T2:.*]] = xegpu.load_nd %[[T0]][0, 0] <{layout = #xegpu.layout<inst_data = [8, 32], lane_layout = [1, 16], lane_data = [1, 2]>}> :
+// CHECK-SAME: !xegpu.tensor_desc<16x1024xf8E5M2, #xegpu.layout<inst_data = [8, 32], lane_layout = [1, 16], lane_data = [1, 2]>> -> vector<16x1024xf8E5M2>
+// CHECK: %[[T3:.*]] = xegpu.load_nd %[[T1]][0, 0] <{layout = #xegpu.layout<inst_data = [32, 16], lane_layout = [1, 16], lane_data = [4, 1]>}> :
+// CHECK-SAME: !xegpu.tensor_desc<1024x32xf8E5M2, #xegpu.layout<inst_data = [32, 16], lane_layout = [1, 16], lane_data = [4, 1]>> -> vector<1024x32xf8E5M2>
+// CHECK: %[[T4:.*]] = xegpu.create_nd_tdesc %[[ARG3]] : memref<16x32xf8E8M0FNU> -> !xegpu.tensor_desc<16x32xf8E8M0FNU, #xegpu.layout<inst_data = [16, 32], lane_layout = [16, 1], lane_data = [1, 1]>>
+// CHECK: %[[T5:.*]] = xegpu.load_nd %[[T4]][0, 0] <{layout = #xegpu.layout<inst_data = [16, 32], lane_layout = [16, 1], lane_data = [1, 1]>}> :
+// CHECK-SAME: !xegpu.tensor_desc<16x32xf8E8M0FNU, #xegpu.layout<inst_data = [16, 32], lane_layout = [16, 1], lane_data = [1, 1]>> -> vector<16x32xf8E8M0FNU>
+// CHECK: %[[T6:.*]] = xegpu.create_nd_tdesc %[[ARG4]] : memref<32x32xf8E8M0FNU> -> !xegpu.tensor_desc<32x32xf8E8M0FNU, #xegpu.layout<inst_data = [1, 16], lane_layout = [1, 16], lane_data = [1, 1]>>
+// CHECK: %[[T7:.*]] = xegpu.load_nd %[[T6]][0, 0] <{layout = #xegpu.layout<inst_data = [1, 16], lane_layout = [1, 16], lane_data = [1, 1]>}> :
+// CHECK-SAME: !xegpu.tensor_desc<32x32xf8E8M0FNU, #xegpu.layout<inst_data = [1, 16], lane_layout = [1, 16], lane_data = [1, 1]>> -> vector<32x32xf8E8M0FNU>
// CHECK: %[[T8:.*]] = xegpu.dpas_mx %[[T2]], %[[T3]], %[[CST]] scale_a = %[[T5]] scale_b = %[[T7]]
-// CHECK-SAME: {layout_a = #xegpu.layout<inst_data = [8, 32]>, layout_a_scale = #xegpu.layout<inst_data = [8, 1]>, layout_b = #xegpu.layout<inst_data = [32, 16]>, layout_b_scale = #xegpu.layout<inst_data = [1, 16]>, layout_cd = #xegpu.layout<inst_data = [8, 16]>} :
-// CHECK-SAME: (vector<16x64xf8E5M2>, vector<64x32xf8E5M2>, vector<16x32xbf16>, vector<16x2xf8E8M0FNU>, vector<2x32xf8E8M0FNU>) -> vector<16x32xbf16>
-// CHECK: %[[T9:.*]] = xegpu.create_nd_tdesc %[[ARG2]] : memref<16x32xbf16> -> !xegpu.tensor_desc<16x32xbf16, #xegpu.layout<inst_data = [8, 16]>>
-// CHECK: xegpu.store_nd %[[T8]], %[[T9]][0, 0] <{layout = #xegpu.layout<inst_data = [8, 16]>}> : vector<16x32xbf16>, !xegpu.tensor_desc<16x32xbf16, #xegpu.layout<inst_data = [8, 16]>>
+// CHECK-SAME: {layout_a = #xegpu.layout<inst_data = [8, 32], lane_layout = [1, 16], lane_data = [1, 2]>, layout_a_scale = #xegpu.layout<inst_data = [8, 1], lane_layout = [8, 1], lane_data = [1, 1]>, layout_b = #xegpu.layout<inst_data = [32, 16], lane_layout = [1, 16], lane_data = [4, 1]>, layout_b_scale = #xegpu.layout<inst_data = [1, 16], lane_layout = [1, 16], lane_data = [1, 1]>, layout_cd = #xegpu.layout<inst_data = [8, 16], lane_layout = [1, 16], lane_data = [1, 1]>} :
+// CHECK-SAME: (vector<16x1024xf8E5M2>, vector<1024x32xf8E5M2>, vector<16x32xbf16>, vector<16x32xf8E8M0FNU>, vector<32x32xf8E8M0FNU>) -> vector<16x32xbf16>
+// CHECK: %[[T9:.*]] = xegpu.create_nd_tdesc %[[ARG2]] : memref<16x32xbf16> -> !xegpu.tensor_desc<16x32xbf16, #xegpu.layout<inst_data = [8, 16], lane_layout = [1, 16], lane_data = [1, 1]>>
+// CHECK: xegpu.store_nd %[[T8]], %[[T9]][0, 0] <{layout = #xegpu.layout<inst_data = [8, 16], lane_layout = [1, 16], lane_data = [1, 1]>}> : vector<16x32xbf16>, !xegpu.tensor_desc<16x32xbf16, #xegpu.layout<inst_data = [8, 16], lane_layout = [1, 16], lane_data = [1, 1]>>
gpu.module @test {
-func.func @dpas_mx_f8e5m2(%arg0: memref<16x64xf8E5M2>, %arg1: memref<64x32xf8E5M2>, %arg2: memref<16x32xbf16>,
- %arg3: memref<16x2xf8E8M0FNU>, %arg4: memref<2x32xf8E8M0FNU>) {
+func.func @dpas_mx_f8e5m2(%arg0: memref<16x1024xf8E5M2>, %arg1: memref<1024x32xf8E5M2>, %arg2: memref<16x32xbf16>,
+ %arg3: memref<16x32xf8E8M0FNU>, %arg4: memref<32x32xf8E8M0FNU>) {
%c0 = arith.constant 0 : index
%cst = arith.constant dense<0.000000e+00> : vector<16x32xbf16>
- %0 = xegpu.create_nd_tdesc %arg0 : memref<16x64xf8E5M2> -> !xegpu.tensor_desc<16x64xf8E5M2>
- %1 = xegpu.create_nd_tdesc %arg1 : memref<64x32xf8E5M2> -> !xegpu.tensor_desc<64x32xf8E5M2>
- %2 = xegpu.load_nd %0[0, 0] : !xegpu.tensor_desc<16x64xf8E5M2> -> vector<16x64xf8E5M2>
- %3 = xegpu.load_nd %1[0, 0] : !xegpu.tensor_desc<64x32xf8E5M2> -> vector<64x32xf8E5M2>
- %4 = xegpu.create_nd_tdesc %arg3 : memref<16x2xf8E8M0FNU> -> !xegpu.tensor_desc<16x2xf8E8M0FNU>
- %5 = xegpu.load_nd %4[0, 0] : !xegpu.tensor_desc<16x2xf8E8M0FNU> -> vector<16x2xf8E8M0FNU>
- %6 = xegpu.create_nd_tdesc %arg4 : memref<2x32xf8E8M0FNU> -> !xegpu.tensor_desc<2x32xf8E8M0FNU>
- %7 = xegpu.load_nd %6[0, 0] : !xegpu.tensor_desc<2x32xf8E8M0FNU> -> vector<2x32xf8E8M0FNU>
- %8 = xegpu.dpas_mx %2, %3, %cst scale_a = %5 scale_b = %7 : (vector<16x64xf8E5M2>, vector<64x32xf8E5M2>, vector<16x32xbf16>, vector<16x2xf8E8M0FNU>, vector<2x32xf8E8M0FNU>) -> vector<16x32xbf16>
+ %0 = xegpu.create_nd_tdesc %arg0 : memref<16x1024xf8E5M2> -> !xegpu.tensor_desc<16x1024xf8E5M2>
+ %1 = xegpu.create_nd_tdesc %arg1 : memref<1024x32xf8E5M2> -> !xegpu.tensor_desc<1024x32xf8E5M2>
+ %2 = xegpu.load_nd %0[0, 0] : !xegpu.tensor_desc<16x1024xf8E5M2> -> vector<16x1024xf8E5M2>
+ %3 = xegpu.load_nd %1[0, 0] : !xegpu.tensor_desc<1024x32xf8E5M2> -> vector<1024x32xf8E5M2>
+ %4 = xegpu.create_nd_tdesc %arg3 : memref<16x32xf8E8M0FNU> -> !xegpu.tensor_desc<16x32xf8E8M0FNU>
+ %5 = xegpu.load_nd %4[0, 0] : !xegpu.tensor_desc<16x32xf8E8M0FNU> -> vector<16x32xf8E8M0FNU>
+ %6 = xegpu.create_nd_tdesc %arg4 : memref<32x32xf8E8M0FNU> -> !xegpu.tensor_desc<32x32xf8E8M0FNU>
+ %7 = xegpu.load_nd %6[0, 0] : !xegpu.tensor_desc<32x32xf8E8M0FNU> -> vector<32x32xf8E8M0FNU>
+ %8 = xegpu.dpas_mx %2, %3, %cst scale_a = %5 scale_b = %7 : (vector<16x1024xf8E5M2>, vector<1024x32xf8E5M2>, vector<16x32xbf16>, vector<16x32xf8E8M0FNU>, vector<32x32xf8E8M0FNU>) -> vector<16x32xbf16>
%9 = xegpu.create_nd_tdesc %arg2 : memref<16x32xbf16> -> !xegpu.tensor_desc<16x32xbf16>
xegpu.store_nd %8, %9[0, 0] : vector<16x32xbf16>, !xegpu.tensor_desc<16x32xbf16>
return
@@ -491,40 +412,40 @@ func.func @dpas_mx_f8e5m2(%arg0: memref<16x64xf8E5M2>, %arg1: memref<64x32xf8E5M
// -----
// CHECK-LABEL: func.func @dpas_mx_f4e2m1
-// CHECK-SAME: %[[ARG0:[0-9a-zA-Z]+]]: memref<16x128xf4E2M1FN>, %[[ARG1:[0-9a-zA-Z]+]]: memref<128x32xf4E2M1FN>, %[[ARG2:[0-9a-zA-Z]+]]: memref<16x32xbf16>
-// CHECK-SAME: %[[ARG3:[0-9a-zA-Z]+]]: memref<16x4xf8E8M0FNU>, %[[ARG4:[0-9a-zA-Z]+]]: memref<4x32xf8E8M0FNU>
-// CHECK: %[[CST:.*]] = arith.constant {layout_result_0 = #xegpu.layout<inst_data = [8, 16]>} dense<0.000000e+00> : vector<16x32xbf16>
-// CHECK: %[[T0:.*]] = xegpu.create_nd_tdesc %[[ARG0]] : memref<16x128xf4E2M1FN> -> !xegpu.tensor_desc<16x128xf4E2M1FN, #xegpu.layout<inst_data = [8, 64]>>
-// CHECK: %[[T1:.*]] = xegpu.create_nd_tdesc %[[ARG1]] : memref<128x32xf4E2M1FN> -> !xegpu.tensor_desc<128x32xf4E2M1FN, #xegpu.layout<inst_data = [64, 16]>>
-// CHECK: %[[T2:.*]] = xegpu.load_nd %[[T0]][0, 0] <{layout = #xegpu.layout<inst_data = [8, 64]>}> :
-// CHECK-SAME: !xegpu.tensor_desc<16x128xf4E2M1FN, #xegpu.layout<inst_data = [8, 64]>> -> vector<16x128xf4E2M1FN>
-// CHECK: %[[T3:.*]] = xegpu.load_nd %[[T1]][0, 0] <{layout = #xegpu.layout<inst_data = [64, 16]>}> :
-// CHECK-SAME: !xegpu.tensor_desc<128x32xf4E2M1FN, #xegpu.layout<inst_data = [64, 16]>> -> vector<128x32xf4E2M1FN>
-// CHECK: %[[T4:.*]] = xegpu.create_nd_tdesc %[[ARG3]] : memref<16x4xf8E8M0FNU> -> !xegpu.tensor_desc<16x4xf8E8M0FNU, #xegpu.layout<inst_data = [8, 2]>>
-// CHECK: %[[T5:.*]] = xegpu.load_nd %[[T4]][0, 0] <{layout = #xegpu.layout<inst_data = [8, 2]>}> :
-// CHECK-SAME: !xegpu.tensor_desc<16x4xf8E8M0FNU, #xegpu.layout<inst_data = [8, 2]>> -> vector<16x4xf8E8M0FNU>
-// CHECK: %[[T6:.*]] = xegpu.create_nd_tdesc %[[ARG4]] : memref<4x32xf8E8M0FNU> -> !xegpu.tensor_desc<4x32xf8E8M0FNU, #xegpu.layout<inst_data = [2, 16]>>
-// CHECK: %[[T7:.*]] = xegpu.load_nd %[[T6]][0, 0] <{layout = #xegpu.layout<inst_data = [2, 16]>}> :
-// CHECK-SAME: !xegpu.tensor_desc<4x32xf8E8M0FNU, #xegpu.layout<inst_data = [2, 16]>> -> vector<4x32xf8E8M0FNU>
+// CHECK-SAME: %[[ARG0:[0-9a-zA-Z]+]]: memref<16x1024xf4E2M1FN>, %[[ARG1:[0-9a-zA-Z]+]]: memref<1024x32xf4E2M1FN>, %[[ARG2:[0-9a-zA-Z]+]]: memref<16x32xbf16>
+// CHECK-SAME: %[[ARG3:[0-9a-zA-Z]+]]: memref<16x32xf8E8M0FNU>, %[[ARG4:[0-9a-zA-Z]+]]: memref<32x32xf8E8M0FNU>
+// CHECK: %[[CST:.*]] = arith.constant {layout_result_0 = #xegpu.layout<inst_data = [8, 16], lane_layout = [1, 16], lane_data = [1, 1]>} dense<0.000000e+00> : vector<16x32xbf16>
+// CHECK: %[[T0:.*]] = xegpu.create_nd_tdesc %[[ARG0]] : memref<16x1024xf4E2M1FN> -> !xegpu.tensor_desc<16x1024xf4E2M1FN, #xegpu.layout<inst_data = [8, 64], lane_layout = [1, 16], lane_data = [1, 4]>>
+// CHECK: %[[T1:.*]] = xegpu.create_nd_tdesc %[[ARG1]] : memref<1024x32xf4E2M1FN> -> !xegpu.tensor_desc<1024x32xf4E2M1FN, #xegpu.layout<inst_data = [64, 16], lane_layout = [1, 16], lane_data = [8, 1]>>
+// CHECK: %[[T2:.*]] = xegpu.load_nd %[[T0]][0, 0] <{layout = #xegpu.layout<inst_data = [8, 64], lane_layout = [1, 16], lane_data = [1, 4]>}> :
+// CHECK-SAME: !xegpu.tensor_desc<16x1024xf4E2M1FN, #xegpu.layout<inst_data = [8, 64], lane_layout = [1, 16], lane_data = [1, 4]>> -> vector<16x1024xf4E2M1FN>
+// CHECK: %[[T3:.*]] = xegpu.load_nd %[[T1]][0, 0] <{layout = #xegpu.layout<inst_data = [64, 16], lane_layout = [1, 16], lane_data = [8, 1]>}> :
+// CHECK-SAME: !xegpu.tensor_desc<1024x32xf4E2M1FN, #xegpu.layout<inst_data = [64, 16], lane_layout = [1, 16], lane_data = [8, 1]>> -> vector<1024x32xf4E2M1FN>
+// CHECK: %[[T4:.*]] = xegpu.create_nd_tdesc %[[ARG3]] : memref<16x32xf8E8M0FNU> -> !xegpu.tensor_desc<16x32xf8E8M0FNU, #xegpu.layout<inst_data = [16, 32], lane_layout = [16, 1], lane_data = [1, 1]>>
+// CHECK: %[[T5:.*]] = xegpu.load_nd %[[T4]][0, 0] <{layout = #xegpu.layout<inst_data = [16, 32], lane_layout = [16, 1], lane_data = [1, 1]>}> :
+// CHECK-SAME: !xegpu.tensor_desc<16x32xf8E8M0FNU, #xegpu.layout<inst_data = [16, 32], lane_layout = [16, 1], lane_data = [1, 1]>> -> vector<16x32xf8E8M0FNU>
+// CHECK: %[[T6:.*]] = xegpu.create_nd_tdesc %[[ARG4]] : memref<32x32xf8E8M0FNU> -> !xegpu.tensor_desc<32x32xf8E8M0FNU, #xegpu.layout<inst_data = [2, 16], lane_layout = [1, 16], lane_data = [1, 1]>>
+// CHECK: %[[T7:.*]] = xegpu.load_nd %[[T6]][0, 0] <{layout = #xegpu.layout<inst_data = [2, 16], lane_layout = [1, 16], lane_data = [1, 1]>}> :
+// CHECK-SAME: !xegpu.tensor_desc<32x32xf8E8M0FNU, #xegpu.layout<inst_data = [2, 16], lane_layout = [1, 16], lane_data = [1, 1]>> -> vector<32x32xf8E8M0FNU>
// CHECK: %[[T8:.*]] = xegpu.dpas_mx %[[T2]], %[[T3]], %[[CST]] scale_a = %[[T5]] scale_b = %[[T7]]
-// CHECK-SAME: {layout_a = #xegpu.layout<inst_data = [8, 64]>, layout_a_scale = #xegpu.layout<inst_data = [8, 2]>, layout_b = #xegpu.layout<inst_data = [64, 16]>, layout_b_scale = #xegpu.layout<inst_data = [2, 16]>, layout_cd = #xegpu.layout<inst_data = [8, 16]>} :
-// CHECK-SAME: (vector<16x128xf4E2M1FN>, vector<128x32xf4E2M1FN>, vector<16x32xbf16>, vector<16x4xf8E8M0FNU>, vector<4x32xf8E8M0FNU>) -> vector<16x32xbf16>
-// CHECK: %[[T9:.*]] = xegpu.create_nd_tdesc %[[ARG2]] : memref<16x32xbf16> -> !xegpu.tensor_desc<16x32xbf16, #xegpu.layout<inst_data = [8, 16]>>
-// CHECK: xegpu.store_nd %[[T8]], %[[T9]][0, 0] <{layout = #xegpu.layout<inst_data = [8, 16]>}> : vector<16x32xbf16>, !xegpu.tensor_desc<16x32xbf16, #xegpu.layout<inst_data = [8, 16]>>
+// CHECK-SAME: {layout_a = #xegpu.layout<inst_data = [8, 64], lane_layout = [1, 16], lane_data = [1, 4]>, layout_a_scale = #xegpu.layout<inst_data = [8, 2], lane_layout = [8, 1], lane_data = [1, 1]>, layout_b = #xegpu.layout<inst_data = [64, 16], lane_layout = [1, 16], lane_data = [8, 1]>, layout_b_scale = #xegpu.layout<inst_data = [2, 16], lane_layout = [1, 16], lane_data = [1, 1]>, layout_cd = #xegpu.layout<inst_data = [8, 16], lane_layout = [1, 16], lane_data = [1, 1]>} :
+// CHECK-SAME: (vector<16x1024xf4E2M1FN>, vector<1024x32xf4E2M1FN>, vector<16x32xbf16>, vector<16x32xf8E8M0FNU>, vector<32x32xf8E8M0FNU>) -> vector<16x32xbf16>
+// CHECK: %[[T9:.*]] = xegpu.create_nd_tdesc %[[ARG2]] : memref<16x32xbf16> -> !xegpu.tensor_desc<16x32xbf16, #xegpu.layout<inst_data = [8, 16], lane_layout = [1, 16], lane_data = [1, 1]>>
+// CHECK: xegpu.store_nd %[[T8]], %[[T9]][0, 0] <{layout = #xegpu.layout<inst_data = [8, 16], lane_layout = [1, 16], lane_data = [1, 1]>}> : vector<16x32xbf16>, !xegpu.tensor_desc<16x32xbf16, #xegpu.layout<inst_data = [8, 16], lane_layout = [1, 16], lane_data = [1, 1]>>
gpu.module @test {
-func.func @dpas_mx_f4e2m1(%arg0: memref<16x128xf4E2M1FN>, %arg1: memref<128x32xf4E2M1FN>, %arg2: memref<16x32xbf16>,
- %arg3: memref<16x4xf8E8M0FNU>, %arg4: memref<4x32xf8E8M0FNU>) {
+func.func @dpas_mx_f4e2m1(%arg0: memref<16x1024xf4E2M1FN>, %arg1: memref<1024x32xf4E2M1FN>, %arg2: memref<16x32xbf16>,
+ %arg3: memref<16x32xf8E8M0FNU>, %arg4: memref<32x32xf8E8M0FNU>) {
%c0 = arith.constant 0 : index
%cst = arith.constant dense<0.000000e+00> : vector<16x32xbf16>
- %0 = xegpu.create_nd_tdesc %arg0 : memref<16x128xf4E2M1FN> -> !xegpu.tensor_desc<16x128xf4E2M1FN>
- %1 = xegpu.create_nd_tdesc %arg1 : memref<128x32xf4E2M1FN> -> !xegpu.tensor_desc<128x32xf4E2M1FN>
- %2 = xegpu.load_nd %0[0, 0] : !xegpu.tensor_desc<16x128xf4E2M1FN> -> vector<16x128xf4E2M1FN>
- %3 = xegpu.load_nd %1[0, 0] : !xegpu.tensor_desc<128x32xf4E2M1FN> -> vector<128x32xf4E2M1FN>
- %4 = xegpu.create_nd_tdesc %arg3 : memref<16x4xf8E8M0FNU> -> !xegpu.tensor_desc<16x4xf8E8M0FNU>
- %5 = xegpu.load_nd %4[0, 0] : !xegpu.tensor_desc<16x4xf8E8M0FNU> -> vector<16x4xf8E8M0FNU>
- %6 = xegpu.create_nd_tdesc %arg4 : memref<4x32xf8E8M0FNU> -> !xegpu.tensor_desc<4x32xf8E8M0FNU>
- %7 = xegpu.load_nd %6[0, 0] : !xegpu.tensor_desc<4x32xf8E8M0FNU> -> vector<4x32xf8E8M0FNU>
- %8 = xegpu.dpas_mx %2, %3, %cst scale_a = %5 scale_b = %7 : (vector<16x128xf4E2M1FN>, vector<128x32xf4E2M1FN>, vector<16x32xbf16>, vector<16x4xf8E8M0FNU>, vector<4x32xf8E8M0FNU>) -> vector<16x32xbf16>
+ %0 = xegpu.create_nd_tdesc %arg0 : memref<16x1024xf4E2M1FN> -> !xegpu.tensor_desc<16x1024xf4E2M1FN>
+ %1 = xegpu.create_nd_tdesc %arg1 : memref<1024x32xf4E2M1FN> -> !xegpu.tensor_desc<1024x32xf4E2M1FN>
+ %2 = xegpu.load_nd %0[0, 0] : !xegpu.tensor_desc<16x1024xf4E2M1FN> -> vector<16x1024xf4E2M1FN>
+ %3 = xegpu.load_nd %1[0, 0] : !xegpu.tensor_desc<1024x32xf4E2M1FN> -> vector<1024x32xf4E2M1FN>
+ %4 = xegpu.create_nd_tdesc %arg3 : memref<16x32xf8E8M0FNU> -> !xegpu.tensor_desc<16x32xf8E8M0FNU>
+ %5 = xegpu.load_nd %4[0, 0] : !xegpu.tensor_desc<16x32xf8E8M0FNU> -> vector<16x32xf8E8M0FNU>
+ %6 = xegpu.create_nd_tdesc %arg4 : memref<32x32xf8E8M0FNU> -> !xegpu.tensor_desc<32x32xf8E8M0FNU>
+ %7 = xegpu.load_nd %6[0, 0] : !xegpu.tensor_desc<32x32xf8E8M0FNU> -> vector<32x32xf8E8M0FNU>
+ %8 = xegpu.dpas_mx %2, %3, %cst scale_a = %5 scale_b = %7 : (vector<16x1024xf4E2M1FN>, vector<1024x32xf4E2M1FN>, vector<16x32xbf16>, vector<16x32xf8E8M0FNU>, vector<32x32xf8E8M0FNU>) -> vector<16x32xbf16>
%9 = xegpu.create_nd_tdesc %arg2 : memref<16x32xbf16> -> !xegpu.tensor_desc<16x32xbf16>
xegpu.store_nd %8, %9[0, 0] : vector<16x32xbf16>, !xegpu.tensor_desc<16x32xbf16>
return
@@ -607,3 +528,128 @@ func.func @vector_shape_cast_collapse_multi_groups(%arg0: memref<8x128xf16>) {
return
}
}
+
+// -----
+// completeBlockStoreLaneLayoutFromInstData: user supplies only inst_data on a
+// store_nd; lane_layout / lane_data are completed from it (data sink, no
+// consumer). inst_data=[8,16] -> lane_layout=[1,16], lane_data=[1,1].
+gpu.module @test {
+// CHECK-LABEL: func.func @complete_store_nd_inst_data(
+// CHECK: xegpu.store_nd %{{.*}} <{layout = #xegpu.layout<inst_data = [8, 16], lane_layout = [1, 16], lane_data = [1, 1]>}> : vector<8x32xf32>, !xegpu.tensor_desc<8x32xf32, #xegpu.layout<inst_data = [8, 16], lane_layout = [1, 16], lane_data = [1, 1]>>
+func.func @complete_store_nd_inst_data(%arg0: memref<8x32xf32>) {
+ %cst = arith.constant dense<0.000000e+00> : vector<8x32xf32>
+ %0 = xegpu.create_nd_tdesc %arg0 : memref<8x32xf32> -> !xegpu.tensor_desc<8x32xf32>
+ xegpu.store_nd %cst, %0[0, 0] <{layout = #xegpu.layout<inst_data = [8, 16]>}> : vector<8x32xf32>, !xegpu.tensor_desc<8x32xf32>
+ return
+}
+}
+
+// -----
+// completeBlockStoreLaneLayoutFromInstData (prefetch path): prefetch_nd is also
+// a data sink served by the same helper. inst_data=[8,16] -> [1,16]/[1,1].
+gpu.module @test {
+// CHECK-LABEL: func.func @complete_prefetch_nd_inst_data(
+// CHECK: xegpu.prefetch_nd %{{.*}} <{{{.*}}layout = #xegpu.layout<inst_data = [8, 16], lane_layout = [1, 16], lane_data = [1, 1]>}> : !xegpu.tensor_desc<8x32xf32, #xegpu.layout<inst_data = [8, 16], lane_layout = [1, 16], lane_data = [1, 1]>>
+func.func @complete_prefetch_nd_inst_data(%arg0: memref<8x32xf32>) {
+ %0 = xegpu.create_nd_tdesc %arg0 : memref<8x32xf32> -> !xegpu.tensor_desc<8x32xf32>
+ xegpu.prefetch_nd %0[0, 0] <{l1_hint = #xegpu.cache_hint<cached>, layout = #xegpu.layout<inst_data = [8, 16]>}> : !xegpu.tensor_desc<8x32xf32>
+ return
+}
+}
+
+// -----
+// completeBlockLoadLaneLayoutFromInstData: load_nd feeds a DPAS, so the consumer
+// supplies the transform / transpose / packing properties while lane_layout /
+// lane_data are recomputed from inst_data. A (inst=[8,16]) -> [1,16]/[1,1];
+// B (inst=[16,16], VNNI packing from the DPAS B consumer) -> [1,16]/[2,1].
+gpu.module @test {
+// CHECK-LABEL: func.func @complete_load_nd_inst_data(
+// CHECK: xegpu.load_nd %{{.*}} <{layout = #xegpu.layout<inst_data = [8, 16], lane_layout = [1, 16], lane_data = [1, 1]>}> : !xegpu.tensor_desc<8x16xf16, #xegpu.layout<inst_data = [8, 16], lane_layout = [1, 16], lane_data = [1, 1]>> -> vector<8x16xf16>
+// CHECK: xegpu.load_nd %{{.*}} <{layout = #xegpu.layout<inst_data = [16, 16], lane_layout = [1, 16], lane_data = [2, 1]>}> : !xegpu.tensor_desc<16x16xf16, #xegpu.layout<inst_data = [16, 16], lane_layout = [1, 16], lane_data = [2, 1]>> -> vector<16x16xf16>
+func.func @complete_load_nd_inst_data(%arg0: memref<8x16xf16>, %arg1: memref<16x16xf16>, %arg2: memref<8x16xf32>) {
+ %cst = arith.constant dense<0.000000e+00> : vector<8x16xf32>
+ %0 = xegpu.create_nd_tdesc %arg0 : memref<8x16xf16> -> !xegpu.tensor_desc<8x16xf16>
+ %1 = xegpu.create_nd_tdesc %arg1 : memref<16x16xf16> -> !xegpu.tensor_desc<16x16xf16>
+ %2 = xegpu.load_nd %0[0, 0] <{layout = #xegpu.layout<inst_data = [8, 16]>}> : !xegpu.tensor_desc<8x16xf16> -> vector<8x16xf16>
+ %3 = xegpu.load_nd %1[0, 0] <{layout = #xegpu.layout<inst_data = [16, 16]>}> : !xegpu.tensor_desc<16x16xf16> -> vector<16x16xf16>
+ %4 = xegpu.dpas %2, %3, %cst : vector<8x16xf16>, vector<16x16xf16>, vector<8x16xf32> -> vector<8x16xf32>
+ %5 = xegpu.create_nd_tdesc %arg2 : memref<8x16xf32> -> !xegpu.tensor_desc<8x16xf32>
+ xegpu.store_nd %4, %5[0, 0] : vector<8x16xf32>, !xegpu.tensor_desc<8x16xf32>
+ return
+}
+}
+
+// -----
+// completeScatterStoreLaneLayoutFromInstData: user supplies only inst_data on a
+// scatter store; lane info derived purely from inst_data (data sink).
+// inst_data=[1,16] -> lane_layout=[1,16], lane_data=[1,1].
+gpu.module @test {
+// CHECK-LABEL: func.func @complete_scatter_store_inst_data(
+// CHECK: xegpu.store %{{.*}} <{layout = #xegpu.layout<inst_data = [1, 16], lane_layout = [1, 16], lane_data = [1, 1]>}> : vector<16x32xf32>, memref<512xf32>, vector<16x32xindex>, vector<16x32xi1>
+func.func @complete_scatter_store_inst_data(%src: memref<512xf32>) {
+ %mask = arith.constant dense<1> : vector<16x32xi1>
+ %offset = arith.constant dense<12> : vector<16x32xindex>
+ %data = arith.constant dense<0.000000e+00> : vector<16x32xf32>
+ xegpu.store %data, %src[%offset], %mask <{layout = #xegpu.layout<inst_data = [1, 16]>}>
+ : vector<16x32xf32>, memref<512xf32>, vector<16x32xindex>, vector<16x32xi1>
+ return
+}
+}
+
+// -----
+// completeScatterLoadLaneLayoutFromInstData: user supplies only inst_data on a
+// scatter load; with no usable consumer lane info, the scatter default is used.
+// inst_data=[1,16] -> lane_layout=[1,16], lane_data=[1,1].
+gpu.module @test {
+// CHECK-LABEL: func.func @complete_scatter_load_inst_data(
+// CHECK: xegpu.load %{{.*}} <{layout = #xegpu.layout<inst_data = [1, 16], lane_layout = [1, 16], lane_data = [1, 1]>}> : memref<512xf32>, vector<16x32xindex>, vector<16x32xi1> -> vector<16x32xf32>
+func.func @complete_scatter_load_inst_data(%src: memref<512xf32>) {
+ %mask = arith.constant dense<1> : vector<16x32xi1>
+ %offset = arith.constant dense<12> : vector<16x32xindex>
+ %0 = xegpu.load %src[%offset], %mask <{layout = #xegpu.layout<inst_data = [1, 16]>}>
+ : memref<512xf32>, vector<16x32xindex>, vector<16x32xi1> -> vector<16x32xf32>
+ xegpu.store %0, %src[%offset], %mask <{layout = #xegpu.layout<inst_data = [1, 16]>}>
+ : vector<16x32xf32>, memref<512xf32>, vector<16x32xindex>, vector<16x32xi1>
+ return
+}
+}
+
+// -----
+// completeDpasLaneLayoutFromInstData: user supplies only inst_data on all three
+// DPAS operands; lane info is completed from each operand's shape / matmul role.
+// A=[8,16]->[1,16]/[1,1]; B=[16,16]->[1,16]/[2,1] (VNNI); CD=[8,16]->[1,16]/[1,1].
+gpu.module @test {
+// CHECK-LABEL: func.func @complete_dpas_inst_data(
+// CHECK: xegpu.dpas %{{.*}}, %{{.*}}, %{{.*}} {layout_a = #xegpu.layout<inst_data = [8, 16], lane_layout = [1, 16], lane_data = [1, 1]>, layout_b = #xegpu.layout<inst_data = [16, 16], lane_layout = [1, 16], lane_data = [2, 1]>, layout_cd = #xegpu.layout<inst_data = [8, 16], lane_layout = [1, 16], lane_data = [1, 1]>} : vector<8x16xf16>, vector<16x16xf16>, vector<8x16xf32> -> vector<8x16xf32>
+func.func @complete_dpas_inst_data(%arg0: vector<8x16xf16>, %arg1: vector<16x16xf16>) {
+ %cst = arith.constant dense<0.000000e+00> : vector<8x16xf32>
+ %0 = xegpu.dpas %arg0, %arg1, %cst {
+ layout_a = #xegpu.layout<inst_data = [8, 16]>,
+ layout_b = #xegpu.layout<inst_data = [16, 16]>,
+ layout_cd = #xegpu.layout<inst_data = [8, 16]>}
+ : vector<8x16xf16>, vector<16x16xf16>, vector<8x16xf32> -> vector<8x16xf32>
+ return
+}
+}
+
+// -----
+// completeDpasMxLaneLayoutFromInstData: user supplies only inst_data on A/B/C-D;
+// lane info completed from shapes and scale layouts re-derived via
+// createScaleLayout. Matches the dpas_mx_f8e5m2 default-path result.
+gpu.module @test {
+// CHECK-LABEL: func.func @complete_dpas_mx_inst_data(
+// CHECK: xegpu.dpas_mx %{{.*}}, %{{.*}}, %{{.*}} scale_a = %{{[0-9a-zA-Z]+}} scale_b = %{{[0-9a-zA-Z]+}}
+// CHECK-SAME: {layout_a = #xegpu.layout<inst_data = [8, 32], lane_layout = [1, 16], lane_data = [1, 2]>, layout_a_scale = #xegpu.layout<inst_data = [8, 1], lane_layout = [8, 1], lane_data = [1, 1]>, layout_b = #xegpu.layout<inst_data = [32, 16], lane_layout = [1, 16], lane_data = [4, 1]>, layout_b_scale = #xegpu.layout<inst_data = [1, 16], lane_layout = [1, 16], lane_data = [1, 1]>, layout_cd = #xegpu.layout<inst_data = [8, 16], lane_layout = [1, 16], lane_data = [1, 1]>} :
+func.func @complete_dpas_mx_inst_data(%arg0: vector<16x1024xf8E5M2>, %arg1: vector<1024x32xf8E5M2>,
+ %arg2: vector<16x32xf8E8M0FNU>, %arg3: vector<32x32xf8E8M0FNU>) {
+ %cst = arith.constant dense<0.000000e+00> : vector<16x32xbf16>
+ %0 = xegpu.dpas_mx %arg0, %arg1, %cst scale_a = %arg2 scale_b = %arg3 {
+ layout_a = #xegpu.layout<inst_data = [8, 32]>,
+ layout_a_scale = #xegpu.layout<inst_data = [8, 1]>,
+ layout_b = #xegpu.layout<inst_data = [32, 16]>,
+ layout_b_scale = #xegpu.layout<inst_data = [1, 16]>,
+ layout_cd = #xegpu.layout<inst_data = [8, 16]>}
+ : (vector<16x1024xf8E5M2>, vector<1024x32xf8E5M2>, vector<16x32xbf16>, vector<16x32xf8E8M0FNU>, vector<32x32xf8E8M0FNU>) -> vector<16x32xbf16>
+ return
+}
+}
diff --git a/mlir/test/Dialect/XeGPU/propagate-layout.mlir b/mlir/test/Dialect/XeGPU/propagate-layout.mlir
index bad956d45d186..25d713ccc8a0f 100644
--- a/mlir/test/Dialect/XeGPU/propagate-layout.mlir
+++ b/mlir/test/Dialect/XeGPU/propagate-layout.mlir
@@ -127,26 +127,29 @@ func.func @extf_truncf(%arg0: !xegpu.tensor_desc<8x16xf16>, %arg1: !xegpu.tensor
%2 = arith.extf %1 : vector<16x16xf16> to vector<16x16xf32>
%3 = arith.truncf %2 : vector<16x16xf32> to vector<16x16xf16>
%4 = xegpu.dpas %0, %3 : vector<8x16xf16>, vector<16x16xf16> -> vector<8x16xf32>
+ %5 = xegpu.convert_layout %4
+ <{input_layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>,
+ target_layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>}>
+ : vector<8x16xf32>
return %4 : vector<8x16xf32>
}
}
// -----
gpu.module @test {
-// CHECK-LABEL: func.func @load_gather_with_chunksize(
+// CHECK-LABEL: func.func @load_gather_with_coalesce_chunksize(
// CHECK-SAME: %[[ARG0:[0-9a-zA-Z]+]]: memref<8x16xf16>, %[[ARG1:[0-9a-zA-Z]+]]: memref<256xf16>, %[[ARG2:[0-9a-zA-Z]+]]: memref<8x16xf32>) {
-// CHECK: %[[OFFSET:.*]] = arith.constant {layout_result_0 = #xegpu.layout<lane_layout = [16], lane_data = [1]>}
-// CHECK-SAME: dense<[0, 16, 32, 48, 64, 80, 96, 112, 128, 144, 160, 176, 192, 208, 224, 240]> : vector<16xindex>
-// CHECK-NEXT: %[[MASK:.*]] = arith.constant {layout_result_0 = #xegpu.layout<lane_layout = [16], lane_data = [1]>} dense<true> : vector<16xi1>
-// CHECK-NEXT: %{{.*}} = xegpu.load %arg1[%[[OFFSET]]], %[[MASK]] <{chunk_size = 16 : i64, layout = #xegpu.layout<lane_layout = [16, 1], lane_data = [1, 2]>}> : memref<256xf16>, vector<16xindex>, vector<16xi1> -> vector<16x16xf16>
-func.func @load_gather_with_chunksize(%arg0: memref<8x16xf16>, %arg1: memref<256xf16>, %arg2: memref<8x16xf32>) {
+// CHECK: %[[OFFSET:.*]] = arith.constant {layout_result_0 = #xegpu.layout<lane_layout = [16, 1], lane_data = [1, 2]>} dense<0> : vector<16x16xindex>
+// CHECK: %[[MASK:.*]] = arith.constant {layout_result_0 = #xegpu.layout<lane_layout = [16, 1], lane_data = [1, 2]>} dense<true> : vector<16x16xi1>
+// CHECK: %{{.*}} = xegpu.load %arg1[%[[OFFSET]]], %[[MASK]] <{layout = #xegpu.layout<lane_layout = [16, 1], lane_data = [1, 2]>}> : memref<256xf16>, vector<16x16xindex>, vector<16x16xi1> -> vector<16x16xf16>
+func.func @load_gather_with_coalesce_chunksize(%arg0: memref<8x16xf16>, %arg1: memref<256xf16>, %arg2: memref<8x16xf32>) {
%c0 = arith.constant 0 : index
%0 = xegpu.create_nd_tdesc %arg0 : memref<8x16xf16> -> !xegpu.tensor_desc<8x16xf16>
%1 = xegpu.load_nd %0[0, 0] : !xegpu.tensor_desc<8x16xf16> -> vector<8x16xf16>
- %offset = arith.constant dense<[0, 16, 32, 48, 64, 80, 96, 112, 128, 144, 160, 176, 192, 208, 224, 240]> : vector<16xindex>
- %mask = arith.constant dense<true> : vector<16xi1>
- %3 = xegpu.load %arg1[%offset], %mask <{chunk_size=16}>
- : memref<256xf16>, vector<16xindex>, vector<16xi1> -> vector<16x16xf16>
+ %offset = arith.constant dense<0> : vector<16x16xindex>
+ %mask = arith.constant dense<true> : vector<16x16xi1>
+ %3 = xegpu.load %arg1[%offset], %mask
+ : memref<256xf16>, vector<16x16xindex>, vector<16x16xi1> -> vector<16x16xf16>
%4 = vector.transpose %3, [1, 0] : vector<16x16xf16> to vector<16x16xf16>
%5 = xegpu.dpas %1, %4 : vector<8x16xf16>, vector<16x16xf16> -> vector<8x16xf32>
%6 = xegpu.create_nd_tdesc %arg2 : memref<8x16xf32> -> !xegpu.tensor_desc<8x16xf32>
@@ -157,37 +160,20 @@ func.func @load_gather_with_chunksize(%arg0: memref<8x16xf16>, %arg1: memref<256
// -----
gpu.module @test {
-// CHECK-LABEL: func.func @store_scatter_with_chunksize(
-// CHECK-SAME: %[[ARG0:[0-9a-zA-Z]+]]: memref<128xf32>) {
-// CHECK-NEXT: %[[CST:.*]] = arith.constant {layout_result_0 = #xegpu.layout<lane_layout = [16, 1], lane_data = [1, 8]>} dense<1.000000e+00> : vector<16x8xf32>
-// CHECK-NEXT: %[[CST_0:.*]] = arith.constant {layout_result_0 = #xegpu.layout<lane_layout = [16], lane_data = [1]>} dense<true> : vector<16xi1>
-// CHECK-NEXT: %[[CST_1:.*]] = arith.constant {layout_result_0 = #xegpu.layout<lane_layout = [16], lane_data = [1]>} dense<[0, 16, 32, 48, 64, 80, 96, 112, 128, 144, 160, 176, 192, 208, 224, 240]> : vector<16xindex>
-// CHECK-NEXT: xegpu.store %[[CST]], %[[ARG0]][%[[CST_1]]], %[[CST_0]] <{chunk_size = 8 : i64, layout = #xegpu.layout<lane_layout = [16, 1], lane_data = [1, 8]>}> : vector<16x8xf32>, memref<128xf32>, vector<16xindex>, vector<16xi1>
-func.func @store_scatter_with_chunksize(%arg0: memref<128xf32>) {
- %val = arith.constant dense<1.000000e+00> : vector<16x8xf32>
- %mask = arith.constant dense<true> : vector<16xi1>
- %offset = arith.constant dense<[0, 16, 32, 48, 64, 80, 96, 112, 128, 144, 160, 176, 192, 208, 224, 240]> : vector<16xindex>
- xegpu.store %val, %arg0[%offset], %mask <{chunk_size = 8}>: vector<16x8xf32>, memref<128xf32>, vector<16xindex>, vector<16xi1>
- return
-}
-}
-
-// -----
-gpu.module @test {
-// CHECK-LABEL: func.func @scatter_ops_chunksize(
-// CHECK-SAME: %[[ARG0:[0-9a-zA-Z]+]]: memref<256xf16>) {
-// CHECK: %[[MASK:.*]] = arith.constant {layout_result_0 = #xegpu.layout<lane_layout = [16], lane_data = [1]>} dense<true> : vector<16xi1>
-// CHECK: %[[OFFSETS:.*]] = arith.constant {layout_result_0 = #xegpu.layout<lane_layout = [16], lane_data = [1]>} dense<12> : vector<16xindex>
-// CHECK: %[[LOAD_VEC:.*]] = xegpu.load %[[ARG0]][%[[OFFSETS]]], %[[MASK]] <{chunk_size = 8 : i64, layout = #xegpu.layout<lane_layout = [16, 1], lane_data = [1, 8]>}>
-// CHECK-SAME: memref<256xf16>, vector<16xindex>, vector<16xi1> -> vector<16x8xf16>
-// CHECK: xegpu.store %[[LOAD_VEC]], %[[ARG0]][%[[OFFSETS]]], %[[MASK]] <{chunk_size = 8 : i64, layout = #xegpu.layout<lane_layout = [16, 1], lane_data = [1, 8]>}> : vector<16x8xf16>, memref<256xf16>, vector<16xindex>, vector<16xi1>
-func.func @scatter_ops_chunksize(%src: memref<256xf16>) {
- %1 = arith.constant dense<1>: vector<16xi1>
- %offset = arith.constant dense<12> : vector<16xindex>
- %3 = xegpu.load %src[%offset], %1 <{chunk_size=8}>
- : memref<256xf16>, vector<16xindex>, vector<16xi1> -> vector<16x8xf16>
- xegpu.store %3, %src[%offset], %1 <{chunk_size=8}>
- : vector<16x8xf16>, memref<256xf16>, vector<16xindex>, vector<16xi1>
+// CHECK-LABEL: func.func @scatter_ops_coalesce_chunksize(
+// CHECK-SAME: %[[ARG0:[0-9a-zA-Z]+]]: memref<128xf16>) {
+// CHECK: %[[MASK:.*]] = arith.constant {layout_result_0 = #xegpu.layout<lane_layout = [1, 8], lane_data = [1, 1]>} dense<true> : vector<16x8xi1>
+// CHECK: %[[OFFSETS:.*]] = arith.constant {layout_result_0 = #xegpu.layout<lane_layout = [1, 8], lane_data = [1, 1]>} dense<12> : vector<16x8xindex>
+// CHECK: %[[LOAD_VEC:.*]] = xegpu.load %[[ARG0]][%[[OFFSETS]]], %[[MASK]] <{layout = #xegpu.layout<lane_layout = [1, 8], lane_data = [1, 1]>}>
+// CHECK-SAME: memref<128xf16>, vector<16x8xindex>, vector<16x8xi1> -> vector<16x8xf16>
+// CHECK: xegpu.store %[[LOAD_VEC]], %[[ARG0]][%[[OFFSETS]]], %[[MASK]] <{layout = #xegpu.layout<lane_layout = [1, 8], lane_data = [1, 1]>}> : vector<16x8xf16>, memref<128xf16>, vector<16x8xindex>, vector<16x8xi1>
+func.func @scatter_ops_coalesce_chunksize(%src: memref<128xf16>) {
+ %1 = arith.constant dense<1>: vector<16x8xi1>
+ %offset = arith.constant dense<12> : vector<16x8xindex>
+ %3 = xegpu.load %src[%offset], %1
+ : memref<128xf16>, vector<16x8xindex>, vector<16x8xi1> -> vector<16x8xf16>
+ xegpu.store %3, %src[%offset], %1
+ : vector<16x8xf16>, memref<128xf16>, vector<16x8xindex>, vector<16x8xi1>
return
}
}
@@ -473,37 +459,45 @@ func.func @if_multiple_uses(%arg0: !xegpu.tensor_desc<8x16xf16>, %arg1: !xegpu.t
// -----
gpu.module @test {
// CHECK-LABEL: func.func @vector_outer_reduction(
-// CHECK-SAME: %[[ARG0:[0-9a-zA-Z]+]]: vector<16x16xf32>, %[[ARG1:[0-9a-zA-Z]+]]: !xegpu.tensor_desc<16xf32, #xegpu.layout<lane_layout = [16], lane_data = [1]>>) {
+// CHECK-SAME: %[[ARG0:[0-9a-zA-Z]+]]: vector<16x16xf32>, %[[ARG1:[0-9a-zA-Z]+]]: memref<256xf32>) {
// CHECK: %{{.*}} = vector.multi_reduction <add>, %[[ARG0]], %{{.*}} {layout_result_0 = #xegpu.slice<#xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>, dims = [0]>} [0] : vector<16x16xf32> to vector<16xf32>
-func.func @vector_outer_reduction(%arg0: vector<16x16xf32>, %arg1: !xegpu.tensor_desc<16xf32>) {
+func.func @vector_outer_reduction(%arg0: vector<16x16xf32>, %arg1: memref<256xf32>) {
%cst = arith.constant dense<0.000000e+00> : vector<16xf32>
+ %mask = arith.constant dense<true> : vector<16xi1>
+ %offset = vector.step : vector<16xindex>
%0 = vector.multi_reduction <add>, %arg0, %cst [0] : vector<16x16xf32> to vector<16xf32>
- xegpu.store_nd %0, %arg1[0] : vector<16xf32>, !xegpu.tensor_desc<16xf32>
+ xegpu.store %0, %arg1[%offset], %mask : vector<16xf32>, memref<256xf32>, vector<16xindex>, vector<16xi1>
return
}
}
// -----
gpu.module @test {
// CHECK-LABEL: func.func @vector_inner_reduction(
-// CHECK-SAME: %[[ARG0:[0-9a-zA-Z]+]]: vector<16x16xf32>, %[[ARG1:[0-9a-zA-Z]+]]: !xegpu.tensor_desc<16xf32, #xegpu.layout<lane_layout = [16], lane_data = [1]>>) {
-// CHECK: %{{.*}} = vector.multi_reduction <add>, %[[ARG0]], %{{.*}} {layout_result_0 = #xegpu.slice<#xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>, dims = [1]>} [1] : vector<16x16xf32> to vector<16xf32>
-func.func @vector_inner_reduction(%arg0: vector<16x16xf32>, %arg1: !xegpu.tensor_desc<16xf32>) {
+// CHECK-SAME: %[[ARG0:[0-9a-zA-Z]+]]: vector<16x16xf32>, %[[ARG1:[0-9a-zA-Z]+]]: memref<256xf32>) {
+// CHECK: %{{.*}} = vector.multi_reduction <add>, %[[ARG0]], %{{.*}} {layout_result_0 = #xegpu.slice<#xegpu.layout<lane_layout = [16, 1], lane_data = [1, 1]>, dims = [1]>} [1] : vector<16x16xf32> to vector<16xf32>
+func.func @vector_inner_reduction(%arg0: vector<16x16xf32>, %arg1: memref<256xf32>) {
%cst = arith.constant dense<0.000000e+00> : vector<16xf32>
+ %mask = arith.constant dense<true> : vector<16xi1>
+ %offset = vector.step : vector<16xindex>
%0 = vector.multi_reduction <add>, %arg0, %cst [1] : vector<16x16xf32> to vector<16xf32>
- xegpu.store_nd %0, %arg1[0] : vector<16xf32>, !xegpu.tensor_desc<16xf32>
+ xegpu.store %0, %arg1[%offset], %mask : vector<16xf32>, memref<256xf32>, vector<16xindex>, vector<16xi1>
return
}
}
// -----
gpu.module @test {
-// CHECK-LABEL: func.func @store_nd_with_offset(
+// CHECK-LABEL: func.func @store_with_offset(
// CHECK-SAME: %[[ARG0:[0-9a-zA-Z]+]]: memref<256xf32>) {
-// CHECK: %[[T0:.*]] = xegpu.create_nd_tdesc %[[ARG0]] : memref<256xf32> -> !xegpu.tensor_desc<16xf32, #xegpu.layout<lane_layout = [16], lane_data = [1]>>
-func.func @store_nd_with_offset(%arg0: memref<256xf32>){
+// CHECK: %[[MASK:.*]] = arith.constant {layout_result_0 = #xegpu.layout<lane_layout = [16], lane_data = [1]>} dense<true> : vector<16xi1>
+// CHECK: %[[OFFSET:.*]] = vector.step {layout_result_0 = #xegpu.layout<lane_layout = [16], lane_data = [1]>} : vector<16xindex>
+// CHECK: %[[VAL:.*]] = arith.constant {layout_result_0 = #xegpu.layout<lane_layout = [16], lane_data = [1]>} dense<1.000000e+00> : vector<16xf32>
+// CHECK: xegpu.store %[[VAL]], %[[ARG0]][%[[OFFSET]]], %[[MASK]] <{layout = #xegpu.layout<lane_layout = [16], lane_data = [1]>}> : vector<16xf32>, memref<256xf32>, vector<16xindex>, vector<16xi1>
+func.func @store_with_offset(%arg0: memref<256xf32>){
%c32 = arith.constant 32 : index
+ %mask = arith.constant dense<true> : vector<16xi1>
+ %offset = vector.step : vector<16xindex>
%1 = arith.constant dense<1.000000e+00> : vector<16xf32>
- %0 = xegpu.create_nd_tdesc %arg0 : memref<256xf32> -> !xegpu.tensor_desc<16xf32>
- xegpu.store_nd %1, %0[%c32] : vector<16xf32>, !xegpu.tensor_desc<16xf32>
+ xegpu.store %1, %arg0[%offset], %mask : vector<16xf32>, memref<256xf32>, vector<16xindex>, vector<16xi1>
return
}
}
@@ -537,12 +531,10 @@ func.func @prefetch_2d(%arg0: memref<256x256xf16>){
gpu.module @test {
// CHECK-LABEL: func.func @prefetch_1d(
// CHECK-SAME: %[[ARG0:[0-9a-zA-Z]+]]: memref<256xf16>) {
-// CHECK: %[[T0:.*]] = xegpu.create_nd_tdesc %[[ARG0]] : memref<256xf16> -> !xegpu.tensor_desc<16xf16, #xegpu.layout<lane_layout = [16], lane_data = [1]>>
-// CHECK-NEXT: xegpu.prefetch_nd %[[T0]][0] <{l1_hint = #xegpu.cache_hint<cached>, l2_hint = #xegpu.cache_hint<uncached>, layout = #xegpu.layout<lane_layout = [16], lane_data = [1]>}> : !xegpu.tensor_desc<16xf16, #xegpu.layout<lane_layout = [16], lane_data = [1]>>
+// CHECK: xegpu.prefetch %[[ARG0]][%{{.*}}] <{l1_hint = #xegpu.cache_hint<cached>, l2_hint = #xegpu.cache_hint<uncached>, layout = #xegpu.layout<lane_layout = [16], lane_data = [1]>}> : memref<256xf16>, vector<16xindex>
func.func @prefetch_1d(%arg0: memref<256xf16>){
- %c0 = arith.constant 0 : index
- %0 = xegpu.create_nd_tdesc %arg0 : memref<256xf16> -> !xegpu.tensor_desc<16xf16>
- xegpu.prefetch_nd %0[0] <{l1_hint = #xegpu.cache_hint<cached>, l2_hint = #xegpu.cache_hint<uncached>}>: !xegpu.tensor_desc<16xf16>
+ %offset = vector.step : vector<16xindex>
+ xegpu.prefetch %arg0[%offset] <{l1_hint = #xegpu.cache_hint<cached>, l2_hint = #xegpu.cache_hint<uncached>, layout = #xegpu.layout<lane_layout = [16], lane_data = [1]>}> : memref<256xf16>, vector<16xindex>
return
}
}
@@ -562,9 +554,9 @@ func.func @scf_while_and_condition(%arg0: memref<256xf32>, %arg1: memref<256xf32
%c16 = arith.constant 16 : i32
%c16_idx = arith.constant 16 : index
%c256 = arith.constant 256 : i32
- %0 = xegpu.create_nd_tdesc %arg0 : memref<256xf32> -> !xegpu.tensor_desc<16xf32>
- %1 = xegpu.load_nd %0[0] : !xegpu.tensor_desc<16xf32> -> vector<16xf32>
- %2 = xegpu.create_nd_tdesc %arg1 : memref<256xf32> -> !xegpu.tensor_desc<16xf32>
+ %mask = arith.constant dense<true> : vector<16xi1>
+ %offset = vector.step : vector<16xindex>
+ %1 = xegpu.load %arg0[%offset], %mask : memref<256xf32>, vector<16xindex>, vector<16xi1> -> vector<16xf32>
%3:2 = scf.while (%arg2 = %1, %arg3 = %c0) : (vector<16xf32>, i32)
-> (vector<16xf32>, i32) {
@@ -572,12 +564,14 @@ func.func @scf_while_and_condition(%arg0: memref<256xf32>, %arg1: memref<256xf32
scf.condition(%4) %arg2, %arg3 : vector<16xf32>, i32
} do {
^bb0(%arg2: vector<16xf32>, %arg3: i32):
- xegpu.store_nd %arg2, %2[0] : vector<16xf32>, !xegpu.tensor_desc<16xf32>
+ xegpu.store %arg2, %arg1[%offset], %mask : vector<16xf32>, memref<256xf32>, vector<16xindex>, vector<16xi1>
%4 = arith.addi %arg3, %c16 : i32
- %offset = arith.index_cast %4 : i32 to index
- %6 = xegpu.load_nd %0[%offset] : !xegpu.tensor_desc<16xf32> -> vector<16xf32>
+ %offset2 = arith.index_cast %4 : i32 to index
+ %offset2_v = vector.broadcast %offset2 : index to vector<16xindex>
+ %6 = xegpu.load %arg0[%offset], %mask : memref<256xf32>, vector<16xindex>, vector<16xi1> -> vector<16xf32>
scf.yield %6, %4 : vector<16xf32>, i32
}
+
return
}
}
@@ -656,20 +650,20 @@ gpu.module @test{
// -----
gpu.module @test {
// CHECK-LABEL: func.func @vector_shape_cast_expand_non_unit_dims(
-// CHECK: %[[LOAD:.*]] = xegpu.load %arg0[%[[STEP:.*]]], %[[CST:.*]] <{layout = #xegpu.layout<lane_layout = [16], lane_data = [1]>}> : memref<1024xf16>, vector<1024xindex>, vector<1024xi1> -> vector<1024xf16>
-// CHECK: %[[CAST:.*]] = vector.shape_cast %[[LOAD]] {layout_result_0 = #xegpu.layout<lane_layout = [1, 1, 16], lane_data = [1, 1, 1]>} : vector<1024xf16> to vector<8x8x16xf16>
+// CHECK: %[[LOAD:.*]] = xegpu.load %arg0[%[[STEP:.*]]], %[[CST:.*]] <{layout = #xegpu.layout<lane_layout = [16], lane_data = [1]>}> : memref<128xf16>, vector<128xindex>, vector<128xi1> -> vector<128xf16>
+// CHECK: %[[CAST:.*]] = vector.shape_cast %[[LOAD]] {layout_result_0 = #xegpu.layout<lane_layout = [1, 1, 16], lane_data = [1, 1, 1]>} : vector<128xf16> to vector<1x8x16xf16>
// CHECK: %[[CST_0:.*]] = arith.constant {layout_result_0 = #xegpu.slice<#xegpu.layout<lane_layout = [1, 1, 16], lane_data = [1, 1, 1]>, dims = [0]>} dense<0.000000e+00> : vector<8x16xf16>
// CHECK: %[[CST_1:.*]] = arith.constant {layout_result_0 = #xegpu.slice<#xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>, dims = [0]>} dense<0.000000e+00> : vector<16xf16>
-// CHECK: %[[REDUCE_0:.*]] = vector.multi_reduction <add>, %[[CAST]], %[[CST_0]] {layout_result_0 = #xegpu.slice<#xegpu.layout<lane_layout = [1, 1, 16], lane_data = [1, 1, 1]>, dims = [0]>} [0] : vector<8x8x16xf16> to vector<8x16xf16>
+// CHECK: %[[REDUCE_0:.*]] = vector.multi_reduction <add>, %[[CAST]], %[[CST_0]] {layout_result_0 = #xegpu.slice<#xegpu.layout<lane_layout = [1, 1, 16], lane_data = [1, 1, 1]>, dims = [0]>} [0] : vector<1x8x16xf16> to vector<8x16xf16>
// CHECK: %[[REDUCE_1:.*]] = vector.multi_reduction <add>, %[[REDUCE_0]], %[[CST_1]] {layout_result_0 = #xegpu.slice<#xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>, dims = [0]>} [0] : vector<8x16xf16> to vector<16xf16>
-func.func @vector_shape_cast_expand_non_unit_dims(%arg0: memref<1024xf16>, %arg1: memref<16xf16>) {
- %cst = arith.constant dense<true> : vector<1024xi1>
- %0 = vector.step : vector<1024xindex>
- %1 = xegpu.load %arg0[%0], %cst : memref<1024xf16>, vector<1024xindex>, vector<1024xi1> -> vector<1024xf16>
- %2 = vector.shape_cast %1 : vector<1024xf16> to vector<8x8x16xf16>
+func.func @vector_shape_cast_expand_non_unit_dims(%arg0: memref<128xf16>, %arg1: memref<16xf16>) {
+ %cst = arith.constant dense<true> : vector<128xi1>
+ %0 = vector.step : vector<128xindex>
+ %1 = xegpu.load %arg0[%0], %cst : memref<128xf16>, vector<128xindex>, vector<128xi1> -> vector<128xf16>
+ %2 = vector.shape_cast %1 : vector<128xf16> to vector<1x8x16xf16>
%cst_0 = arith.constant dense<0.000000e+00> : vector<8x16xf16>
%cst_1 = arith.constant dense<0.000000e+00> : vector<16xf16>
- %3 = vector.multi_reduction <add>, %2, %cst_0 [0] : vector<8x8x16xf16> to vector<8x16xf16>
+ %3 = vector.multi_reduction <add>, %2, %cst_0 [0] : vector<1x8x16xf16> to vector<8x16xf16>
%4 = vector.multi_reduction <add>, %3, %cst_1 [0] : vector<8x16xf16> to vector<16xf16>
%cst_2 = arith.constant dense<true> : vector<16xi1>
%cst_3 = arith.constant dense<1> : vector<16xindex>
@@ -983,40 +977,40 @@ gpu.module @test{
// -----
gpu.module @test {
-// CHECK-LABEL: func.func @dpas_mx_f8e5m2
-// CHECK-SAME: %[[ARG0:[0-9a-zA-Z]+]]: memref<8x32xf8E5M2>, %[[ARG1:[0-9a-zA-Z]+]]: memref<32x16xf8E5M2>, %[[ARG2:[0-9a-zA-Z]+]]: memref<8x16xbf16>
-// CHECK-SAME: %[[ARG3:[0-9a-zA-Z]+]]: memref<8x1xf8E8M0FNU>, %[[ARG4:[0-9a-zA-Z]+]]: memref<1x16xf8E8M0FNU>
+// CHECK-LABEL: func.func @dpas_mx_fp4
+// CHECK-SAME: %[[ARG0:[0-9a-zA-Z]+]]: memref<8x64xf4E2M1FN>, %[[ARG1:[0-9a-zA-Z]+]]: memref<64x16xf4E2M1FN>, %[[ARG2:[0-9a-zA-Z]+]]: memref<8x16xbf16>
+// CHECK-SAME: %[[ARG3:[0-9a-zA-Z]+]]: memref<8x2xf8E8M0FNU>, %[[ARG4:[0-9a-zA-Z]+]]: memref<2x16xf8E8M0FNU>
// CHECK: %[[CST:.*]] = arith.constant {layout_result_0 = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>} dense<0.000000e+00> : vector<8x16xbf16>
-// CHECK: %[[T0:.*]] = xegpu.create_nd_tdesc %[[ARG0]] : memref<8x32xf8E5M2> -> !xegpu.tensor_desc<8x32xf8E5M2, #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 2]>>
-// CHECK: %[[T1:.*]] = xegpu.create_nd_tdesc %[[ARG1]] : memref<32x16xf8E5M2> -> !xegpu.tensor_desc<32x16xf8E5M2, #xegpu.layout<lane_layout = [1, 16], lane_data = [4, 1]>>
-// CHECK: %[[T2:.*]] = xegpu.load_nd %[[T0]][0, 0] <{layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 2]>}> :
-// CHECK-SAME: !xegpu.tensor_desc<8x32xf8E5M2, #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 2]>> -> vector<8x32xf8E5M2>
-// CHECK: %[[T3:.*]] = xegpu.load_nd %[[T1]][0, 0] <{layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [4, 1]>}> :
-// CHECK-SAME: !xegpu.tensor_desc<32x16xf8E5M2, #xegpu.layout<lane_layout = [1, 16], lane_data = [4, 1]>> -> vector<32x16xf8E5M2>
-// CHECK: %[[T4:.*]] = xegpu.create_nd_tdesc %[[ARG3]] : memref<8x1xf8E8M0FNU> -> !xegpu.tensor_desc<8x1xf8E8M0FNU, #xegpu.layout<lane_layout = [8, 1], lane_data = [1, 1]>>
+// CHECK: %[[T0:.*]] = xegpu.create_nd_tdesc %[[ARG0]] : memref<8x64xf4E2M1FN> -> !xegpu.tensor_desc<8x64xf4E2M1FN, #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 4]>>
+// CHECK: %[[T1:.*]] = xegpu.create_nd_tdesc %[[ARG1]] : memref<64x16xf4E2M1FN> -> !xegpu.tensor_desc<64x16xf4E2M1FN, #xegpu.layout<lane_layout = [1, 16], lane_data = [8, 1]>>
+// CHECK: %[[T2:.*]] = xegpu.load_nd %[[T0]][0, 0] <{layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 4]>}> :
+// CHECK-SAME: !xegpu.tensor_desc<8x64xf4E2M1FN, #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 4]>> -> vector<8x64xf4E2M1FN>
+// CHECK: %[[T3:.*]] = xegpu.load_nd %[[T1]][0, 0] <{layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [8, 1]>}> :
+// CHECK-SAME: !xegpu.tensor_desc<64x16xf4E2M1FN, #xegpu.layout<lane_layout = [1, 16], lane_data = [8, 1]>> -> vector<64x16xf4E2M1FN>
+// CHECK: %[[T4:.*]] = xegpu.create_nd_tdesc %[[ARG3]] : memref<8x2xf8E8M0FNU> -> !xegpu.tensor_desc<8x2xf8E8M0FNU, #xegpu.layout<lane_layout = [8, 1], lane_data = [1, 1]>>
// CHECK: %[[T5:.*]] = xegpu.load_nd %[[T4]][0, 0] <{layout = #xegpu.layout<lane_layout = [8, 1], lane_data = [1, 1]>}> :
-// CHECK-SAME: !xegpu.tensor_desc<8x1xf8E8M0FNU, #xegpu.layout<lane_layout = [8, 1], lane_data = [1, 1]>> -> vector<8x1xf8E8M0FNU>
-// CHECK: %[[T6:.*]] = xegpu.create_nd_tdesc %[[ARG4]] : memref<1x16xf8E8M0FNU> -> !xegpu.tensor_desc<1x16xf8E8M0FNU, #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>>
+// CHECK-SAME: !xegpu.tensor_desc<8x2xf8E8M0FNU, #xegpu.layout<lane_layout = [8, 1], lane_data = [1, 1]>> -> vector<8x2xf8E8M0FNU>
+// CHECK: %[[T6:.*]] = xegpu.create_nd_tdesc %[[ARG4]] : memref<2x16xf8E8M0FNU> -> !xegpu.tensor_desc<2x16xf8E8M0FNU, #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>>
// CHECK: %[[T7:.*]] = xegpu.load_nd %[[T6]][0, 0] <{layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>}> :
-// CHECK-SAME: !xegpu.tensor_desc<1x16xf8E8M0FNU, #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>> -> vector<1x16xf8E8M0FNU>
+// CHECK-SAME: !xegpu.tensor_desc<2x16xf8E8M0FNU, #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>> -> vector<2x16xf8E8M0FNU>
// CHECK: %[[T8:.*]] = xegpu.dpas_mx %[[T2]], %[[T3]], %[[CST]] scale_a = %[[T5]] scale_b = %[[T7]]
-// CHECK-SAME: {layout_a = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 2]>, layout_a_scale = #xegpu.layout<lane_layout = [8, 1], lane_data = [1, 1]>, layout_b = #xegpu.layout<lane_layout = [1, 16], lane_data = [4, 1]>, layout_b_scale = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>, layout_cd = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>} :
-// CHECK-SAME: (vector<8x32xf8E5M2>, vector<32x16xf8E5M2>, vector<8x16xbf16>, vector<8x1xf8E8M0FNU>, vector<1x16xf8E8M0FNU>) -> vector<8x16xbf16>
+// CHECK-SAME: {layout_a = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 4]>, layout_a_scale = #xegpu.layout<lane_layout = [8, 1], lane_data = [1, 1]>, layout_b = #xegpu.layout<lane_layout = [1, 16], lane_data = [8, 1]>, layout_b_scale = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>, layout_cd = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>} :
+// CHECK-SAME: (vector<8x64xf4E2M1FN>, vector<64x16xf4E2M1FN>, vector<8x16xbf16>, vector<8x2xf8E8M0FNU>, vector<2x16xf8E8M0FNU>) -> vector<8x16xbf16>
// CHECK: %[[T9:.*]] = xegpu.create_nd_tdesc %[[ARG2]] : memref<8x16xbf16> -> !xegpu.tensor_desc<8x16xbf16, #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>>
// CHECK: xegpu.store_nd %[[T8]], %[[T9]][0, 0] <{layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>}> : vector<8x16xbf16>, !xegpu.tensor_desc<8x16xbf16, #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>>
-func.func @dpas_mx_f8e5m2(%arg0: memref<8x32xf8E5M2>, %arg1: memref<32x16xf8E5M2>, %arg2: memref<8x16xbf16>,
- %arg3: memref<8x1xf8E8M0FNU>, %arg4: memref<1x16xf8E8M0FNU>) {
+func.func @dpas_mx_fp4(%arg0: memref<8x64xf4E2M1FN>, %arg1: memref<64x16xf4E2M1FN>, %arg2: memref<8x16xbf16>,
+ %arg3: memref<8x2xf8E8M0FNU>, %arg4: memref<2x16xf8E8M0FNU>) {
%c0 = arith.constant 0 : index
%cst = arith.constant dense<0.000000e+00> : vector<8x16xbf16>
- %0 = xegpu.create_nd_tdesc %arg0 : memref<8x32xf8E5M2> -> !xegpu.tensor_desc<8x32xf8E5M2>
- %1 = xegpu.create_nd_tdesc %arg1 : memref<32x16xf8E5M2> -> !xegpu.tensor_desc<32x16xf8E5M2>
- %2 = xegpu.load_nd %0[0, 0] : !xegpu.tensor_desc<8x32xf8E5M2> -> vector<8x32xf8E5M2>
- %3 = xegpu.load_nd %1[0, 0] : !xegpu.tensor_desc<32x16xf8E5M2> -> vector<32x16xf8E5M2>
- %4 = xegpu.create_nd_tdesc %arg3 : memref<8x1xf8E8M0FNU> -> !xegpu.tensor_desc<8x1xf8E8M0FNU>
- %5 = xegpu.load_nd %4[0, 0] : !xegpu.tensor_desc<8x1xf8E8M0FNU> -> vector<8x1xf8E8M0FNU>
- %6 = xegpu.create_nd_tdesc %arg4 : memref<1x16xf8E8M0FNU> -> !xegpu.tensor_desc<1x16xf8E8M0FNU>
- %7 = xegpu.load_nd %6[0, 0] : !xegpu.tensor_desc<1x16xf8E8M0FNU> -> vector<1x16xf8E8M0FNU>
- %8 = xegpu.dpas_mx %2, %3, %cst scale_a = %5 scale_b = %7 : (vector<8x32xf8E5M2>, vector<32x16xf8E5M2>, vector<8x16xbf16>, vector<8x1xf8E8M0FNU>, vector<1x16xf8E8M0FNU>) -> vector<8x16xbf16>
+ %0 = xegpu.create_nd_tdesc %arg0 : memref<8x64xf4E2M1FN> -> !xegpu.tensor_desc<8x64xf4E2M1FN>
+ %1 = xegpu.create_nd_tdesc %arg1 : memref<64x16xf4E2M1FN> -> !xegpu.tensor_desc<64x16xf4E2M1FN>
+ %2 = xegpu.load_nd %0[0, 0] : !xegpu.tensor_desc<8x64xf4E2M1FN> -> vector<8x64xf4E2M1FN>
+ %3 = xegpu.load_nd %1[0, 0] : !xegpu.tensor_desc<64x16xf4E2M1FN> -> vector<64x16xf4E2M1FN>
+ %4 = xegpu.create_nd_tdesc %arg3 : memref<8x2xf8E8M0FNU> -> !xegpu.tensor_desc<8x2xf8E8M0FNU>
+ %5 = xegpu.load_nd %4[0, 0] : !xegpu.tensor_desc<8x2xf8E8M0FNU> -> vector<8x2xf8E8M0FNU>
+ %6 = xegpu.create_nd_tdesc %arg4 : memref<2x16xf8E8M0FNU> -> !xegpu.tensor_desc<2x16xf8E8M0FNU>
+ %7 = xegpu.load_nd %6[0, 0] : !xegpu.tensor_desc<2x16xf8E8M0FNU> -> vector<2x16xf8E8M0FNU>
+ %8 = xegpu.dpas_mx %2, %3, %cst scale_a = %5 scale_b = %7 : (vector<8x64xf4E2M1FN>, vector<64x16xf4E2M1FN>, vector<8x16xbf16>, vector<8x2xf8E8M0FNU>, vector<2x16xf8E8M0FNU>) -> vector<8x16xbf16>
%9 = xegpu.create_nd_tdesc %arg2 : memref<8x16xbf16> -> !xegpu.tensor_desc<8x16xbf16>
xegpu.store_nd %8, %9[0, 0] : vector<8x16xbf16>, !xegpu.tensor_desc<8x16xbf16>
return
@@ -1035,14 +1029,14 @@ gpu.module @test {
// CHECK-SAME: !xegpu.tensor_desc<8x64xf4E2M1FN, #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 4]>> -> vector<8x64xf4E2M1FN>
// CHECK: %[[T3:.*]] = xegpu.load_nd %[[T1]][0, 0] <{layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [8, 1]>}> :
// CHECK-SAME: !xegpu.tensor_desc<64x16xf4E2M1FN, #xegpu.layout<lane_layout = [1, 16], lane_data = [8, 1]>> -> vector<64x16xf4E2M1FN>
-// CHECK: %[[T4:.*]] = xegpu.create_nd_tdesc %[[ARG3]] : memref<8x2xf8E8M0FNU> -> !xegpu.tensor_desc<8x2xf8E8M0FNU, #xegpu.layout<lane_layout = [8, 1], lane_data = [1, 2]>>
-// CHECK: %[[T5:.*]] = xegpu.load_nd %[[T4]][0, 0] <{layout = #xegpu.layout<lane_layout = [8, 1], lane_data = [1, 2]>}> :
-// CHECK-SAME: !xegpu.tensor_desc<8x2xf8E8M0FNU, #xegpu.layout<lane_layout = [8, 1], lane_data = [1, 2]>> -> vector<8x2xf8E8M0FNU>
-// CHECK: %[[T6:.*]] = xegpu.create_nd_tdesc %[[ARG4]] : memref<2x16xf8E8M0FNU> -> !xegpu.tensor_desc<2x16xf8E8M0FNU, #xegpu.layout<lane_layout = [1, 16], lane_data = [2, 1]>>
-// CHECK: %[[T7:.*]] = xegpu.load_nd %[[T6]][0, 0] <{layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [2, 1]>}> :
-// CHECK-SAME: !xegpu.tensor_desc<2x16xf8E8M0FNU, #xegpu.layout<lane_layout = [1, 16], lane_data = [2, 1]>> -> vector<2x16xf8E8M0FNU>
+// CHECK: %[[T4:.*]] = xegpu.create_nd_tdesc %[[ARG3]] : memref<8x2xf8E8M0FNU> -> !xegpu.tensor_desc<8x2xf8E8M0FNU, #xegpu.layout<lane_layout = [8, 1], lane_data = [1, 1]>>
+// CHECK: %[[T5:.*]] = xegpu.load_nd %[[T4]][0, 0] <{layout = #xegpu.layout<lane_layout = [8, 1], lane_data = [1, 1]>}> :
+// CHECK-SAME: !xegpu.tensor_desc<8x2xf8E8M0FNU, #xegpu.layout<lane_layout = [8, 1], lane_data = [1, 1]>> -> vector<8x2xf8E8M0FNU>
+// CHECK: %[[T6:.*]] = xegpu.create_nd_tdesc %[[ARG4]] : memref<2x16xf8E8M0FNU> -> !xegpu.tensor_desc<2x16xf8E8M0FNU, #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>>
+// CHECK: %[[T7:.*]] = xegpu.load_nd %[[T6]][0, 0] <{layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>}> :
+// CHECK-SAME: !xegpu.tensor_desc<2x16xf8E8M0FNU, #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>> -> vector<2x16xf8E8M0FNU>
// CHECK: %[[T8:.*]] = xegpu.dpas_mx %[[T2]], %[[T3]], %[[CST]] scale_a = %[[T5]] scale_b = %[[T7]]
-// CHECK-SAME: {layout_a = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 4]>, layout_a_scale = #xegpu.layout<lane_layout = [8, 1], lane_data = [1, 2]>, layout_b = #xegpu.layout<lane_layout = [1, 16], lane_data = [8, 1]>, layout_b_scale = #xegpu.layout<lane_layout = [1, 16], lane_data = [2, 1]>, layout_cd = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>} :
+// CHECK-SAME: {layout_a = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 4]>, layout_a_scale = #xegpu.layout<lane_layout = [8, 1], lane_data = [1, 1]>, layout_b = #xegpu.layout<lane_layout = [1, 16], lane_data = [8, 1]>, layout_b_scale = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>, layout_cd = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>} :
// CHECK-SAME: (vector<8x64xf4E2M1FN>, vector<64x16xf4E2M1FN>, vector<8x16xbf16>, vector<8x2xf8E8M0FNU>, vector<2x16xf8E8M0FNU>) -> vector<8x16xbf16>
// CHECK: %[[T9:.*]] = xegpu.create_nd_tdesc %[[ARG2]] : memref<8x16xbf16> -> !xegpu.tensor_desc<8x16xbf16, #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>>
// CHECK: xegpu.store_nd %[[T8]], %[[T9]][0, 0] <{layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>}> : vector<8x16xbf16>, !xegpu.tensor_desc<8x16xbf16, #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>>
diff --git a/mlir/test/Dialect/XeGPU/resolve-layout-conflicts.mlir b/mlir/test/Dialect/XeGPU/resolve-layout-conflicts.mlir
index 40e907be6e4a1..55d0e64bb2c65 100644
--- a/mlir/test/Dialect/XeGPU/resolve-layout-conflicts.mlir
+++ b/mlir/test/Dialect/XeGPU/resolve-layout-conflicts.mlir
@@ -254,8 +254,7 @@ func.func @conflict_postop() {
}
// CHECK-LABEL: func.func @convert_layout
-// CHECK: %[[V0:.*]] = xegpu.convert_layout %[[CST:.*]] <{input_layout = #xegpu.layout<sg_layout = [8, 4], sg_data = [4, 32]>, target_layout = #xegpu.layout<sg_layout = [8, 4], sg_data = [4, 32]>}> : vector<32x128xf32>
-// CHECK: %[[V1:.*]] = xegpu.convert_layout %[[V0]] <{input_layout = #xegpu.layout<sg_layout = [8, 4], sg_data = [4, 32]>, target_layout = #xegpu.layout<sg_layout = [4, 8], sg_data = [8, 16]>}> : vector<32x128xf32>
+// CHECK: %[[V1:.*]] = xegpu.convert_layout %[[V0:.*]] <{input_layout = #xegpu.layout<sg_layout = [8, 4], sg_data = [4, 32]>, target_layout = #xegpu.layout<sg_layout = [4, 8], sg_data = [8, 16]>}> : vector<32x128xf32>
func.func @convert_layout() {
%src0 = arith.constant
{layout_result_0 = #xegpu.layout<sg_layout=[8, 4], sg_data=[4, 32]>}
@@ -354,10 +353,7 @@ func.func @extract_source_conflict_with_order() -> vector<16x32xf16> {
// CHECK-LABEL: func.func @convert_layout_bridge_input_mismatch
// CHECK: %[[V0:.*]] = "some_op"() {layout_result_0 = #xegpu.layout<inst_data = [8, 16]>} : () -> vector<32x32xf16>
// CHECK-NEXT: %[[BRIDGE:.*]] = xegpu.convert_layout %[[V0]]
-// CHECK-SAME: <{input_layout = #xegpu.layout<inst_data = [8, 16]>, target_layout = #xegpu.layout<inst_data = [16, 16]>}>
-// CHECK-SAME: : vector<32x32xf16>
-// CHECK-NEXT: %[[CVT:.*]] = xegpu.convert_layout %[[BRIDGE]]
-// CHECK-SAME: <{input_layout = #xegpu.layout<inst_data = [16, 16]>, target_layout = #xegpu.layout<inst_data = [32, 16]>}>
+// CHECK-SAME: <{input_layout = #xegpu.layout<inst_data = [8, 16]>, target_layout = #xegpu.layout<inst_data = [32, 16]>}>
// CHECK-SAME: : vector<32x32xf16>
gpu.module @test_convert_layout_bridge {
func.func @convert_layout_bridge_input_mismatch() {
More information about the Mlir-commits
mailing list