289 lines
11 KiB
C++
289 lines
11 KiB
C++
//===- uArch.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
|
|
// Base uArch definition for different architectures.
|
|
//
|
|
//
|
|
//===----------------------------------------------------------------------===//
|
|
#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 <shared_mutex>
|
|
#include <tuple>
|
|
|
|
#include "mlir/IR/Types.h"
|
|
#include "llvm/ADT/SmallVector.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 {
|
|
SubgroupMatrixMultiplyAcc, // Dot Product Accumulate Systolic (DPAS) is a
|
|
// matrix multiply-add operation
|
|
SubgroupScaledMatrixMultiplyAcc, // Scaled Matrix Multiply Accumulate is a
|
|
// DPAS with scaling factor applied to
|
|
// operand A or B before multiplication
|
|
Subgroup2DBlockStore, // Subgroup-level 2D block write instruction
|
|
Subgroup2DBlockLoad, // Subgroup-level 2D block load instruction
|
|
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.
|
|
// The primary purpose of the Instruction struct is to provide a generic way to
|
|
// represent information about an instruction and to use this information to
|
|
// generate the uArch. Specifc instruction in a uArch can inherit from this
|
|
// struct and add more fields as needed.
|
|
struct Instruction {
|
|
Instruction(InstructionKind kind, InstructionScope scope)
|
|
: instKind(kind), scope(scope) {}
|
|
|
|
~Instruction() = default;
|
|
// Get methods
|
|
InstructionKind getInstructionKind() const { return instKind; }
|
|
InstructionScope getScope() const { return scope; }
|
|
static llvm::StringRef toString(InstructionKind instKind) {
|
|
switch (instKind) {
|
|
case InstructionKind::SubgroupMatrixMultiplyAcc:
|
|
return "dpas";
|
|
case InstructionKind::SubgroupScaledMatrixMultiplyAcc:
|
|
return "dpas_mx";
|
|
case InstructionKind::Subgroup2DBlockStore:
|
|
return "store_nd";
|
|
case InstructionKind::Subgroup2DBlockLoad:
|
|
return "load_nd";
|
|
case InstructionKind::Subgroup2DBlockPrefetch:
|
|
return "prefetch_nd";
|
|
case InstructionKind::StoreScatter:
|
|
return "store";
|
|
case InstructionKind::LoadGather:
|
|
return "load";
|
|
}
|
|
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 {
|
|
// Constructor
|
|
uArch(StringRef name, StringRef description,
|
|
llvm::ArrayRef<const Instruction *> instructionRegistry)
|
|
: name(name), description(description) {
|
|
for (const Instruction *instr : instructionRegistry)
|
|
this->instructionRegistry[instr->getInstructionKind()] = instr;
|
|
}
|
|
virtual ~uArch() = default;
|
|
StringRef getName() const { return name; }
|
|
StringRef getDescription() const { return description; }
|
|
virtual int getSubgroupSize() const = 0;
|
|
virtual unsigned getGeneralPackedFormatBitSize() const = 0;
|
|
|
|
const Instruction *getInstruction(InstructionKind instKind) const {
|
|
auto it = instructionRegistry.find(instKind);
|
|
assert(it != instructionRegistry.end() &&
|
|
"Instruction not found in registry");
|
|
return it->second;
|
|
}
|
|
|
|
bool isSupportedInstruction(InstructionKind instr) const {
|
|
return instructionRegistry.contains(instr);
|
|
}
|
|
|
|
protected:
|
|
StringRef name;
|
|
StringRef description;
|
|
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
|
|
//===----------------------------------------------------------------------===//
|
|
enum class MMAOpndKind { MatrixA, MatrixB, MatrixC, MatrixD };
|
|
struct MMAInstructionInterface {
|
|
// Get supported Matrix shapes
|
|
virtual llvm::SmallVector<std::pair<uint32_t, uint32_t>, 16>
|
|
getSupportedShapes(Type dataType, MMAOpndKind matrixType) = 0;
|
|
// @TODO: This method takes an context object as a parameter, this is to
|
|
// create the Type objects from the same context. Since type objects are
|
|
// uniqued in a specific context, to do things like "aType == bType" (where
|
|
// aType and bType are both same type) kind of checks, the both types should
|
|
// be from the same context.
|
|
//
|
|
// One alternative to this is to create enum to represent each types, but this
|
|
// adds an extra burden to user to convert these enums to specific types. In
|
|
// fact the utility that would convert enumToType() and vice versa would still
|
|
// have to use the context object.
|
|
//
|
|
// Untill we have a better solution, we stick to passing context object to
|
|
// 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;
|
|
virtual bool isLaneLayoutRowMajorOrder() const = 0;
|
|
virtual ~MMAInstructionInterface() = default;
|
|
};
|
|
|
|
//===----------------------------------------------------------------------===//
|
|
// Common instructions (shared across architectures)
|
|
//===----------------------------------------------------------------------===//
|
|
|
|
struct LoadGatherInstructionInterface : public Instruction {
|
|
LoadGatherInstructionInterface()
|
|
: Instruction(InstructionKind::LoadGather, InstructionScope::Lane) {}
|
|
static bool classof(const Instruction *B) {
|
|
return B->getInstructionKind() == InstructionKind::LoadGather;
|
|
}
|
|
|
|
virtual int32_t getMaxLaneLoadSize(int32_t bitWidth) const = 0;
|
|
virtual ~LoadGatherInstructionInterface() = default;
|
|
};
|
|
|
|
struct StoreScatterInstructionInterface : public Instruction {
|
|
StoreScatterInstructionInterface()
|
|
: Instruction(InstructionKind::StoreScatter, InstructionScope::Lane) {}
|
|
static bool classof(const Instruction *B) {
|
|
return B->getInstructionKind() == InstructionKind::StoreScatter;
|
|
}
|
|
|
|
virtual int32_t getMaxLaneStoreSize(int32_t bitWidth) const = 0;
|
|
virtual ~StoreScatterInstructionInterface() = default;
|
|
};
|
|
|
|
} // namespace uArch
|
|
} // namespace xegpu
|
|
} // namespace mlir
|
|
|
|
#endif // MLIR_DIALECT_XEGPU_UARCH_UARCHBASE_H
|