[Mlir-commits] [mlir] [MLIR][XeGPU] Enable `isa<>` check for uarch (PR #204577)
Artem Kroviakov
llvmlistbot at llvm.org
Thu Jun 25 02:21:29 PDT 2026
================
@@ -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) ||
----------------
akroviakov wrote:
This is also true for Xe2, pinging the author @mshahneo.
https://github.com/llvm/llvm-project/pull/204577
More information about the Mlir-commits
mailing list