[Mlir-commits] [mlir] 4226250 - [MLIR][XeGPU] Fix matrix ops layout propagation (#182268)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Sun Feb 22 04:40:53 PST 2026
Author: Artem Kroviakov
Date: 2026-02-22T13:40:48+01:00
New Revision: 4226250a421b6dbb734d4358fcb36a97da2278d7
URL: https://github.com/llvm/llvm-project/commit/4226250a421b6dbb734d4358fcb36a97da2278d7
DIFF: https://github.com/llvm/llvm-project/commit/4226250a421b6dbb734d4358fcb36a97da2278d7.diff
LOG: [MLIR][XeGPU] Fix matrix ops layout propagation (#182268)
Added:
Modified:
mlir/include/mlir/Dialect/XeGPU/uArch/IntelGpuXe2.h
mlir/include/mlir/Dialect/XeGPU/uArch/uArchBase.h
mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
mlir/test/Dialect/XeGPU/propagate-layout-subgroup.mlir
Removed:
################################################################################
diff --git a/mlir/include/mlir/Dialect/XeGPU/uArch/IntelGpuXe2.h b/mlir/include/mlir/Dialect/XeGPU/uArch/IntelGpuXe2.h
index c60a744d6b858..721cb74823718 100644
--- a/mlir/include/mlir/Dialect/XeGPU/uArch/IntelGpuXe2.h
+++ b/mlir/include/mlir/Dialect/XeGPU/uArch/IntelGpuXe2.h
@@ -223,14 +223,6 @@ struct SpirvStoreScatterInstruction : public StoreScatterInstructionInterface {
int32_t getMaxLaneStoreSize(int32_t bitWidth) const override { return 16; }
};
-struct LoadMatrixInstruction : public LoadMatrixInstructionInterface {
- int32_t getMaxLaneLoadSize(int32_t bitWidth) const override { return 16; }
-};
-
-struct StoreMatrixInstruction : public StoreMatrixInstructionInterface {
- int32_t getMaxLaneStoreSize(int32_t bitWidth) const override { return 16; }
-};
-
//===----------------------------------------------------------------------===//
// uArch instances
//===----------------------------------------------------------------------===//
@@ -243,11 +235,9 @@ struct PVCuArch final : public Xe2Plus {
static const Subgroup2DBlockPrefetchInstruction prefetchNdInst;
static const SpirvStoreScatterInstruction storeScatterInst;
static const SpirvLoadGatherInstruction loadGatherInst;
- static const StoreMatrixInstruction storeMatrixInst;
- static const LoadMatrixInstruction loadMatrixInst;
- static const Instruction *arr[] = {
- &dpasInst, &loadNdInst, &storeNdInst, &prefetchNdInst,
- &storeScatterInst, &loadGatherInst, &storeMatrixInst, &loadMatrixInst};
+ static const Instruction *arr[] = {&dpasInst, &loadNdInst,
+ &storeNdInst, &prefetchNdInst,
+ &storeScatterInst, &loadGatherInst};
return arr;
}
@@ -271,11 +261,9 @@ struct BMGuArch : public Xe2Plus {
static const Subgroup2DBlockPrefetchInstruction prefetchNdInst;
static const SpirvStoreScatterInstruction storeScatterInst;
static const SpirvLoadGatherInstruction loadGatherInst;
- static const StoreMatrixInstruction storeMatrixInst;
- static const LoadMatrixInstruction loadMatrixInst;
- static const Instruction *arr[] = {
- &dpasInst, &loadNdInst, &storeNdInst, &prefetchNdInst,
- &storeScatterInst, &loadGatherInst, &storeMatrixInst, &loadMatrixInst};
+ static const Instruction *arr[] = {&dpasInst, &loadNdInst,
+ &storeNdInst, &prefetchNdInst,
+ &storeScatterInst, &loadGatherInst};
return arr;
}
diff --git a/mlir/include/mlir/Dialect/XeGPU/uArch/uArchBase.h b/mlir/include/mlir/Dialect/XeGPU/uArch/uArchBase.h
index 0c8673e602c46..0f9d052e11147 100644
--- a/mlir/include/mlir/Dialect/XeGPU/uArch/uArchBase.h
+++ b/mlir/include/mlir/Dialect/XeGPU/uArch/uArchBase.h
@@ -41,8 +41,6 @@ enum class InstructionKind {
Subgroup2DBlockPrefetch, // Subgroup-level 2D block prefetch instruction
StoreScatter, // Lane-level store (scalar, vector)
LoadGather, // Lane-level load (scalar, vector)
- StoreMatrix, // Lane-level matrix store to slm
- LoadMatrix // Lane-level matrix load to slm
// @TODO: Add more instructions as needed
};
@@ -73,10 +71,6 @@ struct Instruction {
return "store";
case InstructionKind::LoadGather:
return "load";
- case InstructionKind::StoreMatrix:
- return "store_matrix";
- case InstructionKind::LoadMatrix:
- return "load_matrix";
}
llvm_unreachable("Unknown InstructionKind");
}
@@ -282,28 +276,6 @@ struct StoreScatterInstructionInterface : public Instruction {
virtual ~StoreScatterInstructionInterface() = default;
};
-struct LoadMatrixInstructionInterface : public Instruction {
- LoadMatrixInstructionInterface()
- : Instruction(InstructionKind::LoadMatrix, InstructionScope::Lane) {}
- static bool classof(const Instruction *B) {
- return B->getInstructionKind() == InstructionKind::LoadMatrix;
- }
-
- virtual int32_t getMaxLaneLoadSize(int32_t bitWidth) const = 0;
- virtual ~LoadMatrixInstructionInterface() = default;
-};
-
-struct StoreMatrixInstructionInterface : public Instruction {
- StoreMatrixInstructionInterface()
- : Instruction(InstructionKind::StoreMatrix, InstructionScope::Lane) {}
- static bool classof(const Instruction *B) {
- return B->getInstructionKind() == InstructionKind::StoreMatrix;
- }
-
- virtual int32_t getMaxLaneStoreSize(int32_t bitWidth) const = 0;
- virtual ~StoreMatrixInstructionInterface() = default;
-};
-
} // namespace uArch
} // namespace xegpu
} // namespace mlir
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
index eb7fab3610218..aa4195c2ca279 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
@@ -735,7 +735,7 @@ xegpu::DistributeLayoutAttr xegpu::setupLoadGatherAnchorLayout(
auto elemBitWidth = resVecTy.getElementType().getIntOrFloatBitWidth();
const auto *uArchInstruction =
- dyn_cast<xegpu::uArch::SpirvLoadGatherInstruction>(
+ dyn_cast<xegpu::uArch::LoadGatherInstructionInterface>(
uArch->getInstruction(xegpu::uArch::InstructionKind::LoadGather));
int maxChunkSize = uArchInstruction->getMaxLaneLoadSize(elemBitWidth);
@@ -757,8 +757,9 @@ xegpu::setupLoadMatrixAnchorLayout(xegpu::LayoutKind layoutKind,
auto context = resVecTy.getContext();
auto elemBitWidth = resVecTy.getElementType().getIntOrFloatBitWidth();
- const auto *uArchInstruction = dyn_cast<xegpu::uArch::LoadMatrixInstruction>(
- uArch->getInstruction(xegpu::uArch::InstructionKind::LoadMatrix));
+ const auto *uArchInstruction =
+ dyn_cast<xegpu::uArch::LoadGatherInstructionInterface>(
+ uArch->getInstruction(xegpu::uArch::InstructionKind::LoadGather));
int maxChunkSize = uArchInstruction->getMaxLaneLoadSize(elemBitWidth);
return setupGenericLoadAnchorLayout(layoutKind, context, consumerLayout,
false, maxChunkSize, resShapeSize,
@@ -828,7 +829,7 @@ xegpu::setupStoreScatterAnchorLayout(xegpu::LayoutKind layoutKind,
auto elemBitWidth = srcVecTy.getElementType().getIntOrFloatBitWidth();
const auto *uArchInstruction =
- dyn_cast<xegpu::uArch::SpirvStoreScatterInstruction>(
+ dyn_cast<xegpu::uArch::StoreScatterInstructionInterface>(
uArch->getInstruction(xegpu::uArch::InstructionKind::StoreScatter));
int maxChunkSize = uArchInstruction->getMaxLaneStoreSize(elemBitWidth);
return setupGenericStoreAnchorLayout(layoutKind, context, (chunkSize > 1),
@@ -846,8 +847,9 @@ xegpu::setupStoreMatrixAnchorLayout(xegpu::LayoutKind layoutKind,
auto context = srcVecTy.getContext();
auto elemBitWidth = srcVecTy.getElementType().getIntOrFloatBitWidth();
- const auto *uArchInstruction = dyn_cast<xegpu::uArch::StoreMatrixInstruction>(
- uArch->getInstruction(xegpu::uArch::InstructionKind::StoreMatrix));
+ const auto *uArchInstruction =
+ dyn_cast<xegpu::uArch::StoreScatterInstructionInterface>(
+ uArch->getInstruction(xegpu::uArch::InstructionKind::StoreScatter));
int maxChunkSize = uArchInstruction->getMaxLaneStoreSize(elemBitWidth);
return setupGenericStoreAnchorLayout(layoutKind, context, false, maxChunkSize,
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
index bc309c9029878..9b4ecad409815 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
@@ -1102,6 +1102,9 @@ void LayoutInfoPropagation::visitLoadMatrixOp(
ArrayRef<const LayoutInfoLattice *> results) {
LayoutInfo resLayoutInfo = results[0]->getValue();
+ if (!resLayoutInfo.isAssigned())
+ return;
+
auto consumerLayoutAttr =
dyn_cast<xegpu::DistributeLayoutAttr>(resLayoutInfo.get());
diff --git a/mlir/test/Dialect/XeGPU/propagate-layout-subgroup.mlir b/mlir/test/Dialect/XeGPU/propagate-layout-subgroup.mlir
index 6e6ed01845ff1..c073045691f56 100644
--- a/mlir/test/Dialect/XeGPU/propagate-layout-subgroup.mlir
+++ b/mlir/test/Dialect/XeGPU/propagate-layout-subgroup.mlir
@@ -254,3 +254,18 @@ gpu.module @test {
gpu.return
}
}
+
+// -----
+gpu.module @xevm_module{
+ // CHECK-LABEL: load_store_matrix
+ gpu.func @load_store_matrix(%arg0: !xegpu.mem_desc<64x128xf32>, %sg_id_lt_2: i1) {
+ %c0 = arith.constant 0 : index
+ scf.if %sg_id_lt_2 {
+ // CHECK: xegpu.load_matrix %{{.*}} <{layout = #xegpu.layout<sg_layout = [4, 2], sg_data = [8, 16]>}>
+ %1 = xegpu.load_matrix %arg0[%c0, %c0] : !xegpu.mem_desc<64x128xf32>, index, index -> vector<32x32xf32>
+ // CHECK: xegpu.store_matrix %{{.*}} <{layout = #xegpu.layout<sg_layout = [4, 2], sg_data = [8, 16]>}>
+ xegpu.store_matrix %1, %arg0[%c0, %c0] <{layout = #xegpu.layout<sg_layout = [4, 2], sg_data = [8, 16]>}> : vector<32x32xf32>, !xegpu.mem_desc<64x128xf32>, index, index
+ }
+ gpu.return
+ }
+}
More information about the Mlir-commits
mailing list