[Mlir-commits] [mlir] [MLIR][XeGPU] Enable `isa<>` check for uarch (PR #204577)
Artem Kroviakov
llvmlistbot at llvm.org
Thu Jun 25 03:07:05 PDT 2026
https://github.com/akroviakov updated https://github.com/llvm/llvm-project/pull/204577
>From a615e5f8fcbba5cf88afe4b633abef4972535832 Mon Sep 17 00:00:00 2001
From: Artem Kroviakov <artem.kroviakov at intel.com>
Date: Thu, 18 Jun 2026 12:28:30 +0000
Subject: [PATCH 1/4] [MLIR][XeGPU] Enable isa<> check for uarch
---
.../mlir/Dialect/XeGPU/uArch/IntelGpuXe2.h | 26 +++++-----
.../mlir/Dialect/XeGPU/uArch/uArchBase.h | 47 ++++++++++++++++---
mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp | 4 +-
3 files changed, 54 insertions(+), 23 deletions(-)
diff --git a/mlir/include/mlir/Dialect/XeGPU/uArch/IntelGpuXe2.h b/mlir/include/mlir/Dialect/XeGPU/uArch/IntelGpuXe2.h
index ff80a77b28d37..46747af25c38e 100644
--- a/mlir/include/mlir/Dialect/XeGPU/uArch/IntelGpuXe2.h
+++ b/mlir/include/mlir/Dialect/XeGPU/uArch/IntelGpuXe2.h
@@ -31,13 +31,17 @@ namespace xegpu {
namespace uArch {
struct Xe2Plus : public uArch {
- Xe2Plus(StringRef archName, StringRef archDescription,
- llvm::ArrayRef<const Instruction *> instructionRegistry,
+ Xe2Plus(Kind kind, llvm::ArrayRef<const Instruction *> instructionRegistry,
const XeCoreInfo &xeCore)
- : uArch(archName, archDescription, instructionRegistry), xeCore(xeCore) {}
+ : uArch(kind, instructionRegistry), xeCore(xeCore) {}
int getSubgroupSize() const override { return 16; }
unsigned getGeneralPackedFormatBitSize() const override { return 32; }
+ static bool classof(const uArch *u) {
+ return u->getKind() >= Kind::Xe2Plus_First &&
+ u->getKind() <= Kind::Xe2Plus_Last;
+ }
+
protected:
XeCoreInfo xeCore;
};
@@ -315,11 +319,10 @@ struct PVCuArch final : public Xe2Plus {
}
PVCuArch()
- : Xe2Plus("pvc", // archName
- "Ponte Vecchio Architecture", // archDescription
- getInstructionRegistryArr(),
+ : Xe2Plus(Kind::PVC, getInstructionRegistryArr(),
XeCoreInfo(8, SharedMemory(512 * 1024, 4), 8, 8) // xeCore
) {}
+ static bool classof(const uArch *u) { return u->getKind() == Kind::PVC; }
static const uArch *getInstance() {
static const PVCuArch instance;
return reinterpret_cast<const uArch *>(&instance);
@@ -341,11 +344,10 @@ struct BMGuArch : public Xe2Plus {
}
BMGuArch()
- : Xe2Plus("bmg", // archName
- "Battlemage Architecture", // archDescription
- getInstructionRegistryArr(),
+ : Xe2Plus(Kind::BMG, getInstructionRegistryArr(),
XeCoreInfo(8, SharedMemory(256 * 1024, 4), 8, 8) // xeCore
) {}
+ static bool classof(const uArch *u) { return u->getKind() == Kind::BMG; }
static const uArch *getInstance() {
static const BMGuArch instance;
return reinterpret_cast<const uArch *>(&instance);
@@ -368,13 +370,12 @@ struct CRIuArch : public Xe2Plus {
}
CRIuArch()
- : Xe2Plus("cri", // archName
- "Crescent Island Architecture", // archDescription
- getInstructionRegistryArr(),
+ : Xe2Plus(Kind::CRI, getInstructionRegistryArr(),
// Using bmg config as placeholder
// TODO: Update to actual XeCore and SharedMemory config
XeCoreInfo(8, SharedMemory(256 * 1024, 4), 8, 8) // xeCore
) {}
+ static bool classof(const uArch *u) { return u->getKind() == Kind::CRI; }
static const uArch *getInstance() {
static const CRIuArch instance;
return reinterpret_cast<const uArch *>(&instance);
@@ -390,7 +391,6 @@ inline const uArch *getUArch(llvm::StringRef archName) {
return CRIuArch::getInstance();
return nullptr;
}
-
} // namespace uArch
} // namespace xegpu
} // namespace mlir
diff --git a/mlir/include/mlir/Dialect/XeGPU/uArch/uArchBase.h b/mlir/include/mlir/Dialect/XeGPU/uArch/uArchBase.h
index 61db4605e85fa..1702c7b46e0eb 100644
--- a/mlir/include/mlir/Dialect/XeGPU/uArch/uArchBase.h
+++ b/mlir/include/mlir/Dialect/XeGPU/uArch/uArchBase.h
@@ -25,6 +25,7 @@
#include "mlir/IR/Types.h"
#include "llvm/ADT/SmallVector.h"
+#include "llvm/Support/Casting.h"
namespace mlir {
namespace xegpu {
@@ -153,16 +154,49 @@ struct CacheInfo {
};
struct uArch {
+ enum class Kind {
+ // Xe2 family
+ Xe2Plus_First,
+ PVC = Xe2Plus_First,
+ BMG,
+ CRI,
+ Xe2Plus_Last = CRI,
+ };
+
// Constructor
- uArch(StringRef name, StringRef description,
- llvm::ArrayRef<const Instruction *> instructionRegistry)
- : name(name), description(description) {
+ uArch(Kind kind, llvm::ArrayRef<const Instruction *> instructionRegistry)
+ : kind(kind) {
for (const Instruction *instr : instructionRegistry)
this->instructionRegistry[instr->getInstructionKind()] = instr;
}
virtual ~uArch() = default;
- StringRef getName() const { return name; }
- StringRef getDescription() const { return description; }
+ Kind getKind() const { return kind; }
+ StringRef getName() const { return getUArchName(kind); }
+ StringRef getDescription() const { return getUArchDescription(kind); }
+
+ static StringRef getUArchName(Kind k) {
+ switch (k) {
+ case Kind::PVC:
+ return "pvc";
+ case Kind::BMG:
+ return "bmg";
+ case Kind::CRI:
+ return "cri";
+ }
+ llvm_unreachable("Unknown uArch::Kind");
+ }
+
+ static StringRef getUArchDescription(Kind k) {
+ switch (k) {
+ case Kind::PVC:
+ return "Ponte Vecchio Architecture";
+ case Kind::BMG:
+ return "Battlemage Architecture";
+ case Kind::CRI:
+ return "Crescent Island Architecture";
+ }
+ llvm_unreachable("Unknown uArch::Kind");
+ }
virtual int getSubgroupSize() const = 0;
virtual unsigned getGeneralPackedFormatBitSize() const = 0;
@@ -178,8 +212,7 @@ struct uArch {
}
protected:
- StringRef name;
- StringRef description;
+ Kind kind;
llvm::SmallDenseMap<InstructionKind, const Instruction *, 32>
instructionRegistry;
};
diff --git a/mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp b/mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp
index 735dd93e2ac14..24999f85c1268 100644
--- a/mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp
+++ b/mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp
@@ -786,9 +786,7 @@ bool xegpu::requireTranspose(const xegpu::DistributeLayoutAttr layout,
const xegpu::uArch::uArch *uArch) {
// Return false for unsupported targets.
// TODO: Add more support or move to target info.
- if (uArch->getName().equals_insensitive("pvc") &&
- uArch->getName().equals_insensitive("bmg") &&
- uArch->getName().equals_insensitive("cri"))
+ if (!isa<xegpu::uArch::Xe2Plus>(uArch))
return false;
if (!layout)
return false;
>From 6101a6fd410ccb3057041c83bfb44cf1e90d989e Mon Sep 17 00:00:00 2001
From: Artem Kroviakov <artem.kroviakov at intel.com>
Date: Thu, 18 Jun 2026 16:49:42 +0000
Subject: [PATCH 2/4] Split to Xe2 and Xe3
---
.../XeGPU/Transforms/XeGPULayoutImpl.h | 2 +-
.../mlir/Dialect/XeGPU/uArch/IntelGpuXe2.h | 66 +-
.../mlir/Dialect/XeGPU/uArch/IntelGpuXe3.h | 603 ++++++++++++++++++
.../mlir/Dialect/XeGPU/uArch/uArchBase.h | 20 +-
.../mlir/Dialect/XeGPU/uArch/uArchCommon.h | 38 ++
.../Conversion/XeGPUToXeVM/XeGPUToXeVM.cpp | 2 +-
mlir/lib/Dialect/XeGPU/IR/XeGPUDialect.cpp | 2 +-
.../XeGPUArrayLengthOptimization.cpp | 2 +-
.../Transforms/XeGPUPeepHoleOptimizer.cpp | 2 +-
.../XeGPU/Transforms/XeGPUPropagateLayout.cpp | 2 +-
.../Transforms/XeGPUSgToLaneDistribute.cpp | 2 +-
mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp | 4 +-
12 files changed, 676 insertions(+), 69 deletions(-)
create mode 100644 mlir/include/mlir/Dialect/XeGPU/uArch/IntelGpuXe3.h
create mode 100644 mlir/include/mlir/Dialect/XeGPU/uArch/uArchCommon.h
diff --git a/mlir/include/mlir/Dialect/XeGPU/Transforms/XeGPULayoutImpl.h b/mlir/include/mlir/Dialect/XeGPU/Transforms/XeGPULayoutImpl.h
index bf61f0ced6048..027a70d9579ab 100644
--- a/mlir/include/mlir/Dialect/XeGPU/Transforms/XeGPULayoutImpl.h
+++ b/mlir/include/mlir/Dialect/XeGPU/Transforms/XeGPULayoutImpl.h
@@ -11,7 +11,7 @@
#include "mlir/Dialect/XeGPU/IR/XeGPU.h"
#include "mlir/Dialect/XeGPU/Utils/XeGPUUtils.h"
-#include "mlir/Dialect/XeGPU/uArch/IntelGpuXe2.h"
+#include "mlir/Dialect/XeGPU/uArch/uArchCommon.h"
#include "mlir/IR/BuiltinTypes.h"
#include "mlir/IR/OpDefinition.h"
#include "mlir/Interfaces/ControlFlowInterfaces.h"
diff --git a/mlir/include/mlir/Dialect/XeGPU/uArch/IntelGpuXe2.h b/mlir/include/mlir/Dialect/XeGPU/uArch/IntelGpuXe2.h
index 46747af25c38e..ed8ae5289882a 100644
--- a/mlir/include/mlir/Dialect/XeGPU/uArch/IntelGpuXe2.h
+++ b/mlir/include/mlir/Dialect/XeGPU/uArch/IntelGpuXe2.h
@@ -30,16 +30,15 @@ namespace mlir {
namespace xegpu {
namespace uArch {
-struct Xe2Plus : public uArch {
- Xe2Plus(Kind kind, llvm::ArrayRef<const Instruction *> instructionRegistry,
- const XeCoreInfo &xeCore)
+struct Xe2 : public uArch {
+ Xe2(Kind kind, llvm::ArrayRef<const Instruction *> instructionRegistry,
+ const XeCoreInfo &xeCore)
: uArch(kind, instructionRegistry), xeCore(xeCore) {}
int getSubgroupSize() const override { return 16; }
unsigned getGeneralPackedFormatBitSize() const override { return 32; }
static bool classof(const uArch *u) {
- return u->getKind() >= Kind::Xe2Plus_First &&
- u->getKind() <= Kind::Xe2Plus_Last;
+ return u->getKind() >= Kind::Xe2_First && u->getKind() <= Kind::Xe2_Last;
}
protected:
@@ -292,19 +291,11 @@ struct SubgroupScaledMatrixMultiplyAcc : public Instruction,
const unsigned packedFormatBitSizeB;
};
-struct SpirvLoadGatherInstruction : public LoadGatherInstructionInterface {
- int32_t getMaxLaneLoadSize(int32_t bitWidth) const override { return 16; }
-};
-
-struct SpirvStoreScatterInstruction : public StoreScatterInstructionInterface {
- int32_t getMaxLaneStoreSize(int32_t bitWidth) const override { return 16; }
-};
-
//===----------------------------------------------------------------------===//
// uArch instances
//===----------------------------------------------------------------------===//
-struct PVCuArch final : public Xe2Plus {
+struct PVCuArch final : public Xe2 {
static llvm::ArrayRef<const Instruction *> getInstructionRegistryArr() {
static const SubgroupMatrixMultiplyAcc dpasInst{16, 32};
static const Subgroup2DBlockLoadInstruction loadNdInst;
@@ -319,8 +310,8 @@ struct PVCuArch final : public Xe2Plus {
}
PVCuArch()
- : Xe2Plus(Kind::PVC, getInstructionRegistryArr(),
- XeCoreInfo(8, SharedMemory(512 * 1024, 4), 8, 8) // xeCore
+ : Xe2(Kind::PVC, getInstructionRegistryArr(),
+ XeCoreInfo(8, SharedMemory(512 * 1024, 4), 8, 8) // xeCore
) {}
static bool classof(const uArch *u) { return u->getKind() == Kind::PVC; }
static const uArch *getInstance() {
@@ -329,7 +320,7 @@ struct PVCuArch final : public Xe2Plus {
}
};
-struct BMGuArch : public Xe2Plus {
+struct BMGuArch : public Xe2 {
static llvm::ArrayRef<const Instruction *> getInstructionRegistryArr() {
static const SubgroupMatrixMultiplyAcc dpasInst{16, 32};
static const Subgroup2DBlockLoadInstruction loadNdInst;
@@ -344,8 +335,8 @@ struct BMGuArch : public Xe2Plus {
}
BMGuArch()
- : Xe2Plus(Kind::BMG, getInstructionRegistryArr(),
- XeCoreInfo(8, SharedMemory(256 * 1024, 4), 8, 8) // xeCore
+ : Xe2(Kind::BMG, getInstructionRegistryArr(),
+ XeCoreInfo(8, SharedMemory(256 * 1024, 4), 8, 8) // xeCore
) {}
static bool classof(const uArch *u) { return u->getKind() == Kind::BMG; }
static const uArch *getInstance() {
@@ -354,43 +345,6 @@ struct BMGuArch : public Xe2Plus {
}
};
-struct CRIuArch : public Xe2Plus {
- static llvm::ArrayRef<const Instruction *> getInstructionRegistryArr() {
- static const SubgroupMatrixMultiplyAcc dpasInst{16, 32};
- static const SubgroupScaledMatrixMultiplyAcc dpasMxInst{16, 32};
- static const Subgroup2DBlockLoadInstruction loadNdInst;
- static const Subgroup2DBlockStoreInstruction storeNdInst;
- static const Subgroup2DBlockPrefetchInstruction prefetchNdInst;
- static const SpirvStoreScatterInstruction storeScatterInst;
- static const SpirvLoadGatherInstruction loadGatherInst;
- static const Instruction *arr[] = {
- &dpasInst, &dpasMxInst, &loadNdInst, &storeNdInst,
- &prefetchNdInst, &storeScatterInst, &loadGatherInst};
- return arr;
- }
-
- CRIuArch()
- : Xe2Plus(Kind::CRI, getInstructionRegistryArr(),
- // Using bmg config as placeholder
- // TODO: Update to actual XeCore and SharedMemory config
- XeCoreInfo(8, SharedMemory(256 * 1024, 4), 8, 8) // xeCore
- ) {}
- static bool classof(const uArch *u) { return u->getKind() == Kind::CRI; }
- static const uArch *getInstance() {
- static const CRIuArch instance;
- return reinterpret_cast<const uArch *>(&instance);
- }
-};
-
-inline const uArch *getUArch(llvm::StringRef archName) {
- if (archName.equals_insensitive("pvc"))
- return PVCuArch::getInstance();
- if (archName.equals_insensitive("bmg"))
- return BMGuArch::getInstance();
- if (archName.equals_insensitive("cri"))
- return CRIuArch::getInstance();
- return nullptr;
-}
} // namespace uArch
} // namespace xegpu
} // namespace mlir
diff --git a/mlir/include/mlir/Dialect/XeGPU/uArch/IntelGpuXe3.h b/mlir/include/mlir/Dialect/XeGPU/uArch/IntelGpuXe3.h
new file mode 100644
index 0000000000000..a831196d2a8df
--- /dev/null
+++ b/mlir/include/mlir/Dialect/XeGPU/uArch/IntelGpuXe3.h
@@ -0,0 +1,603 @@
+//===--- IntelGpuXe3.h ------------------------------------------*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// \file
+// Xe2 uArch definition. Xe3 is the second generation of Intel Xe GPUs.
+// This file defines the uArch details for Xe3 and its derived architectures.
+// This includes Crescent Island Architecture.
+//
+//===----------------------------------------------------------------------===//
+#ifndef MLIR_DIALECT_XEGPU_UARCH_INTELGPUXE3_H
+#define MLIR_DIALECT_XEGPU_UARCH_INTELGPUXE3_H
+
+#include "mlir/Dialect/XeGPU/uArch/uArchBase.h"
+#include "mlir/IR/BuiltinTypes.h"
+#include "mlir/IR/TypeUtilities.h"
+#include "llvm/ADT/SmallVector.h"
+#include "llvm/Support/DebugLog.h"
+#include <map>
+#include <string>
+
+using namespace mlir;
+using namespace mlir::xegpu::uArch;
+
+namespace mlir {
+namespace xegpu {
+namespace uArch {
+
+struct Xe3 : public uArch {
+ Xe3(Kind kind, llvm::ArrayRef<const Instruction *> instructionRegistry,
+ const XeCoreInfo &xeCore)
+ : uArch(kind, instructionRegistry), xeCore(xeCore) {}
+ int getSubgroupSize() const override { return 16; }
+ unsigned getGeneralPackedFormatBitSize() const override { return 32; }
+
+ static bool classof(const uArch *u) {
+ return u->getKind() >= Kind::Xe3_First && u->getKind() <= Kind::Xe3_Last;
+ }
+
+protected:
+ XeCoreInfo xeCore;
+};
+
+//===----------------------------------------------------------------------===//
+// uArch instructions
+//===----------------------------------------------------------------------===//
+struct CriSubgroup2DBlockStoreInstruction : public Instruction {
+ CriSubgroup2DBlockStoreInstruction()
+ : Instruction(InstructionKind::Subgroup2DBlockStore,
+ InstructionScope::Subgroup) {}
+ static bool classof(const Instruction *B) {
+ return B->getInstructionKind() == InstructionKind::Subgroup2DBlockStore;
+ }
+ // 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
+ std::optional<
+ std::tuple<llvm::ArrayRef<int>, llvm::ArrayRef<int>, llvm::ArrayRef<int>>>
+ getBlockWidthHeightCount(Type elemTy) const {
+ const static int kHeight[] = {1, 2, 4, 8};
+ const static int kWidth16[] = {16};
+ const static int kWidth32[] = {16};
+ const static int kCount[] = {1};
+ const int elemByteSize = elemTy.getIntOrFloatBitWidth() / 8;
+ if (elemByteSize == 1)
+ return std::make_tuple(llvm::ArrayRef<int>(kWidth32),
+ llvm::ArrayRef<int>(kHeight),
+ llvm::ArrayRef<int>(kCount));
+ else if (elemByteSize == 2 || elemByteSize == 4)
+ return std::make_tuple(llvm::ArrayRef<int>(kWidth16),
+ llvm::ArrayRef<int>(kHeight),
+ llvm::ArrayRef<int>(kCount));
+ return std::nullopt;
+ }
+
+ int32_t getPackedFormatBitSize() const { return 16; }
+};
+
+struct CriSubgroup2DBlockLoadInstruction : public Instruction {
+ CriSubgroup2DBlockLoadInstruction()
+ : Instruction(InstructionKind::Subgroup2DBlockLoad,
+ InstructionScope::Subgroup) {}
+ static bool classof(const Instruction *B) {
+ return B->getInstructionKind() == InstructionKind::Subgroup2DBlockLoad;
+ }
+
+ // 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
+ std::optional<
+ std::tuple<llvm::ArrayRef<int>, llvm::ArrayRef<int>, llvm::ArrayRef<int>>>
+ getBlockWidthHeightCount(Type elemTy, bool hasTransform, bool hasTranspose,
+ bool upConv = false) const {
+ 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 kWidth32[] = {32};
+ static const int kWidth16[] = {16};
+ 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)
+ 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>>;
+ 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}},
+ // Block Loads with Transform:
+ {{1, true, false, false}, {kWidth16, kHeightAtLeast32, kCount4}},
+ {{2, 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});
+ if (it != kMap.end())
+ return it->second;
+ return std::nullopt;
+ }
+
+ int32_t getPackedFormatBitSize() const { return 16; }
+};
+
+struct CriSubgroup2DBlockPrefetchInstruction : public Instruction {
+ CriSubgroup2DBlockPrefetchInstruction()
+ : Instruction(InstructionKind::Subgroup2DBlockPrefetch,
+ InstructionScope::Subgroup) {}
+ static bool classof(const Instruction *B) {
+ return B->getInstructionKind() == InstructionKind::Subgroup2DBlockPrefetch;
+ }
+ // 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
+ std::optional<
+ std::tuple<llvm::ArrayRef<int>, llvm::ArrayRef<int>, llvm::ArrayRef<int>>>
+ getBlockWidthHeightCount(Type elemTy) const {
+ static const int kHeightAtLeast1[] = {1, 2, 4, 8, 16, 32};
+
+ static const int kWidth32[] = {32};
+ static const int kWidth16[] = {16};
+
+ static const int32_t kCount1[] = {1};
+ static const int32_t kCount2[] = {1, 2};
+ // elemBytes
+ using Key = int;
+ // (widths, heights, counts)
+ using Value = std::tuple<llvm::ArrayRef<int32_t>, llvm::ArrayRef<int32_t>,
+ llvm::ArrayRef<int32_t>>;
+ static const llvm::DenseMap<Key, Value> kMap = {
+ {1, {kWidth32, kHeightAtLeast1, kCount2}},
+ {2, {kWidth16, kHeightAtLeast1, kCount2}},
+ {4, {kWidth16, kHeightAtLeast1, kCount1}},
+ };
+ const int elemByteSize = elemTy.getIntOrFloatBitWidth() / 8;
+ auto it = kMap.find(elemByteSize);
+ if (it != kMap.end())
+ return it->second;
+ return std::nullopt;
+ }
+ int32_t getPackedFormatBitSize() const { return 16; }
+};
+
+struct CriSubgroupMatrixMultiplyAcc : public Instruction,
+ public MMAInstructionInterface {
+ CriSubgroupMatrixMultiplyAcc(unsigned packedFormatBitSizeA,
+ unsigned packedFormatBitSizeB)
+ : Instruction(InstructionKind::SubgroupMatrixMultiplyAcc,
+ InstructionScope::Subgroup),
+ packedFormatBitSizeA(packedFormatBitSizeA),
+ packedFormatBitSizeB(packedFormatBitSizeB) {}
+ static bool classof(const Instruction *B) {
+ return B->getInstructionKind() ==
+ InstructionKind::SubgroupMatrixMultiplyAcc;
+ }
+ // Source:
+ // https://registry.khronos.org/OpenCL/extensions/intel/cl_intel_subgroup_matrix_multiply_accumulate.html
+
+ // Override all virtuals from MatrixOpInterface
+ virtual llvm::SmallVector<std::pair<uint32_t, uint32_t>, 16>
+ getSupportedShapes(Type dataType, MMAOpndKind matrixType) override;
+ virtual llvm::SmallVector<Type, 8>
+ getSupportedTypes(MLIRContext &context, MMAOpndKind matrixType) override;
+ virtual bool
+ checkSupportedShapesAndTypes(std::pair<uint32_t, uint32_t> AShape,
+ std::pair<uint32_t, uint32_t> BShape,
+ std::pair<uint32_t, uint32_t> CShape,
+ std::pair<uint32_t, uint32_t> DShape, Type AType,
+ Type BType, Type CType, Type DType) override;
+ virtual bool checkSupportedTypes(Type AType, Type BType, Type CType,
+ Type DType) override;
+ virtual bool validate(std::pair<uint32_t, uint32_t> AShape,
+ std::pair<uint32_t, uint32_t> BShape,
+ std::pair<uint32_t, uint32_t> CShape,
+ std::pair<uint32_t, uint32_t> DShape, Type AType,
+ Type BType, Type CType, Type DType) override;
+ virtual llvm::SmallVector<uint32_t, 8>
+ getSupportedM(Type type) const override;
+ virtual llvm::SmallVector<uint32_t, 8>
+ getSupportedK(Type type) const override;
+ virtual llvm::SmallVector<uint32_t, 8>
+ getSupportedN(Type type) const override;
+
+ unsigned getPackedFormatBitSizeA() const { return packedFormatBitSizeA; }
+ unsigned getPackedFormatBitSizeB() const { return packedFormatBitSizeB; }
+ bool isLaneLayoutRowMajorOrder() const override { return true; }
+
+protected:
+ const unsigned packedFormatBitSizeA;
+ const unsigned packedFormatBitSizeB;
+};
+
+struct CriSubgroupScaledMatrixMultiplyAcc : public Instruction,
+ public MMAInstructionInterface {
+ CriSubgroupScaledMatrixMultiplyAcc(unsigned packedFormatBitSizeA,
+ unsigned packedFormatBitSizeB)
+ : Instruction(InstructionKind::SubgroupScaledMatrixMultiplyAcc,
+ InstructionScope::Subgroup),
+ packedFormatBitSizeA(packedFormatBitSizeA),
+ packedFormatBitSizeB(packedFormatBitSizeB) {}
+ static bool classof(const Instruction *B) {
+ return B->getInstructionKind() ==
+ InstructionKind::SubgroupScaledMatrixMultiplyAcc;
+ }
+ // Source:
+ // https://github.com/intel/llvm/blob/sycl/sycl/doc/design/spirv-extensions/SPV_INTEL_subgroup_scaled_matrix_multiply_accumulate.asciidoc
+
+ // Override all virtuals from MatrixOpInterface
+ virtual llvm::SmallVector<std::pair<uint32_t, uint32_t>, 16>
+ getSupportedShapes(Type dataType, MMAOpndKind matrixType) override;
+ virtual llvm::SmallVector<Type, 8>
+ getSupportedTypes(MLIRContext &context, MMAOpndKind matrixType) override;
+ virtual bool
+ checkSupportedShapesAndTypes(std::pair<uint32_t, uint32_t> AShape,
+ std::pair<uint32_t, uint32_t> BShape,
+ std::pair<uint32_t, uint32_t> CShape,
+ std::pair<uint32_t, uint32_t> DShape, Type AType,
+ Type BType, Type CType, Type DType) override;
+ virtual bool checkSupportedTypes(Type AType, Type BType, Type CType,
+ Type DType) override;
+ virtual bool validate(std::pair<uint32_t, uint32_t> AShape,
+ std::pair<uint32_t, uint32_t> BShape,
+ std::pair<uint32_t, uint32_t> CShape,
+ std::pair<uint32_t, uint32_t> DShape, Type AType,
+ Type BType, Type CType, Type DType) override;
+ virtual llvm::SmallVector<uint32_t, 8>
+ getSupportedM(Type type) const override;
+ virtual llvm::SmallVector<uint32_t, 8>
+ getSupportedK(Type type) const override;
+ virtual llvm::SmallVector<uint32_t, 8>
+ getSupportedN(Type type) const override;
+
+ unsigned getPackedFormatBitSizeA() const { return packedFormatBitSizeA; }
+ unsigned getPackedFormatBitSizeB() const { return packedFormatBitSizeB; }
+ bool isLaneLayoutRowMajorOrder() const override { return true; }
+
+protected:
+ const unsigned packedFormatBitSizeA;
+ const unsigned packedFormatBitSizeB;
+};
+
+//===----------------------------------------------------------------------===//
+// uArch instances
+//===----------------------------------------------------------------------===//
+
+struct CRIuArch : public Xe3 {
+ static llvm::ArrayRef<const Instruction *> getInstructionRegistryArr() {
+ static const CriSubgroupMatrixMultiplyAcc dpasInst{16, 32};
+ static const CriSubgroupScaledMatrixMultiplyAcc dpasMxInst{16, 32};
+ static const CriSubgroup2DBlockLoadInstruction loadNdInst;
+ static const CriSubgroup2DBlockStoreInstruction storeNdInst;
+ static const CriSubgroup2DBlockPrefetchInstruction prefetchNdInst;
+ static const SpirvStoreScatterInstruction storeScatterInst;
+ static const SpirvLoadGatherInstruction loadGatherInst;
+ static const Instruction *arr[] = {
+ &dpasInst, &dpasMxInst, &loadNdInst, &storeNdInst,
+ &prefetchNdInst, &storeScatterInst, &loadGatherInst};
+ return arr;
+ }
+
+ CRIuArch()
+ : Xe3(Kind::CRI, getInstructionRegistryArr(),
+ // Using bmg config as placeholder
+ // TODO: Update to actual XeCore and SharedMemory config
+ XeCoreInfo(8, SharedMemory(256 * 1024, 4), 8, 8) // xeCore
+ ) {}
+ static bool classof(const uArch *u) { return u->getKind() == Kind::CRI; }
+ static const uArch *getInstance() {
+ static const CRIuArch instance;
+ return reinterpret_cast<const uArch *>(&instance);
+ }
+};
+
+} // namespace uArch
+} // namespace xegpu
+} // namespace mlir
+
+//===----------------------------------------------------------------------===//
+// Instruction implementations
+//===----------------------------------------------------------------------===//
+
+inline llvm::SmallVector<std::pair<uint32_t, uint32_t>, 16>
+CriSubgroupMatrixMultiplyAcc::getSupportedShapes(Type dataType,
+ MMAOpndKind matrixType) {
+ auto combineVectors = [](const llvm::SmallVector<uint32_t, 8> &a,
+ const llvm::SmallVector<uint32_t, 8> &b)
+ -> llvm::SmallVector<std::pair<uint32_t, uint32_t>, 16> {
+ llvm::SmallVector<std::pair<uint32_t, uint32_t>, 16> result;
+ for (unsigned x : a) {
+ for (unsigned y : b) {
+ result.emplace_back(x, y);
+ }
+ }
+ return result;
+ };
+
+ auto M = getSupportedM(dataType);
+ auto K = getSupportedK(dataType);
+ auto N = getSupportedN(dataType);
+ llvm::SmallVector<std::pair<unsigned, unsigned>, 16> resultMatrix;
+
+ switch (matrixType) {
+ case MMAOpndKind::MatrixA:
+ resultMatrix = combineVectors(M, K);
+ break;
+ case MMAOpndKind::MatrixB:
+ resultMatrix = combineVectors(K, N);
+ break;
+ case MMAOpndKind::MatrixC:
+ resultMatrix = combineVectors(M, N);
+ break;
+ case MMAOpndKind::MatrixD:
+ resultMatrix = combineVectors(M, N);
+ break;
+ }
+ return resultMatrix;
+}
+
+inline llvm::SmallVector<Type, 8>
+CriSubgroupMatrixMultiplyAcc::getSupportedTypes(MLIRContext &context,
+ MMAOpndKind matrixType) {
+ Type bf16Type = BFloat16Type::get(&context);
+ Type f16Type = Float16Type::get(&context);
+ Type tf32Type = FloatTF32Type::get(&context);
+ Type f32Type = Float32Type::get(&context);
+
+ switch (matrixType) {
+ case MMAOpndKind::MatrixA:
+ return {bf16Type, f16Type, tf32Type};
+ case MMAOpndKind::MatrixB:
+ return {bf16Type, f16Type, tf32Type};
+ case MMAOpndKind::MatrixC:
+ return {bf16Type, f16Type, f32Type};
+ case MMAOpndKind::MatrixD:
+ return {bf16Type, f16Type, f32Type};
+ }
+ return {};
+}
+
+inline bool CriSubgroupMatrixMultiplyAcc::checkSupportedTypes(Type AType,
+ Type BType,
+ Type CType,
+ Type DType) {
+ if (AType.isF16() || BType.isF16()) {
+ if (AType != BType || (CType && (!CType.isF32() && !CType.isF16())) ||
+ (!DType.isF32() && !DType.isF16())) {
+ LDBG() << "Unsupported dpas combinations of Dst, Acc, A and B matrices.";
+ return false;
+ }
+ } else if (AType.isBF16() || BType.isBF16()) {
+ if (AType != BType || (CType && (!CType.isF32() && !CType.isBF16())) ||
+ (!DType.isF32() && !DType.isBF16())) {
+ LDBG() << "Unsupported dpas combinations of Dst, Acc, A and B matrices.";
+ return false;
+ }
+ } else if (AType.isTF32() || BType.isTF32()) {
+ if (AType != BType || (CType && (!CType.isF32() && !DType.isF32())) ||
+ (!DType.isF32())) {
+ LDBG() << "Unsupported dpas combinations of Dst, Acc, A and B matrices.";
+ return false;
+ }
+ } else if (!(AType.isInteger(2) || AType.isInteger(4) ||
+ AType.isInteger(8)) &&
+ !(BType.isInteger(2) || BType.isInteger(4) ||
+ BType.isInteger(8))) {
+ LDBG() << "Unsupported dpas combinations of Dst, Acc, A and B matrices.";
+ return false;
+ }
+
+ return true;
+}
+
+inline bool CriSubgroupMatrixMultiplyAcc::checkSupportedShapesAndTypes(
+ std::pair<uint32_t, uint32_t> AShape, std::pair<uint32_t, uint32_t> BShape,
+ std::pair<uint32_t, uint32_t> CShape, std::pair<uint32_t, uint32_t> DShape,
+ Type AType, Type BType, Type CType, Type DType) {
+ auto supportedAShapes = getSupportedShapes(AType, MMAOpndKind::MatrixA);
+ auto supportedBShapes = getSupportedShapes(BType, MMAOpndKind::MatrixB);
+ auto supportedCShapes = getSupportedShapes(CType, MMAOpndKind::MatrixC);
+ auto supportedDShapes = getSupportedShapes(DType, MMAOpndKind::MatrixD);
+ return llvm::is_contained(supportedAShapes, AShape) &&
+ llvm::is_contained(supportedBShapes, BShape) &&
+ llvm::is_contained(supportedCShapes, CShape) &&
+ llvm::is_contained(supportedDShapes, DShape) &&
+ checkSupportedTypes(AType, BType, CType, DType);
+}
+
+inline bool CriSubgroupMatrixMultiplyAcc::validate(
+ std::pair<uint32_t, uint32_t> AShape, std::pair<uint32_t, uint32_t> BShape,
+ std::pair<uint32_t, uint32_t> CShape, std::pair<uint32_t, uint32_t> DShape,
+ Type AType, Type BType, Type CType, Type DType) {
+ return checkSupportedShapesAndTypes(AShape, BShape, CShape, DShape, AType,
+ BType, CType, DType);
+}
+
+inline llvm::SmallVector<uint32_t, 8>
+CriSubgroupMatrixMultiplyAcc::getSupportedM(Type type) const {
+ return {1, 2, 3, 4, 5, 6, 7, 8};
+}
+
+inline llvm::SmallVector<uint32_t, 8>
+CriSubgroupMatrixMultiplyAcc::getSupportedK(Type type) const {
+ // assert if data type is not int or float type
+ assert(type.isIntOrFloat() && "Matrix type must be int or float");
+ auto bitWidth = type.getIntOrFloatBitWidth();
+ uint32_t kSize = 0;
+ switch (bitWidth) {
+ case 2:
+ kSize = 64;
+ break;
+ case 4:
+ kSize = 64;
+ break;
+ case 8:
+ kSize = 32;
+ break;
+ case 16:
+ kSize = 16;
+ break;
+ case 32:
+ kSize = 8;
+ break;
+ default:
+ llvm_unreachable("Invalid int or float");
+ }
+ return {kSize};
+}
+
+inline llvm::SmallVector<uint32_t, 8>
+CriSubgroupMatrixMultiplyAcc::getSupportedN(Type type) const {
+ return {16};
+}
+
+//===----------------------------------------------------------------------===//
+// SubgroupScaledMatrixMultiplyAcc implementations
+//===----------------------------------------------------------------------===//
+
+inline llvm::SmallVector<std::pair<uint32_t, uint32_t>, 16>
+CriSubgroupScaledMatrixMultiplyAcc::getSupportedShapes(Type dataType,
+ MMAOpndKind matrixType) {
+ auto combineVectors = [](const llvm::SmallVector<uint32_t, 8> &a,
+ const llvm::SmallVector<uint32_t, 8> &b)
+ -> llvm::SmallVector<std::pair<uint32_t, uint32_t>, 16> {
+ llvm::SmallVector<std::pair<uint32_t, uint32_t>, 16> result;
+ for (unsigned x : a) {
+ for (unsigned y : b) {
+ result.emplace_back(x, y);
+ }
+ }
+ return result;
+ };
+
+ // Avoid calling getSupportedK for C/D types (which are f32/bf16
+ // and not valid for the K-dimension bit-width calculation).
+ switch (matrixType) {
+ case MMAOpndKind::MatrixA:
+ return combineVectors(getSupportedM(dataType), getSupportedK(dataType));
+ case MMAOpndKind::MatrixB:
+ return combineVectors(getSupportedK(dataType), getSupportedN(dataType));
+ case MMAOpndKind::MatrixC:
+ case MMAOpndKind::MatrixD:
+ return combineVectors(getSupportedM(dataType), getSupportedN(dataType));
+ }
+ return {};
+}
+
+inline llvm::SmallVector<Type, 8>
+CriSubgroupScaledMatrixMultiplyAcc::getSupportedTypes(MLIRContext &context,
+ MMAOpndKind matrixType) {
+ Type f8E4M3FNType = Float8E4M3FNType::get(&context);
+ Type f8E5M2Type = Float8E5M2Type::get(&context);
+ Type f4E2M1FNType = Float4E2M1FNType::get(&context);
+ Type bf16Type = BFloat16Type::get(&context);
+ Type f32Type = Float32Type::get(&context);
+
+ switch (matrixType) {
+ case MMAOpndKind::MatrixA:
+ return {f8E4M3FNType, f8E5M2Type, f4E2M1FNType};
+ case MMAOpndKind::MatrixB:
+ return {f8E4M3FNType, f8E5M2Type, f4E2M1FNType};
+ case MMAOpndKind::MatrixC:
+ return {bf16Type, f32Type};
+ case MMAOpndKind::MatrixD:
+ return {bf16Type, f32Type};
+ }
+ return {};
+}
+
+inline bool CriSubgroupScaledMatrixMultiplyAcc::checkSupportedTypes(
+ Type AType, Type BType, Type CType, Type DType) {
+ auto isSupportedLowPrecision = [](Type t) {
+ return t.isF8E4M3FN() || t.isF8E5M2() || llvm::isa<Float4E2M1FNType>(t);
+ };
+ auto isSupportedAccum = [](Type t) { return t.isF32() || t.isBF16(); };
+
+ if (!isSupportedLowPrecision(AType) || !isSupportedLowPrecision(BType)) {
+ LDBG() << "Unsupported scaled dpas: A and B must be FP8 or FP4 types.";
+ return false;
+ }
+
+ // A and B must have the same bit width for K dimension compatibility.
+ if (AType.getIntOrFloatBitWidth() != BType.getIntOrFloatBitWidth()) {
+ LDBG() << "Unsupported scaled dpas: A and B must have the same bit width.";
+ return false;
+ }
+
+ if (CType && !isSupportedAccum(CType)) {
+ LDBG() << "Unsupported scaled dpas: C must be f32 or bf16.";
+ return false;
+ }
+
+ if (!isSupportedAccum(DType)) {
+ LDBG() << "Unsupported scaled dpas: D must be f32 or bf16.";
+ return false;
+ }
+
+ return true;
+}
+
+inline bool CriSubgroupScaledMatrixMultiplyAcc::checkSupportedShapesAndTypes(
+ std::pair<uint32_t, uint32_t> AShape, std::pair<uint32_t, uint32_t> BShape,
+ std::pair<uint32_t, uint32_t> CShape, std::pair<uint32_t, uint32_t> DShape,
+ Type AType, Type BType, Type CType, Type DType) {
+ auto supportedAShapes = getSupportedShapes(AType, MMAOpndKind::MatrixA);
+ auto supportedBShapes = getSupportedShapes(BType, MMAOpndKind::MatrixB);
+ auto supportedCShapes = getSupportedShapes(CType, MMAOpndKind::MatrixC);
+ auto supportedDShapes = getSupportedShapes(DType, MMAOpndKind::MatrixD);
+ return llvm::is_contained(supportedAShapes, AShape) &&
+ llvm::is_contained(supportedBShapes, BShape) &&
+ llvm::is_contained(supportedCShapes, CShape) &&
+ llvm::is_contained(supportedDShapes, DShape) &&
+ checkSupportedTypes(AType, BType, CType, DType);
+}
+
+inline bool CriSubgroupScaledMatrixMultiplyAcc::validate(
+ std::pair<uint32_t, uint32_t> AShape, std::pair<uint32_t, uint32_t> BShape,
+ std::pair<uint32_t, uint32_t> CShape, std::pair<uint32_t, uint32_t> DShape,
+ Type AType, Type BType, Type CType, Type DType) {
+ return checkSupportedShapesAndTypes(AShape, BShape, CShape, DShape, AType,
+ BType, CType, DType);
+}
+
+inline llvm::SmallVector<uint32_t, 8>
+CriSubgroupScaledMatrixMultiplyAcc::getSupportedM(Type type) const {
+ return {8};
+}
+
+inline llvm::SmallVector<uint32_t, 8>
+CriSubgroupScaledMatrixMultiplyAcc::getSupportedK(Type type) const {
+ assert(type.isIntOrFloat() && "Matrix type must be int or float");
+ auto bitWidth = type.getIntOrFloatBitWidth();
+ uint32_t kSize = 0;
+ switch (bitWidth) {
+ case 4:
+ kSize = 64; // FP4: scale K by 4 (base 16-bit K=16 -> 64)
+ break;
+ case 8:
+ kSize = 32; // FP8: scale K by 2 (base 16-bit K=16 -> 32)
+ break;
+ default:
+ // Scaled dpas only supports FP8 (8-bit) and FP4 (4-bit) types for A/B
+ // matrices. Return empty so callers can gracefully reject unsupported
+ // types instead of aborting.
+ return {};
+ }
+ return {kSize};
+}
+
+inline llvm::SmallVector<uint32_t, 8>
+CriSubgroupScaledMatrixMultiplyAcc::getSupportedN(Type type) const {
+ return {16};
+}
+
+#endif // MLIR_DIALECT_XEGPU_UARCH_INTELGPUXE3_H
diff --git a/mlir/include/mlir/Dialect/XeGPU/uArch/uArchBase.h b/mlir/include/mlir/Dialect/XeGPU/uArch/uArchBase.h
index 1702c7b46e0eb..2bb45e38b0fe2 100644
--- a/mlir/include/mlir/Dialect/XeGPU/uArch/uArchBase.h
+++ b/mlir/include/mlir/Dialect/XeGPU/uArch/uArchBase.h
@@ -156,11 +156,13 @@ struct CacheInfo {
struct uArch {
enum class Kind {
// Xe2 family
- Xe2Plus_First,
- PVC = Xe2Plus_First,
+ Xe2_First,
+ PVC = Xe2_First,
BMG,
- CRI,
- Xe2Plus_Last = CRI,
+ Xe2_Last = BMG,
+ Xe3_First,
+ CRI = Xe3_First,
+ Xe3_Last = CRI
};
// Constructor
@@ -182,6 +184,8 @@ struct uArch {
return "bmg";
case Kind::CRI:
return "cri";
+ default:
+ return "";
}
llvm_unreachable("Unknown uArch::Kind");
}
@@ -333,6 +337,14 @@ struct StoreScatterInstructionInterface : public Instruction {
virtual ~StoreScatterInstructionInterface() = default;
};
+struct SpirvLoadGatherInstruction : public LoadGatherInstructionInterface {
+ int32_t getMaxLaneLoadSize(int32_t bitWidth) const override { return 16; }
+};
+
+struct SpirvStoreScatterInstruction : public StoreScatterInstructionInterface {
+ int32_t getMaxLaneStoreSize(int32_t bitWidth) const override { return 16; }
+};
+
} // namespace uArch
} // namespace xegpu
} // namespace mlir
diff --git a/mlir/include/mlir/Dialect/XeGPU/uArch/uArchCommon.h b/mlir/include/mlir/Dialect/XeGPU/uArch/uArchCommon.h
new file mode 100644
index 0000000000000..82cfb5538f23f
--- /dev/null
+++ b/mlir/include/mlir/Dialect/XeGPU/uArch/uArchCommon.h
@@ -0,0 +1,38 @@
+//===- uArchCommon.h --------------------------------------------*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// \file
+// Common functionality related to uArch instances.
+//
+//
+//===----------------------------------------------------------------------===//
+#ifndef MLIR_DIALECT_XEGPU_UARCH_UARCHCOMMON_H
+#define MLIR_DIALECT_XEGPU_UARCH_UARCHCOMMON_H
+
+#include "IntelGpuXe2.h"
+#include "IntelGpuXe3.h"
+
+namespace mlir {
+namespace xegpu {
+namespace uArch {
+
+inline const uArch *getUArch(llvm::StringRef archName) {
+ if (archName.equals_insensitive("pvc"))
+ return PVCuArch::getInstance();
+ if (archName.equals_insensitive("bmg"))
+ return BMGuArch::getInstance();
+ if (archName.equals_insensitive("cri"))
+ return CRIuArch::getInstance();
+ return nullptr;
+}
+
+} // namespace uArch
+} // namespace xegpu
+} // namespace mlir
+
+#endif // MLIR_DIALECT_XEGPU_UARCH_UARCHCOMMON_H
diff --git a/mlir/lib/Conversion/XeGPUToXeVM/XeGPUToXeVM.cpp b/mlir/lib/Conversion/XeGPUToXeVM/XeGPUToXeVM.cpp
index c7d593c74f264..6144d7c0c1a15 100644
--- a/mlir/lib/Conversion/XeGPUToXeVM/XeGPUToXeVM.cpp
+++ b/mlir/lib/Conversion/XeGPUToXeVM/XeGPUToXeVM.cpp
@@ -23,7 +23,7 @@
#include "mlir/Dialect/Vector/IR/VectorOps.h"
#include "mlir/Dialect/XeGPU/IR/XeGPU.h"
#include "mlir/Dialect/XeGPU/Utils/XeGPUUtils.h"
-#include "mlir/Dialect/XeGPU/uArch/IntelGpuXe2.h"
+#include "mlir/Dialect/XeGPU/uArch/uArchCommon.h"
#include "mlir/Pass/Pass.h"
#include "mlir/Support/LLVM.h"
#include "llvm/ADT/STLExtras.h"
diff --git a/mlir/lib/Dialect/XeGPU/IR/XeGPUDialect.cpp b/mlir/lib/Dialect/XeGPU/IR/XeGPUDialect.cpp
index 5d55342afea15..f9c258aba5bc7 100644
--- a/mlir/lib/Dialect/XeGPU/IR/XeGPUDialect.cpp
+++ b/mlir/lib/Dialect/XeGPU/IR/XeGPUDialect.cpp
@@ -12,7 +12,7 @@
#include "mlir/Dialect/LLVMIR/XeVMDialect.h"
#include "mlir/Dialect/Utils/IndexingUtils.h"
#include "mlir/Dialect/XeGPU/IR/XeGPU.h"
-#include "mlir/Dialect/XeGPU/uArch/IntelGpuXe2.h"
+#include "mlir/Dialect/XeGPU/uArch/uArchCommon.h"
#include "mlir/IR/Builders.h"
#include "mlir/IR/DialectImplementation.h"
#include "llvm/ADT/SmallVectorExtras.h"
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUArrayLengthOptimization.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUArrayLengthOptimization.cpp
index 53da46cbdee73..590d1804167a5 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUArrayLengthOptimization.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUArrayLengthOptimization.cpp
@@ -10,8 +10,8 @@
#include "mlir/Dialect/XeGPU/IR/XeGPU.h"
#include "mlir/Dialect/XeGPU/Transforms/Transforms.h"
#include "mlir/Dialect/XeGPU/Utils/XeGPUUtils.h"
-#include "mlir/Dialect/XeGPU/uArch/IntelGpuXe2.h"
#include "mlir/Dialect/XeGPU/uArch/uArchBase.h"
+#include "mlir/Dialect/XeGPU/uArch/uArchCommon.h"
#include "mlir/IR/PatternMatch.h"
#include "llvm/ADT/SmallVector.h"
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPeepHoleOptimizer.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPeepHoleOptimizer.cpp
index 566d0eb433606..1ed418d446886 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPeepHoleOptimizer.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPeepHoleOptimizer.cpp
@@ -18,8 +18,8 @@
#include "mlir/Dialect/XeGPU/Transforms/Transforms.h"
#include "mlir/Dialect/XeGPU/Transforms/XeGPULayoutImpl.h"
#include "mlir/Dialect/XeGPU/Utils/XeGPUUtils.h"
-#include "mlir/Dialect/XeGPU/uArch/IntelGpuXe2.h"
#include "mlir/Dialect/XeGPU/uArch/uArchBase.h"
+#include "mlir/Dialect/XeGPU/uArch/uArchCommon.h"
#include "mlir/IR/BuiltinTypes.h"
#include "mlir/IR/OpDefinition.h"
#include "mlir/IR/Types.h"
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
index 8600492e4bf41..57cf1d0669cb1 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
@@ -17,7 +17,7 @@
#include "mlir/Dialect/XeGPU/Transforms/Passes.h"
#include "mlir/Dialect/XeGPU/Transforms/XeGPULayoutImpl.h"
#include "mlir/Dialect/XeGPU/Utils/XeGPUUtils.h"
-#include "mlir/Dialect/XeGPU/uArch/IntelGpuXe2.h"
+#include "mlir/Dialect/XeGPU/uArch/uArchCommon.h"
#include "mlir/IR/Attributes.h"
#include "mlir/IR/Builders.h"
#include "mlir/IR/BuiltinAttributes.h"
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
index 75a87f84b3da8..766754af76cb7 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
@@ -16,7 +16,7 @@
#include "mlir/Dialect/XeGPU/Transforms/Transforms.h"
#include "mlir/Dialect/XeGPU/Transforms/XeGPULayoutImpl.h"
#include "mlir/Dialect/XeGPU/Utils/XeGPUUtils.h"
-#include "mlir/Dialect/XeGPU/uArch/IntelGpuXe2.h"
+#include "mlir/Dialect/XeGPU/uArch/uArchCommon.h"
#include "mlir/IR/Builders.h"
#include "mlir/IR/BuiltinAttributes.h"
#include "mlir/IR/BuiltinOps.h"
diff --git a/mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp b/mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp
index 24999f85c1268..1aad2aa77741b 100644
--- a/mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp
+++ b/mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp
@@ -18,7 +18,7 @@
#include "mlir/Dialect/Utils/IndexingUtils.h"
#include "mlir/Dialect/Vector/IR/VectorOps.h"
#include "mlir/Dialect/XeGPU/IR/XeGPU.h"
-#include "mlir/Dialect/XeGPU/uArch/IntelGpuXe2.h"
+#include "mlir/Dialect/XeGPU/uArch/uArchCommon.h"
#include "mlir/IR/Builders.h"
#include "mlir/IR/BuiltinOps.h"
#include "mlir/IR/Operation.h"
@@ -786,7 +786,7 @@ bool xegpu::requireTranspose(const xegpu::DistributeLayoutAttr layout,
const xegpu::uArch::uArch *uArch) {
// Return false for unsupported targets.
// TODO: Add more support or move to target info.
- if (!isa<xegpu::uArch::Xe2Plus>(uArch))
+ if (!isa<xegpu::uArch::Xe2>(uArch) && !isa<xegpu::uArch::Xe3>(uArch))
return false;
if (!layout)
return false;
>From 324cee8a6ca1c74d2e49dc5765e37b32ca803531 Mon Sep 17 00:00:00 2001
From: Artem Kroviakov <artem.kroviakov at intel.com>
Date: Thu, 25 Jun 2026 09:47:23 +0000
Subject: [PATCH 3/4] Fix uarch tables
---
.../mlir/Dialect/XeGPU/uArch/IntelGpuXe3.h | 67 +++++++++++++------
1 file changed, 45 insertions(+), 22 deletions(-)
diff --git a/mlir/include/mlir/Dialect/XeGPU/uArch/IntelGpuXe3.h b/mlir/include/mlir/Dialect/XeGPU/uArch/IntelGpuXe3.h
index a831196d2a8df..982dbdbea5473 100644
--- a/mlir/include/mlir/Dialect/XeGPU/uArch/IntelGpuXe3.h
+++ b/mlir/include/mlir/Dialect/XeGPU/uArch/IntelGpuXe3.h
@@ -48,7 +48,8 @@ struct Xe3 : public uArch {
//===----------------------------------------------------------------------===//
// uArch instructions
//===----------------------------------------------------------------------===//
-struct CriSubgroup2DBlockStoreInstruction : public Instruction {
+struct CriSubgroup2DBlockStoreInstruction : public Instruction,
+ public BlockIOInstructionInterface {
CriSubgroup2DBlockStoreInstruction()
: Instruction(InstructionKind::Subgroup2DBlockStore,
InstructionScope::Subgroup) {}
@@ -57,9 +58,12 @@ struct CriSubgroup2DBlockStoreInstruction : 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};
@@ -76,10 +80,11 @@ struct CriSubgroup2DBlockStoreInstruction : public Instruction {
return std::nullopt;
}
- int32_t getPackedFormatBitSize() const { return 16; }
+ int32_t getPackedFormatBitSize() const override { return 16; }
};
-struct CriSubgroup2DBlockLoadInstruction : public Instruction {
+struct CriSubgroup2DBlockLoadInstruction : public Instruction,
+ public BlockIOInstructionInterface {
CriSubgroup2DBlockLoadInstruction()
: Instruction(InstructionKind::Subgroup2DBlockLoad,
InstructionScope::Subgroup) {}
@@ -91,48 +96,64 @@ struct CriSubgroup2DBlockLoadInstruction : 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 CriSubgroup2DBlockPrefetchInstruction : public Instruction {
+struct CriSubgroup2DBlockPrefetchInstruction
+ : public Instruction,
+ public BlockIOInstructionInterface {
CriSubgroup2DBlockPrefetchInstruction()
: Instruction(InstructionKind::Subgroup2DBlockPrefetch,
InstructionScope::Subgroup) {}
@@ -143,7 +164,9 @@ struct CriSubgroup2DBlockPrefetchInstruction : public Instruction {
// https://registry.khronos.org/OpenCL/extensions/intel/cl_intel_subgroup_buffer_prefetch.html#_add_a_new_section_6_15_x_sub_group_prefetch_functions
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 {
static const int kHeightAtLeast1[] = {1, 2, 4, 8, 16, 32};
static const int kWidth32[] = {32};
@@ -167,7 +190,7 @@ struct CriSubgroup2DBlockPrefetchInstruction : public Instruction {
return it->second;
return std::nullopt;
}
- int32_t getPackedFormatBitSize() const { return 16; }
+ int32_t getPackedFormatBitSize() const override { return 16; }
};
struct CriSubgroupMatrixMultiplyAcc : public Instruction,
>From 01aa010cb66cbef33f9ffcba28963d19fca9df84 Mon Sep 17 00:00:00 2001
From: Artem Kroviakov <artem.kroviakov at intel.com>
Date: Thu, 25 Jun 2026 10:06:38 +0000
Subject: [PATCH 4/4] Cleanup
---
.../mlir/Dialect/XeGPU/uArch/IntelGpuXe2.h | 9 ++-------
.../mlir/Dialect/XeGPU/uArch/IntelGpuXe3.h | 9 +++------
.../mlir/Dialect/XeGPU/uArch/uArchBase.h | 17 -----------------
3 files changed, 5 insertions(+), 30 deletions(-)
diff --git a/mlir/include/mlir/Dialect/XeGPU/uArch/IntelGpuXe2.h b/mlir/include/mlir/Dialect/XeGPU/uArch/IntelGpuXe2.h
index ed8ae5289882a..dea5240248089 100644
--- a/mlir/include/mlir/Dialect/XeGPU/uArch/IntelGpuXe2.h
+++ b/mlir/include/mlir/Dialect/XeGPU/uArch/IntelGpuXe2.h
@@ -20,7 +20,6 @@
#include "mlir/IR/TypeUtilities.h"
#include "llvm/ADT/SmallVector.h"
#include "llvm/Support/DebugLog.h"
-#include <map>
#include <string>
using namespace mlir;
@@ -316,7 +315,7 @@ struct PVCuArch final : public Xe2 {
static bool classof(const uArch *u) { return u->getKind() == Kind::PVC; }
static const uArch *getInstance() {
static const PVCuArch instance;
- return reinterpret_cast<const uArch *>(&instance);
+ return &instance;
}
};
@@ -341,7 +340,7 @@ struct BMGuArch : public Xe2 {
static bool classof(const uArch *u) { return u->getKind() == Kind::BMG; }
static const uArch *getInstance() {
static const BMGuArch instance;
- return reinterpret_cast<const uArch *>(&instance);
+ return &instance;
}
};
@@ -381,8 +380,6 @@ SubgroupMatrixMultiplyAcc::getSupportedShapes(Type dataType,
resultMatrix = combineVectors(K, N);
break;
case MMAOpndKind::MatrixC:
- resultMatrix = combineVectors(M, N);
- break;
case MMAOpndKind::MatrixD:
resultMatrix = combineVectors(M, N);
break;
@@ -480,8 +477,6 @@ SubgroupMatrixMultiplyAcc::getSupportedK(Type type) const {
uint32_t kSize = 0;
switch (bitWidth) {
case 2:
- kSize = 64;
- break;
case 4:
kSize = 64;
break;
diff --git a/mlir/include/mlir/Dialect/XeGPU/uArch/IntelGpuXe3.h b/mlir/include/mlir/Dialect/XeGPU/uArch/IntelGpuXe3.h
index 982dbdbea5473..ca16d8ebdddd5 100644
--- a/mlir/include/mlir/Dialect/XeGPU/uArch/IntelGpuXe3.h
+++ b/mlir/include/mlir/Dialect/XeGPU/uArch/IntelGpuXe3.h
@@ -7,7 +7,7 @@
//===----------------------------------------------------------------------===//
//
// \file
-// Xe2 uArch definition. Xe3 is the second generation of Intel Xe GPUs.
+// Xe3 uArch definition. Xe3 is the second generation of Intel Xe GPUs.
// This file defines the uArch details for Xe3 and its derived architectures.
// This includes Crescent Island Architecture.
//
@@ -20,7 +20,6 @@
#include "mlir/IR/TypeUtilities.h"
#include "llvm/ADT/SmallVector.h"
#include "llvm/Support/DebugLog.h"
-#include <map>
#include <string>
using namespace mlir;
@@ -166,7 +165,7 @@ struct CriSubgroup2DBlockPrefetchInstruction
std::tuple<llvm::ArrayRef<int>, llvm::ArrayRef<int>, llvm::ArrayRef<int>>>
getBlockWidthHeightCount(Type elemTy, bool /*hasTransform*/ = false,
bool /*hasTranspose*/ = false,
- bool /*upConv*/ = false) const {
+ bool /*upConv*/ = false) const override {
static const int kHeightAtLeast1[] = {1, 2, 4, 8, 16, 32};
static const int kWidth32[] = {32};
@@ -319,7 +318,7 @@ struct CRIuArch : public Xe3 {
static bool classof(const uArch *u) { return u->getKind() == Kind::CRI; }
static const uArch *getInstance() {
static const CRIuArch instance;
- return reinterpret_cast<const uArch *>(&instance);
+ return &instance;
}
};
@@ -359,8 +358,6 @@ CriSubgroupMatrixMultiplyAcc::getSupportedShapes(Type dataType,
resultMatrix = combineVectors(K, N);
break;
case MMAOpndKind::MatrixC:
- resultMatrix = combineVectors(M, N);
- break;
case MMAOpndKind::MatrixD:
resultMatrix = combineVectors(M, N);
break;
diff --git a/mlir/include/mlir/Dialect/XeGPU/uArch/uArchBase.h b/mlir/include/mlir/Dialect/XeGPU/uArch/uArchBase.h
index 2bb45e38b0fe2..d5239cc0fdf67 100644
--- a/mlir/include/mlir/Dialect/XeGPU/uArch/uArchBase.h
+++ b/mlir/include/mlir/Dialect/XeGPU/uArch/uArchBase.h
@@ -14,12 +14,6 @@
#ifndef MLIR_DIALECT_XEGPU_UARCH_UARCHBASE_H
#define MLIR_DIALECT_XEGPU_UARCH_UARCHBASE_H
-#include <any>
-#include <functional>
-#include <iostream>
-#include <map>
-#include <mutex>
-#include <optional>
#include <shared_mutex>
#include <tuple>
@@ -31,8 +25,6 @@ namespace mlir {
namespace xegpu {
namespace uArch {
-constexpr unsigned generalPackedFormatBitSize{32};
-
// An enum class to represent the scope of an instruction
enum class InstructionScope { Lane, Subgroup, Workgroup, Cluster };
enum class InstructionKind {
@@ -82,18 +74,10 @@ struct Instruction {
llvm_unreachable("Unknown InstructionKind");
}
- static std::optional<InstructionKind>
- parseInstructionKind(llvm::StringRef str) {
- if (str.equals_insensitive("dpas"))
- return InstructionKind::SubgroupMatrixMultiplyAcc;
- return std::nullopt;
- }
-
protected:
const InstructionKind instKind; // Specific InstructionKind (e.g., DPAS)
const InstructionScope scope; // scope of the instruction (e.g., lane,
// subgroup, workgroup, cluster)
- // @TODO: Add more fields as needed
};
enum class RegisterFileMode : uint8_t { Small, Large };
@@ -187,7 +171,6 @@ struct uArch {
default:
return "";
}
- llvm_unreachable("Unknown uArch::Kind");
}
static StringRef getUArchDescription(Kind k) {
More information about the Mlir-commits
mailing list