[Mlir-commits] [mlir] bda6db4 - [MLIR][XeGPU] Enable `isa<>` check for uarch (#204577)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Fri Jun 26 01:58:08 PDT 2026
Author: Artem Kroviakov
Date: 2026-06-26T10:58:03+02:00
New Revision: bda6db4a9bb54d10d269a39a26e8fb4263b751a4
URL: https://github.com/llvm/llvm-project/commit/bda6db4a9bb54d10d269a39a26e8fb4263b751a4
DIFF: https://github.com/llvm/llvm-project/commit/bda6db4a9bb54d10d269a39a26e8fb4263b751a4.diff
LOG: [MLIR][XeGPU] Enable `isa<>` check for uarch (#204577)
Added:
mlir/include/mlir/Dialect/XeGPU/uArch/IntelGpuXe3.h
mlir/include/mlir/Dialect/XeGPU/uArch/uArchCommon.h
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/Conversion/XeGPUToXeVM/XeGPUToXeVM.cpp
mlir/lib/Dialect/XeGPU/IR/XeGPUDialect.cpp
mlir/lib/Dialect/XeGPU/Transforms/XeGPUArrayLengthOptimization.cpp
mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
mlir/lib/Dialect/XeGPU/Transforms/XeGPUPeepHoleOptimizer.cpp
mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp
Removed:
################################################################################
diff --git a/mlir/include/mlir/Dialect/XeGPU/Transforms/XeGPULayoutImpl.h b/mlir/include/mlir/Dialect/XeGPU/Transforms/XeGPULayoutImpl.h
index bf61f0ced6048..5f313e480c918 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"
@@ -256,7 +256,7 @@ DistributeLayoutAttr setupStoreMatrixAnchorLayout(LayoutKind layoutKind,
std::optional<DistributeLayoutAttr> completeScatterLoadLaneLayoutFromInstData(
DistributeLayoutAttr userSpecifiedLayout,
DistributeLayoutAttr consumerLayout, Type elemTy,
- const xegpu::uArch::LoadGatherInstructionInterface *uArchInstruction,
+ const xegpu::uArch::LoadGatherInstruction *uArchInstruction,
const int subgroupSize);
/// Like completeScatterLoadLaneLayoutFromInstData, but for scatter stores
@@ -265,7 +265,7 @@ std::optional<DistributeLayoutAttr> completeScatterLoadLaneLayoutFromInstData(
/// with no consumer layout to reuse.
std::optional<DistributeLayoutAttr> completeScatterStoreLaneLayoutFromInstData(
DistributeLayoutAttr specifiedLayout, Type elemTy,
- const xegpu::uArch::StoreScatterInstructionInterface *uArchInstruction,
+ const xegpu::uArch::StoreScatterInstruction *uArchInstruction,
const int subgroupSize);
/// Completes a user-provided 2D-block store_nd / prefetch_nd anchor that has
diff --git a/mlir/include/mlir/Dialect/XeGPU/uArch/IntelGpuXe2.h b/mlir/include/mlir/Dialect/XeGPU/uArch/IntelGpuXe2.h
index ff80a77b28d37..ae042e71e362f 100644
--- a/mlir/include/mlir/Dialect/XeGPU/uArch/IntelGpuXe2.h
+++ b/mlir/include/mlir/Dialect/XeGPU/uArch/IntelGpuXe2.h
@@ -16,681 +16,63 @@
#define MLIR_DIALECT_XEGPU_UARCH_INTELGPUXE2_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 Xe2Plus : public uArch {
- Xe2Plus(StringRef archName, StringRef archDescription,
- llvm::ArrayRef<const Instruction *> instructionRegistry,
- const XeCoreInfo &xeCore)
- : uArch(archName, archDescription, instructionRegistry), xeCore(xeCore) {}
+struct Xe2 : public uArch {
+ Xe2(Kind kind, llvm::ArrayRef<const Instruction *> instructionRegistry)
+ : uArch(kind, instructionRegistry) {}
int getSubgroupSize() const override { return 16; }
unsigned getGeneralPackedFormatBitSize() const override { return 32; }
-protected:
- XeCoreInfo xeCore;
-};
-
-//===----------------------------------------------------------------------===//
-// uArch instructions
-//===----------------------------------------------------------------------===//
-struct Subgroup2DBlockStoreInstruction : public Instruction,
- public BlockIOInstructionInterface {
- Subgroup2DBlockStoreInstruction()
- : 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
- // Stores ignore the transform / transpose / upConv flags.
- 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 override {
- 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 override { return 16; }
-};
-
-struct Subgroup2DBlockLoadInstruction : public Instruction,
- public BlockIOInstructionInterface {
- Subgroup2DBlockLoadInstruction()
- : 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 = 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 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};
- // (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 = {
- {{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:
- {{8, true, false, false}, {kWidth16, kHeight32, kCount4}},
- {{16, true, false, false}, {kWidth16, kHeightAtLeast16, kCount2}},
- // Block Loads with Transpose:
- {{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;
+ static bool classof(const uArch *u) {
+ return u->getKind() >= Kind::Xe2_First && u->getKind() <= Kind::Xe2_Last;
}
-
- int32_t getPackedFormatBitSize() const override { return 16; }
-};
-
-struct Subgroup2DBlockPrefetchInstruction : public Instruction,
- public BlockIOInstructionInterface {
- Subgroup2DBlockPrefetchInstruction()
- : 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
- // Prefetches ignore the transform / transpose / upConv flags.
- 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 override {
- 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 override { return 16; }
-};
-
-struct SubgroupMatrixMultiplyAcc : public Instruction,
- public MMAInstructionInterface {
- SubgroupMatrixMultiplyAcc(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 SubgroupScaledMatrixMultiplyAcc : public Instruction,
- public MMAInstructionInterface {
- SubgroupScaledMatrixMultiplyAcc(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;
-};
-
-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
+//
+// PVC and BMG share the same Khronos-extension instruction set.
//===----------------------------------------------------------------------===//
-struct PVCuArch final : public Xe2Plus {
- static llvm::ArrayRef<const Instruction *> getInstructionRegistryArr() {
- static const SubgroupMatrixMultiplyAcc dpasInst{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, &loadNdInst,
- &storeNdInst, &prefetchNdInst,
- &storeScatterInst, &loadGatherInst};
- return arr;
- }
-
- PVCuArch()
- : Xe2Plus("pvc", // archName
- "Ponte Vecchio Architecture", // archDescription
- getInstructionRegistryArr(),
- XeCoreInfo(8, SharedMemory(512 * 1024, 4), 8, 8) // xeCore
- ) {}
+namespace detail {
+inline llvm::ArrayRef<const Instruction *> getXe2InstructionRegistry() {
+ static const SubgroupMatrixMultiplyAcc dpasInst{16, 32};
+ static const Subgroup2DBlockLoadInstruction loadNdInst;
+ static const Subgroup2DBlockStoreInstruction storeNdInst;
+ static const Subgroup2DBlockPrefetchInstruction prefetchNdInst;
+ static const StoreScatterInstruction storeScatterInst;
+ static const LoadGatherInstruction loadGatherInst;
+ static const Instruction *arr[] = {&dpasInst, &loadNdInst,
+ &storeNdInst, &prefetchNdInst,
+ &storeScatterInst, &loadGatherInst};
+ return arr;
+}
+} // namespace detail
+
+struct PVCuArch final : public Xe2 {
+ PVCuArch() : Xe2(Kind::PVC, detail::getXe2InstructionRegistry()) {}
+ 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;
}
};
-struct BMGuArch : public Xe2Plus {
- static llvm::ArrayRef<const Instruction *> getInstructionRegistryArr() {
- static const SubgroupMatrixMultiplyAcc dpasInst{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, &loadNdInst,
- &storeNdInst, &prefetchNdInst,
- &storeScatterInst, &loadGatherInst};
- return arr;
- }
-
- BMGuArch()
- : Xe2Plus("bmg", // archName
- "Battlemage Architecture", // archDescription
- getInstructionRegistryArr(),
- XeCoreInfo(8, SharedMemory(256 * 1024, 4), 8, 8) // xeCore
- ) {}
+struct BMGuArch final : public Xe2 {
+ BMGuArch() : Xe2(Kind::BMG, detail::getXe2InstructionRegistry()) {}
+ 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;
}
};
-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("cri", // archName
- "Crescent Island Architecture", // archDescription
- getInstructionRegistryArr(),
- // Using bmg config as placeholder
- // TODO: Update to actual XeCore and SharedMemory config
- XeCoreInfo(8, SharedMemory(256 * 1024, 4), 8, 8) // xeCore
- ) {}
- 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
-//===----------------------------------------------------------------------===//
-// Instruction implementations
-//===----------------------------------------------------------------------===//
-
-inline llvm::SmallVector<std::pair<uint32_t, uint32_t>, 16>
-SubgroupMatrixMultiplyAcc::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>
-SubgroupMatrixMultiplyAcc::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 SubgroupMatrixMultiplyAcc::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 SubgroupMatrixMultiplyAcc::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 SubgroupMatrixMultiplyAcc::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>
-SubgroupMatrixMultiplyAcc::getSupportedM(Type type) const {
- return {1, 2, 3, 4, 5, 6, 7, 8};
-}
-
-inline llvm::SmallVector<uint32_t, 8>
-SubgroupMatrixMultiplyAcc::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>
-SubgroupMatrixMultiplyAcc::getSupportedN(Type type) const {
- return {16};
-}
-
-//===----------------------------------------------------------------------===//
-// SubgroupScaledMatrixMultiplyAcc implementations
-//===----------------------------------------------------------------------===//
-
-inline llvm::SmallVector<std::pair<uint32_t, uint32_t>, 16>
-SubgroupScaledMatrixMultiplyAcc::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>
-SubgroupScaledMatrixMultiplyAcc::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 SubgroupScaledMatrixMultiplyAcc::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 SubgroupScaledMatrixMultiplyAcc::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 SubgroupScaledMatrixMultiplyAcc::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>
-SubgroupScaledMatrixMultiplyAcc::getSupportedM(Type type) const {
- return {8};
-}
-
-inline llvm::SmallVector<uint32_t, 8>
-SubgroupScaledMatrixMultiplyAcc::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>
-SubgroupScaledMatrixMultiplyAcc::getSupportedN(Type type) const {
- return {16};
-}
-
#endif // MLIR_DIALECT_XEGPU_UARCH_INTELGPUXE2_H
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..6cd777519d950
--- /dev/null
+++ b/mlir/include/mlir/Dialect/XeGPU/uArch/IntelGpuXe3.h
@@ -0,0 +1,68 @@
+//===--- 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
+// Xe3 uArch definition. Xe3 is the third generation of Intel Xe GPUs and
+// includes Crescent Island (CRI). The base instruction set comes from the
+// shared Khronos OpenCL extensions defined in uArchBase.h; Xe3 only adds the
+// scaled DPAS extension. Subclass and override here only when a CRI-specific
+// instruction diverges from the SPIRV defaults.
+//
+//===----------------------------------------------------------------------===//
+#ifndef MLIR_DIALECT_XEGPU_UARCH_INTELGPUXE3_H
+#define MLIR_DIALECT_XEGPU_UARCH_INTELGPUXE3_H
+
+#include "mlir/Dialect/XeGPU/uArch/uArchBase.h"
+
+namespace mlir {
+namespace xegpu {
+namespace uArch {
+
+struct Xe3 : public uArch {
+ Xe3(Kind kind, llvm::ArrayRef<const Instruction *> instructionRegistry)
+ : uArch(kind, instructionRegistry) {}
+ 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;
+ }
+};
+
+//===----------------------------------------------------------------------===//
+// uArch instances
+//===----------------------------------------------------------------------===//
+
+struct CRIuArch final : public Xe3 {
+ static llvm::ArrayRef<const Instruction *> getCriInstructionRegistry() {
+ 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 StoreScatterInstruction storeScatterInst;
+ static const LoadGatherInstruction loadGatherInst;
+ static const Instruction *arr[] = {
+ &dpasInst, &dpasMxInst, &loadNdInst, &storeNdInst,
+ &prefetchNdInst, &storeScatterInst, &loadGatherInst};
+ return arr;
+ }
+
+ CRIuArch() : Xe3(Kind::CRI, getCriInstructionRegistry()) {}
+ static bool classof(const uArch *u) { return u->getKind() == Kind::CRI; }
+ static const uArch *getInstance() {
+ static const CRIuArch instance;
+ return &instance;
+ }
+};
+
+} // namespace uArch
+} // namespace xegpu
+} // namespace mlir
+
+#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 61db4605e85fa..e1879bb3ffe1c 100644
--- a/mlir/include/mlir/Dialect/XeGPU/uArch/uArchBase.h
+++ b/mlir/include/mlir/Dialect/XeGPU/uArch/uArchBase.h
@@ -7,31 +7,33 @@
//===----------------------------------------------------------------------===//
//
// \file
-// Base uArch definition for
diff erent architectures.
-//
+// Base uArch definition for
diff erent architectures, plus the SPIRV / Khronos
+// OpenCL extension instruction defaults shared across Intel Xe uArchs.
//
//===----------------------------------------------------------------------===//
#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 <cassert>
#include <optional>
-#include <shared_mutex>
#include <tuple>
+#include <utility>
+#include "mlir/IR/BuiltinTypes.h"
+#include "mlir/IR/TypeUtilities.h"
#include "mlir/IR/Types.h"
+#include "llvm/ADT/DenseMap.h"
+#include "llvm/ADT/STLExtras.h"
#include "llvm/ADT/SmallVector.h"
+#include "llvm/ADT/StringRef.h"
+#include "llvm/Support/Casting.h"
+#include "llvm/Support/DebugLog.h"
+#include "llvm/Support/ErrorHandling.h"
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 {
@@ -45,7 +47,6 @@ enum class InstructionKind {
Subgroup2DBlockPrefetch, // Subgroup-level 2D block prefetch instruction
StoreScatter, // Lane-level store (scalar, vector)
LoadGather, // Lane-level load (scalar, vector)
- // @TODO: Add more instructions as needed
};
// A struct to represent basic information about an instruction.
@@ -81,88 +82,33 @@ 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 };
-enum class RegisterFileType : uint8_t { GRF, ARF };
-
-// A struct to represent register file information
-struct RegisterFileInfo {
- // Constructor
- RegisterFileInfo() = default;
- RegisterFileInfo(uint32_t size,
- const llvm::SmallVector<RegisterFileMode, 4> &mode,
- const llvm::SmallVector<uint32_t, 4> &numRegs)
- : size(size), mode(mode), numRegsPerThreadPerMode(numRegs) {}
-
- // Get methods
- uint32_t getSize() const { return size; }
-
- const llvm::SmallVector<RegisterFileMode, 4> &getModes() const {
- return mode;
- }
-
- const llvm::SmallVector<uint32_t, 4> &getNumRegsPerThreadPerMode() const {
- return numRegsPerThreadPerMode;
- }
-
-protected:
- uint32_t size; // size per register in bits
- llvm::SmallVector<RegisterFileMode, 4>
- mode; // e.g., "small", "large" GRF modes
- llvm::SmallVector<uint32_t, 4>
- numRegsPerThreadPerMode; // number of registers per thread per mode
-};
-
-enum class CacheHierarchyLevel { L1 = 1, L2 = 2, L3 = 3 };
-
-// A struct to represent cache information
-struct CacheInfo {
- // Constructor
- CacheInfo() = default;
- CacheInfo(uint32_t size, uint32_t line_size,
- CacheHierarchyLevel hierarchy_level)
- : size(size), line_size(line_size), hierarchy_level(hierarchy_level) {}
-
- virtual ~CacheInfo() = default;
-
- // Get methods
- uint32_t getSize() const { return size; }
- uint32_t getLineSize() const { return line_size; }
- CacheHierarchyLevel getHierarchyLevel() const { return hierarchy_level; }
-
-protected:
- uint32_t size;
- uint32_t line_size;
- CacheHierarchyLevel hierarchy_level;
- // @TODO: Add more fields as needed (e.g., associativity, num_banks,
- // bank_size, num_ports, port_width, bank_conflicts, hierarchy_level,
- // latency, throughput, bandwidth)
};
struct uArch {
+ enum class Kind {
+ // Xe2 family
+ Xe2_First,
+ PVC = Xe2_First,
+ BMG,
+ Xe2_Last = BMG,
+ Xe3_First,
+ CRI = Xe3_First,
+ Xe3_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; }
+
virtual int getSubgroupSize() const = 0;
virtual unsigned getGeneralPackedFormatBitSize() const = 0;
@@ -178,41 +124,11 @@ struct uArch {
}
protected:
- StringRef name;
- StringRef description;
+ Kind kind;
llvm::SmallDenseMap<InstructionKind, const Instruction *, 32>
instructionRegistry;
};
-// A struct to represent shared memory information
-struct SharedMemory {
- // Constructor
- SharedMemory(uint32_t size, uint32_t alignment)
- : size(size), alignment(alignment) {}
-
- // Get methods
- uint32_t getSize() const { return size; }
- uint32_t getAlignment() const { return alignment; }
-
-protected:
- uint32_t size; // in bytes
- uint32_t alignment; // in bytes
- // @TODO: Add more fields as needed (e.g., latency, throughput, bandwidth)
-};
-
-struct XeCoreInfo {
- uint32_t num_threads;
- SharedMemory shared_memory;
- uint32_t num_vector_units;
- uint32_t num_matrix_units;
-
- XeCoreInfo(uint32_t num_threads, const SharedMemory &shared_memory,
- uint32_t num_vector_units, uint32_t num_matrix_units)
- : num_threads(num_threads), shared_memory(shared_memory),
- num_vector_units(num_vector_units), num_matrix_units(num_matrix_units) {
- }
-};
-
//===----------------------------------------------------------------------===//
// Interfaces
//===----------------------------------------------------------------------===//
@@ -236,19 +152,7 @@ struct MMAInstructionInterface {
// this method.
virtual llvm::SmallVector<Type, 8>
getSupportedTypes(MLIRContext &context, MMAOpndKind matrixType) = 0;
- 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) = 0;
- virtual bool checkSupportedTypes(Type AType, Type BType, Type CType,
- Type DType) = 0;
- 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) = 0;
+
virtual llvm::SmallVector<uint32_t, 8> getSupportedM(Type type) const = 0;
virtual llvm::SmallVector<uint32_t, 8> getSupportedK(Type type) const = 0;
virtual llvm::SmallVector<uint32_t, 8> getSupportedN(Type type) const = 0;
@@ -262,44 +166,426 @@ struct MMAInstructionInterface {
// transform / transpose / upConv flags are only meaningful for loads; store
// and prefetch implementations ignore them.
struct BlockIOInstructionInterface {
+ using BlockShapes =
+ std::tuple<llvm::ArrayRef<int>, llvm::ArrayRef<int>, llvm::ArrayRef<int>>;
+
// 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>>>
+ std::optional<BlockShapes>
getBlockWidthHeightCount(Type elemTy, bool hasTransform = false,
bool hasTranspose = false,
- bool upConv = false) const = 0;
+ bool upConv = false) const {
+ return computeBlockWidthHeightCount(elemTy, hasTransform, hasTranspose,
+ upConv);
+ }
+
// Bit size of the packed format used by this block instruction.
virtual int32_t getPackedFormatBitSize() const = 0;
virtual ~BlockIOInstructionInterface() = default;
+
+protected:
+ virtual std::optional<BlockShapes>
+ computeBlockWidthHeightCount(Type elemTy, bool hasTransform,
+ bool hasTranspose, bool upConv) const = 0;
};
//===----------------------------------------------------------------------===//
-// Common instructions (shared across architectures)
+// Common virtual ISA instructions (shared across architectures)
//===----------------------------------------------------------------------===//
-struct LoadGatherInstructionInterface : public Instruction {
- LoadGatherInstructionInterface()
- : Instruction(InstructionKind::LoadGather, InstructionScope::Lane) {}
+//===----------------------------------------------------------------------===//
+// SPIRV
+//===----------------------------------------------------------------------===//
+template <InstructionKind Kind>
+struct ScatterIoInstructionInterface : public Instruction {
+ static_assert(Kind == InstructionKind::LoadGather ||
+ Kind == InstructionKind::StoreScatter,
+ "ScatterIO only supports LoadGather / StoreScatter");
+
+ ScatterIoInstructionInterface() : Instruction(Kind, InstructionScope::Lane) {}
+
static bool classof(const Instruction *B) {
- return B->getInstructionKind() == InstructionKind::LoadGather;
+ return B->getInstructionKind() == Kind;
}
- virtual int32_t getMaxLaneLoadSize(int32_t bitWidth) const = 0;
- virtual ~LoadGatherInstructionInterface() = default;
+ virtual int32_t getMaxLaneAccessSizeBytes() const = 0;
+ virtual ~ScatterIoInstructionInterface() = default;
+};
+struct LoadGatherInstruction
+ : public ScatterIoInstructionInterface<InstructionKind::LoadGather> {
+ int32_t getMaxLaneAccessSizeBytes() const override { return 16; }
+};
+
+struct StoreScatterInstruction
+ : public ScatterIoInstructionInterface<InstructionKind::StoreScatter> {
+ int32_t getMaxLaneAccessSizeBytes() const override { return 16; }
};
-struct StoreScatterInstructionInterface : public Instruction {
- StoreScatterInstructionInterface()
- : Instruction(InstructionKind::StoreScatter, InstructionScope::Lane) {}
+//===----------------------------------------------------------------------===//
+// SPIRV / OpenCL-extension subgroup instructions
+//
+// These come from cl_intel_subgroup_2d_block_io and
+// cl_intel_subgroup_matrix_multiply_accumulate. A uArch only needs to
+// subclass when it diverges from the extension defaults.
+//===----------------------------------------------------------------------===//
+
+struct Subgroup2DBlockStoreInstruction : public Instruction,
+ public BlockIOInstructionInterface {
+ Subgroup2DBlockStoreInstruction()
+ : 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
+ // Stores ignore the transform / transpose / upConv flags.
+ int32_t getPackedFormatBitSize() const override { return 16; }
+
+protected:
+ std::optional<BlockShapes>
+ computeBlockWidthHeightCount(Type elemTy, bool /*hasTransform*/,
+ bool /*hasTranspose*/,
+ bool /*upConv*/) const override {
+ static const int kHeight[] = {1, 2, 4, 8};
+ static const int kWidth16[] = {16};
+ static const int kCount[] = {1};
+ const int elemByteSize = elemTy.getIntOrFloatBitWidth() / 8;
+ if (elemByteSize == 1 || elemByteSize == 2 || elemByteSize == 4)
+ return std::make_tuple(llvm::ArrayRef<int>(kWidth16),
+ llvm::ArrayRef<int>(kHeight),
+ llvm::ArrayRef<int>(kCount));
+ return std::nullopt;
+ }
+};
+
+struct Subgroup2DBlockLoadInstruction : public Instruction,
+ public BlockIOInstructionInterface {
+ Subgroup2DBlockLoadInstruction()
+ : 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
+ int32_t getPackedFormatBitSize() const override { return 16; }
+
+protected:
+ std::optional<BlockShapes>
+ computeBlockWidthHeightCount(Type elemTy, bool hasTransform,
+ bool hasTranspose, bool upConv) 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 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};
+ // (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 = {
+ {{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:
+ {{8, true, false, false}, {kWidth16, kHeight32, kCount4}},
+ {{16, true, false, false}, {kWidth16, kHeightAtLeast16, kCount2}},
+ // Block Loads with Transpose:
+ {{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;
+ }
+};
+
+struct Subgroup2DBlockPrefetchInstruction : public Instruction,
+ public BlockIOInstructionInterface {
+ Subgroup2DBlockPrefetchInstruction()
+ : 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
+ // Prefetches ignore the transform / transpose / upConv flags.
+ int32_t getPackedFormatBitSize() const override { return 16; }
+
+protected:
+ std::optional<BlockShapes>
+ computeBlockWidthHeightCount(Type elemTy, bool /*hasTransform*/,
+ bool /*hasTranspose*/,
+ bool /*upConv*/) const override {
+ 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;
+ }
+};
+
+struct SubgroupMatrixMultiplyAcc : public Instruction,
+ public MMAInstructionInterface {
+ SubgroupMatrixMultiplyAcc(unsigned packedFormatBitSizeA,
+ unsigned packedFormatBitSizeB)
+ : Instruction(InstructionKind::SubgroupMatrixMultiplyAcc,
+ InstructionScope::Subgroup),
+ packedFormatBitSizeA(packedFormatBitSizeA),
+ packedFormatBitSizeB(packedFormatBitSizeB) {}
static bool classof(const Instruction *B) {
- return B->getInstructionKind() == InstructionKind::StoreScatter;
+ return B->getInstructionKind() ==
+ InstructionKind::SubgroupMatrixMultiplyAcc;
}
+ // Source:
+ // https://registry.khronos.org/OpenCL/extensions/intel/cl_intel_subgroup_matrix_multiply_accumulate.html
+
+ llvm::SmallVector<std::pair<uint32_t, uint32_t>, 16>
+ getSupportedShapes(Type dataType, MMAOpndKind matrixType) override;
+ llvm::SmallVector<Type, 8> getSupportedTypes(MLIRContext &context,
+ MMAOpndKind matrixType) override;
+
+ llvm::SmallVector<uint32_t, 8> getSupportedM(Type type) const override;
+ llvm::SmallVector<uint32_t, 8> getSupportedK(Type type) const override;
+ 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 SubgroupScaledMatrixMultiplyAcc : public Instruction,
+ public MMAInstructionInterface {
+ SubgroupScaledMatrixMultiplyAcc(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
+
+ llvm::SmallVector<std::pair<uint32_t, uint32_t>, 16>
+ getSupportedShapes(Type dataType, MMAOpndKind matrixType) override;
+ llvm::SmallVector<Type, 8> getSupportedTypes(MLIRContext &context,
+ MMAOpndKind matrixType) override;
+
+ llvm::SmallVector<uint32_t, 8> getSupportedM(Type type) const override;
+ llvm::SmallVector<uint32_t, 8> getSupportedK(Type type) const override;
+ llvm::SmallVector<uint32_t, 8> getSupportedN(Type type) const override;
- virtual int32_t getMaxLaneStoreSize(int32_t bitWidth) const = 0;
- virtual ~StoreScatterInstructionInterface() = default;
+ unsigned getPackedFormatBitSizeA() const { return packedFormatBitSizeA; }
+ unsigned getPackedFormatBitSizeB() const { return packedFormatBitSizeB; }
+ bool isLaneLayoutRowMajorOrder() const override { return true; }
+
+protected:
+ const unsigned packedFormatBitSizeA;
+ const unsigned packedFormatBitSizeB;
};
+//===----------------------------------------------------------------------===//
+// Inline implementations
+//===----------------------------------------------------------------------===//
+
+namespace util {
+inline llvm::SmallVector<std::pair<uint32_t, uint32_t>, 16>
+crossProduct(const llvm::SmallVector<uint32_t, 8> &a,
+ const llvm::SmallVector<uint32_t, 8> &b) {
+ 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;
+}
+} // namespace util
+
+inline llvm::SmallVector<std::pair<uint32_t, uint32_t>, 16>
+SubgroupMatrixMultiplyAcc::getSupportedShapes(Type dataType,
+ MMAOpndKind matrixType) {
+ auto M = getSupportedM(dataType);
+ auto K = getSupportedK(dataType);
+ auto N = getSupportedN(dataType);
+ switch (matrixType) {
+ case MMAOpndKind::MatrixA:
+ return util::crossProduct(M, K);
+ case MMAOpndKind::MatrixB:
+ return util::crossProduct(K, N);
+ case MMAOpndKind::MatrixC:
+ case MMAOpndKind::MatrixD:
+ return util::crossProduct(M, N);
+ }
+ return {};
+}
+
+inline llvm::SmallVector<Type, 8>
+SubgroupMatrixMultiplyAcc::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:
+ case MMAOpndKind::MatrixB:
+ return {bf16Type, f16Type, tf32Type};
+ case MMAOpndKind::MatrixC:
+ case MMAOpndKind::MatrixD:
+ return {bf16Type, f16Type, f32Type};
+ }
+ return {};
+}
+
+inline llvm::SmallVector<uint32_t, 8>
+SubgroupMatrixMultiplyAcc::getSupportedM(Type type) const {
+ return {1, 2, 3, 4, 5, 6, 7, 8};
+}
+
+inline llvm::SmallVector<uint32_t, 8>
+SubgroupMatrixMultiplyAcc::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;
+ 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>
+SubgroupMatrixMultiplyAcc::getSupportedN(Type type) const {
+ return {16};
+}
+
+inline llvm::SmallVector<std::pair<uint32_t, uint32_t>, 16>
+SubgroupScaledMatrixMultiplyAcc::getSupportedShapes(Type dataType,
+ MMAOpndKind matrixType) {
+ // 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 util::crossProduct(getSupportedM(dataType), getSupportedK(dataType));
+ case MMAOpndKind::MatrixB:
+ return util::crossProduct(getSupportedK(dataType), getSupportedN(dataType));
+ case MMAOpndKind::MatrixC:
+ case MMAOpndKind::MatrixD:
+ return util::crossProduct(getSupportedM(dataType), getSupportedN(dataType));
+ }
+ return {};
+}
+
+inline llvm::SmallVector<Type, 8>
+SubgroupScaledMatrixMultiplyAcc::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:
+ case MMAOpndKind::MatrixB:
+ return {f8E4M3FNType, f8E5M2Type, f4E2M1FNType};
+ case MMAOpndKind::MatrixC:
+ case MMAOpndKind::MatrixD:
+ return {bf16Type, f32Type};
+ }
+ return {};
+}
+
+inline llvm::SmallVector<uint32_t, 8>
+SubgroupScaledMatrixMultiplyAcc::getSupportedM(Type type) const {
+ return {8};
+}
+
+inline llvm::SmallVector<uint32_t, 8>
+SubgroupScaledMatrixMultiplyAcc::getSupportedK(Type type) const {
+ assert(type.isIntOrFloat() && "Matrix type must be int or float");
+ auto bitWidth = type.getIntOrFloatBitWidth();
+ switch (bitWidth) {
+ case 4:
+ return {64}; // FP4: scale K by 4 (base 16-bit K=16 -> 64)
+ case 8:
+ return {32}; // FP8: scale K by 2 (base 16-bit K=16 -> 32)
+ 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 {};
+ }
+}
+
+inline llvm::SmallVector<uint32_t, 8>
+SubgroupScaledMatrixMultiplyAcc::getSupportedN(Type type) const {
+ 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/XeGPULayoutImpl.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
index 574eeabac1836..61cd253508357 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
@@ -1710,13 +1710,11 @@ xegpu::DistributeLayoutAttr xegpu::setupLoadGatherAnchorLayout(
const int subgroupSize = uArch->getSubgroupSize();
ArrayRef<int64_t> resShape = resVecTy.getShape();
auto context = resVecTy.getContext();
- auto elemBitWidth = resVecTy.getElementType().getIntOrFloatBitWidth();
- const auto *uArchInstruction =
- dyn_cast<xegpu::uArch::LoadGatherInstructionInterface>(
- uArch->getInstruction(xegpu::uArch::InstructionKind::LoadGather));
- int maxChunkSize = std::min(
- uArchInstruction->getMaxLaneLoadSize(elemBitWidth), contigChunkSize);
+ const auto *uArchInstruction = dyn_cast<xegpu::uArch::LoadGatherInstruction>(
+ uArch->getInstruction(xegpu::uArch::InstructionKind::LoadGather));
+ int maxChunkSize =
+ std::min(uArchInstruction->getMaxLaneAccessSizeBytes(), contigChunkSize);
return setupGenericLoadAnchorLayout(layoutKind, context, consumerLayout,
maxChunkSize, resShape, subgroupSize);
@@ -1733,13 +1731,11 @@ xegpu::setupLoadMatrixAnchorLayout(xegpu::LayoutKind layoutKind,
const int subgroupSize = uArch->getSubgroupSize();
ArrayRef<int64_t> resShape = resVecTy.getShape();
auto context = resVecTy.getContext();
- auto elemBitWidth = resVecTy.getElementType().getIntOrFloatBitWidth();
- const auto *uArchInstruction =
- dyn_cast<xegpu::uArch::LoadGatherInstructionInterface>(
- uArch->getInstruction(xegpu::uArch::InstructionKind::LoadGather));
- int maxChunkSize = std::min(
- uArchInstruction->getMaxLaneLoadSize(elemBitWidth), contigChunkSize);
+ const auto *uArchInstruction = dyn_cast<xegpu::uArch::LoadGatherInstruction>(
+ uArch->getInstruction(xegpu::uArch::InstructionKind::LoadGather));
+ int maxChunkSize =
+ std::min(uArchInstruction->getMaxLaneAccessSizeBytes(), contigChunkSize);
return setupGenericLoadAnchorLayout(layoutKind, context, consumerLayout,
maxChunkSize, resShape, subgroupSize);
}
@@ -1786,13 +1782,12 @@ xegpu::setupStoreScatterAnchorLayout(xegpu::LayoutKind layoutKind,
const int subgroupSize = uArch->getSubgroupSize();
ArrayRef<int64_t> srcShape = srcVecTy.getShape();
auto context = srcVecTy.getContext();
- auto elemBitWidth = srcVecTy.getElementType().getIntOrFloatBitWidth();
const auto *uArchInstruction =
- dyn_cast<xegpu::uArch::StoreScatterInstructionInterface>(
+ dyn_cast<xegpu::uArch::StoreScatterInstruction>(
uArch->getInstruction(xegpu::uArch::InstructionKind::StoreScatter));
- int maxChunkSize = std::min(
- uArchInstruction->getMaxLaneStoreSize(elemBitWidth), contigChunkSize);
+ int maxChunkSize =
+ std::min(uArchInstruction->getMaxLaneAccessSizeBytes(), contigChunkSize);
return setupGenericStoreAnchorLayout(layoutKind, context, maxChunkSize,
srcShape, subgroupSize);
}
@@ -1806,13 +1801,12 @@ xegpu::setupStoreMatrixAnchorLayout(xegpu::LayoutKind layoutKind,
const int subgroupSize = uArch->getSubgroupSize();
ArrayRef<int64_t> srcShape = srcVecTy.getShape();
auto context = srcVecTy.getContext();
- auto elemBitWidth = srcVecTy.getElementType().getIntOrFloatBitWidth();
const auto *uArchInstruction =
- dyn_cast<xegpu::uArch::StoreScatterInstructionInterface>(
+ dyn_cast<xegpu::uArch::StoreScatterInstruction>(
uArch->getInstruction(xegpu::uArch::InstructionKind::StoreScatter));
- int maxChunkSize = std::min(
- uArchInstruction->getMaxLaneStoreSize(elemBitWidth), contigChunkSize);
+ int maxChunkSize =
+ std::min(uArchInstruction->getMaxLaneAccessSizeBytes(), contigChunkSize);
return setupGenericStoreAnchorLayout(layoutKind, context, maxChunkSize,
srcShape, subgroupSize);
@@ -1831,13 +1825,13 @@ xegpu::setupStoreMatrixAnchorLayout(xegpu::LayoutKind layoutKind,
/// - 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`).
+/// (`getMaxLaneAccessSizeBytes`).
///
std::optional<xegpu::DistributeLayoutAttr>
xegpu::completeScatterLoadLaneLayoutFromInstData(
xegpu::DistributeLayoutAttr specifiedLayout,
xegpu::DistributeLayoutAttr consumerLayout, Type elemTy,
- const xegpu::uArch::LoadGatherInstructionInterface *uArchInstruction,
+ const xegpu::uArch::LoadGatherInstruction *uArchInstruction,
const int subgroupSize) {
if (!specifiedLayout)
return specifiedLayout;
@@ -1851,8 +1845,7 @@ xegpu::completeScatterLoadLaneLayoutFromInstData(
// 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);
+ int maxChunkSize = uArchInstruction->getMaxLaneAccessSizeBytes();
if (consumerLayout) {
auto consumerLaneLayout = consumerLayout.getEffectiveLaneLayoutAsInt();
auto consumerLaneData = consumerLayout.getEffectiveLaneDataAsInt();
@@ -1876,7 +1869,7 @@ xegpu::completeScatterLoadLaneLayoutFromInstData(
std::optional<xegpu::DistributeLayoutAttr>
xegpu::completeScatterStoreLaneLayoutFromInstData(
xegpu::DistributeLayoutAttr specifiedLayout, Type elemTy,
- const xegpu::uArch::StoreScatterInstructionInterface *uArchInstruction,
+ const xegpu::uArch::StoreScatterInstruction *uArchInstruction,
const int subgroupSize) {
if (!specifiedLayout)
return specifiedLayout;
@@ -1890,8 +1883,7 @@ xegpu::completeScatterStoreLaneLayoutFromInstData(
// 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);
+ int maxChunkSize = uArchInstruction->getMaxLaneAccessSizeBytes();
auto [defLaneLayout, defLaneData] = computeScatterIOLaneLayoutAndData(
specifiedInstData, subgroupSize, maxChunkSize);
if (!isValidLaneLayout(specifiedInstData, defLaneLayout, defLaneData))
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPeepHoleOptimizer.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPeepHoleOptimizer.cpp
index 566d0eb433606..1bf04f1f095a8 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"
@@ -106,8 +106,9 @@ static bool canBeOptimizedForTranspose(xegpu::TensorDescType tdescType) {
/// Check if a tensor desc type can be optimized for transpose, if so return the
/// new optimized tensor desc type with a valid transpose layout.
-static xegpu::TensorDescType tryOptimize(xegpu::TensorDescType tdescType,
- const uArch *targetuArch) {
+static xegpu::TensorDescType
+tryOptimize(xegpu::TensorDescType tdescType,
+ const xegpu::uArch::uArch *targetuArch) {
if (!canBeOptimizedForTranspose(tdescType))
return tdescType;
auto laneData = getMaybeLaneData(tdescType)
@@ -124,8 +125,10 @@ static xegpu::TensorDescType tryOptimize(xegpu::TensorDescType tdescType,
Type newElemTy = IntegerType::get(tdescType.getContext(), newBitWidth);
// Supported shape is the max transpose shape that can be supported by
// hardware that is less than or equal to required shape.
- auto *blockLoadTarget = dyn_cast<Subgroup2DBlockLoadInstruction>(
- targetuArch->getInstruction(InstructionKind::Subgroup2DBlockLoad));
+ auto *blockLoadTarget =
+ dyn_cast<xegpu::uArch::Subgroup2DBlockLoadInstruction>(
+ targetuArch->getInstruction(
+ xegpu::uArch::InstructionKind::Subgroup2DBlockLoad));
auto maybeHWParams = blockLoadTarget->getBlockWidthHeightCount(
newElemTy, /** has transform */ false, /** has transpose */ true);
// If no HW params found, return the original type.
@@ -261,7 +264,7 @@ class XeGPUCreateNdDescOpPattern final
chipStr.value() == "cri") &&
"Expecting target chip to be pvc, bmg or cri for transpose "
"optimization.");
- const uArch *targetuArch = xegpu::uArch::getUArch(chipStr.value());
+ const auto *targetuArch = xegpu::uArch::getUArch(chipStr.value());
auto convertType = tryOptimize(tdescTy, targetuArch);
if (convertType == tdescTy)
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
index 8600492e4bf41..64d0d8063b7ff 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"
@@ -452,7 +452,7 @@ void LayoutInfoPropagation::visitPrefetchNdOp(
ArrayRef<const LayoutInfoLattice *> results) {
LayoutInfo prefetchLayout;
- const uArch *uArch = getUArch(getChipStr(prefetch).value_or(""));
+ const auto *uArch = xegpu::uArch::getUArch(getChipStr(prefetch).value_or(""));
if (!uArch)
return;
xegpu::DistributeLayoutAttr anchorLayout = prefetch.getLayoutAttr();
@@ -518,7 +518,8 @@ void LayoutInfoPropagation::visitVectorMultiReductionOp(
VectorType sourceTy = reduction.getSourceVectorType();
SmallVector<int64_t> reductionDims(reduction.getReductionDims());
- const uArch *uArch = getUArch(xegpu::getChipStr(reduction).value_or(""));
+ const auto *uArch =
+ xegpu::uArch::getUArch(xegpu::getChipStr(reduction).value_or(""));
if (!uArch)
return;
@@ -556,7 +557,8 @@ void LayoutInfoPropagation::visitVectorReductionOp(
ArrayRef<const LayoutInfoLattice *> results) {
VectorType sourceTy = reduction.getSourceVectorType();
- const uArch *uArch = getUArch(xegpu::getChipStr(reduction).value_or(""));
+ const auto *uArch =
+ xegpu::uArch::getUArch(xegpu::getChipStr(reduction).value_or(""));
if (!uArch)
return;
@@ -631,7 +633,7 @@ void LayoutInfoPropagation::visitDpasOp(
LayoutInfo dpasBLayout;
LayoutInfo dpasCDLayout;
- const uArch *uArch = getUArch(getChipStr(dpas).value_or(""));
+ const auto *uArch = xegpu::uArch::getUArch(getChipStr(dpas).value_or(""));
if (!uArch)
return;
VectorType aTy = dpas.getLhsType();
@@ -725,7 +727,7 @@ void LayoutInfoPropagation::visitDpasMxOp(
xegpu::DistributeLayoutAttr anchorLayoutB = dpasMx.getLayoutBAttr();
xegpu::DistributeLayoutAttr anchorLayoutCD = dpasMx.getLayoutCdAttr();
- const uArch *uArch = getUArch(getChipStr(dpasMx).value_or(""));
+ const auto *uArch = xegpu::uArch::getUArch(getChipStr(dpasMx).value_or(""));
if (!uArch)
return;
@@ -866,7 +868,7 @@ void LayoutInfoPropagation::visitStoreNdOp(
xegpu::StoreNdOp store, ArrayRef<LayoutInfoLattice *> operands,
ArrayRef<const LayoutInfoLattice *> results) {
LayoutInfo storeLayout;
- const uArch *uArch = getUArch(getChipStr(store).value_or(""));
+ const auto *uArch = xegpu::uArch::getUArch(getChipStr(store).value_or(""));
if (!uArch)
return;
xegpu::DistributeLayoutAttr anchorLayout = store.getLayoutAttr();
@@ -921,7 +923,7 @@ void LayoutInfoPropagation::visitLoadNdOp(
ArrayRef<const LayoutInfoLattice *> results) {
LayoutInfo loadLayout;
- const uArch *uArch = getUArch(getChipStr(load).value_or(""));
+ const auto *uArch = xegpu::uArch::getUArch(getChipStr(load).value_or(""));
if (!uArch)
return;
LayoutInfo valueLayout = results[0]->getValue();
@@ -1061,7 +1063,8 @@ void LayoutInfoPropagation::visitVectorBitcastOp(
auto consumerLayoutAttr =
dyn_cast<xegpu::DistributeLayoutAttr>(resLayoutInfo.get());
- const uArch *uArch = getUArch(xegpu::getChipStr(bitcast).value_or(""));
+ const auto *uArch =
+ xegpu::uArch::getUArch(xegpu::getChipStr(bitcast).value_or(""));
if (!uArch)
return;
auto requiredResLayoutAttr = setupBitCastResultLayout(
@@ -1095,7 +1098,8 @@ void LayoutInfoPropagation::visitVectorInterleaveOp(
auto consumerLayoutAttr =
dyn_cast<xegpu::DistributeLayoutAttr>(resLayoutInfo.get());
- const uArch *uArch = getUArch(xegpu::getChipStr(interleave).value_or(""));
+ const auto *uArch =
+ xegpu::uArch::getUArch(xegpu::getChipStr(interleave).value_or(""));
if (!uArch)
return;
@@ -1150,8 +1154,8 @@ void LayoutInfoPropagation::visitInsertStridedSliceOp(
auto consumerLayoutAttr =
dyn_cast<xegpu::DistributeLayoutAttr>(resLayoutInfo.get());
- const uArch *uArch =
- getUArch(xegpu::getChipStr(insertStridedSlice).value_or(""));
+ const auto *uArch = xegpu::uArch::getUArch(
+ xegpu::getChipStr(insertStridedSlice).value_or(""));
if (!uArch)
return;
@@ -1174,7 +1178,7 @@ void LayoutInfoPropagation::visitLoadGatherOp(
ArrayRef<const LayoutInfoLattice *> results) {
xegpu::DistributeLayoutAttr requiredAnchorLayoutAttr;
xegpu::DistributeLayoutAttr anchorLayoutAttr = load.getLayoutAttr();
- const uArch *uArch = getUArch(getChipStr(load).value_or(""));
+ const auto *uArch = xegpu::uArch::getUArch(getChipStr(load).value_or(""));
if (!uArch)
return;
VectorType resVecTy = load.getValueType();
@@ -1191,7 +1195,7 @@ void LayoutInfoPropagation::visitLoadGatherOp(
if (layoutKind == xegpu::LayoutKind::InstData &&
!consumerLayoutAttr.getEffectiveLaneLayoutAsInt().empty()) {
const auto uArchInstruction =
- dyn_cast<xegpu::uArch::LoadGatherInstructionInterface>(
+ dyn_cast<xegpu::uArch::LoadGatherInstruction>(
uArch->getInstruction(xegpu::uArch::InstructionKind::LoadGather));
if (!uArchInstruction)
return;
@@ -1238,7 +1242,8 @@ void LayoutInfoPropagation::visitStoreScatterOp(
xegpu::DistributeLayoutAttr requiredAnchorLayoutAttr;
xegpu::DistributeLayoutAttr anchorLayoutAttr = storeScatter.getLayoutAttr();
- const uArch *uArch = getUArch(getChipStr(storeScatter).value_or(""));
+ const auto *uArch =
+ xegpu::uArch::getUArch(getChipStr(storeScatter).value_or(""));
if (!uArch)
return;
VectorType srcVecTy = storeScatter.getValueType();
@@ -1248,9 +1253,8 @@ void LayoutInfoPropagation::visitStoreScatterOp(
requiredAnchorLayoutAttr = anchorLayoutAttr;
if (layoutKind == xegpu::LayoutKind::InstData) {
const auto uArchInstruction =
- dyn_cast<xegpu::uArch::StoreScatterInstructionInterface>(
- uArch->getInstruction(
- xegpu::uArch::InstructionKind::StoreScatter));
+ dyn_cast<xegpu::uArch::StoreScatterInstruction>(uArch->getInstruction(
+ xegpu::uArch::InstructionKind::StoreScatter));
if (!uArchInstruction)
return;
auto completed = xegpu::completeScatterStoreLaneLayoutFromInstData(
@@ -1308,7 +1312,8 @@ void LayoutInfoPropagation::visitLoadMatrixOp(
if (!hasParamsOfLayoutKind(anchorLayout)) {
VectorType resVecTy =
llvm::cast<VectorType>(loadMatrixOp.getRes().getType());
- const uArch *uArch = getUArch(getChipStr(loadMatrixOp).value_or(""));
+ const auto *uArch =
+ xegpu::uArch::getUArch(getChipStr(loadMatrixOp).value_or(""));
if (!uArch)
return;
int chunkSize =
@@ -1326,16 +1331,16 @@ void LayoutInfoPropagation::visitStoreMatrixOp(
xegpu::DistributeLayoutAttr anchorLayoutAttr = storeMatrix.getLayoutAttr();
LayoutInfo layout;
VectorType srcVecTy = llvm::cast<VectorType>(storeMatrix.getData().getType());
- const uArch *uArch = getUArch(getChipStr(storeMatrix).value_or(""));
+ const auto *uArch =
+ xegpu::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));
+ dyn_cast<xegpu::uArch::StoreScatterInstruction>(uArch->getInstruction(
+ xegpu::uArch::InstructionKind::StoreScatter));
if (!uArchInstruction)
return;
auto completed = xegpu::completeScatterStoreLaneLayoutFromInstData(
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
index 75a87f84b3da8..23d3877ea63dd 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"
@@ -152,7 +152,8 @@ struct SgToLaneLoadNd : public OpConversionPattern<xegpu::LoadNdOp> {
if (op.getTensorDescType().getLayout() != layout)
return rewriter.notifyMatchFailure(
op, "conflicting layout attributes on tensor descriptor and anchor");
- auto uArch = getUArch(xegpu::getChipStr(op).value_or(""));
+ const auto *uArch =
+ xegpu::uArch::getUArch(xegpu::getChipStr(op).value_or(""));
if (!uArch)
return rewriter.notifyMatchFailure(
op, "xegpu::LoadNdOp require target attribute attached to "
@@ -259,7 +260,8 @@ struct SgToLaneDpas : public OpConversionPattern<xegpu::DpasOp> {
"lane layout");
// Validate bit widths match uArch packed format requirements
- const uArch *uArch = getUArch(xegpu::getChipStr(op).value_or(""));
+ const auto *uArch =
+ xegpu::uArch::getUArch(xegpu::getChipStr(op).value_or(""));
if (uArch) {
const auto *uArchInstruction =
dyn_cast<xegpu::uArch::SubgroupMatrixMultiplyAcc>(
@@ -543,7 +545,8 @@ struct SgToLaneVectorReduction
// Get the subgroup size from the layout.
int64_t sgSize = layout.getEffectiveLaneLayoutAsInt()[0];
- const uArch *uArch = getUArch(xegpu::getChipStr(op).value_or(""));
+ const auto *uArch =
+ xegpu::uArch::getUArch(xegpu::getChipStr(op).value_or(""));
if (!uArch)
return rewriter.notifyMatchFailure(
op, "xegpu::ReductionOp require target attribute attached to "
@@ -1214,7 +1217,8 @@ struct SgToLaneVectorExtractStridedSlice
return rewriter.notifyMatchFailure(
op, "only single dimension distribution is supported");
int64_t distDim = distributedDims[0];
- const uArch *uArch = getUArch(xegpu::getChipStr(op).value_or(""));
+ const auto *uArch =
+ xegpu::uArch::getUArch(xegpu::getChipStr(op).value_or(""));
if (!uArch)
return rewriter.notifyMatchFailure(
op, "target attribute required to determine subgroup size");
@@ -1418,7 +1422,8 @@ struct SgToLaneVectorInsertStridedSlice
op, "only single dimension distribution is supported");
int64_t destDistDim = destDistributedDims[0];
- const uArch *uArch = getUArch(xegpu::getChipStr(op).value_or(""));
+ const auto *uArch =
+ xegpu::uArch::getUArch(xegpu::getChipStr(op).value_or(""));
if (!uArch)
return rewriter.notifyMatchFailure(
op, "target attribute required to determine subgroup size");
@@ -1681,7 +1686,8 @@ struct SgToLaneDpasMx : public OpConversionPattern<xegpu::DpasMxOp> {
LogicalResult
matchAndRewrite(xegpu::DpasMxOp op, OpAdaptor adaptor,
ConversionPatternRewriter &rewriter) const override {
- const uArch *uArch = getUArch(xegpu::getChipStr(op).value_or(""));
+ const auto *uArch =
+ xegpu::uArch::getUArch(xegpu::getChipStr(op).value_or(""));
if (!uArch)
return failure();
if (!uArch->isSupportedInstruction(
diff --git a/mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp b/mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp
index 735dd93e2ac14..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,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::Xe2>(uArch) && !isa<xegpu::uArch::Xe3>(uArch))
return false;
if (!layout)
return false;
More information about the Mlir-commits
mailing list