MLIR 24.0.0git
AMDGPUToROCDL.cpp File Reference
#include "mlir/Conversion/AMDGPUToROCDL/AMDGPUToROCDL.h"
#include "mlir/Conversion/GPUCommon/GPUCommonPass.h"
#include "mlir/Conversion/LLVMCommon/ConversionTarget.h"
#include "mlir/Conversion/LLVMCommon/Pattern.h"
#include "mlir/Conversion/LLVMCommon/TypeConverter.h"
#include "mlir/Dialect/AMDGPU/IR/AMDGPUDialect.h"
#include "mlir/Dialect/AMDGPU/IR/AMDGPUEnums.h"
#include "mlir/Dialect/LLVMIR/LLVMDialect.h"
#include "mlir/Dialect/LLVMIR/LLVMTypes.h"
#include "mlir/Dialect/LLVMIR/ROCDLDialect.h"
#include "mlir/IR/Attributes.h"
#include "mlir/IR/BuiltinAttributes.h"
#include "mlir/IR/BuiltinTypes.h"
#include "mlir/IR/Matchers.h"
#include "mlir/IR/TypeUtilities.h"
#include "mlir/Pass/Pass.h"
#include "../LLVMCommon/MemRefDescriptor.h"
#include "llvm/ADT/STLExtras.h"
#include "llvm/ADT/TypeSwitch.h"
#include "llvm/Support/AMDGPUAddrSpace.h"
#include "llvm/Support/Casting.h"
#include "llvm/Support/ErrorHandling.h"
#include "llvm/Support/MathExtras.h"
#include <cstdint>
#include <optional>
#include "mlir/Conversion/Passes.h.inc"

Go to the source code of this file.

Classes

class  mlir::impl::ConvertAMDGPUToROCDLPassBase< DerivedT >
struct  SparseWMMAOpInfo
 Returns the rocdl intrinsic corresponding to a SparseWMMA operation swmmac if one exists. More...

Namespaces

namespace  mlir
 Include the generated interface declarations.
namespace  mlir::impl
 Attribute collections provide a dictionary-like interface.

Macros

#define GEN_PASS_DEF_CONVERTAMDGPUTOROCDLPASS

Typedefs

using ScaledMFMAIntrinsic
 If there is a scaled MFMA instruction for the input element types aType and bType, output type destType, problem size M, N, K, and B (number of blocks) on the given chipset, return a tuple consisting of the OperationName of the intrinsic and the type codes that need to be passed to that intrinsic.

Functions

std::unique_ptr<::mlir::Pass > mlir::impl::createConvertAMDGPUToROCDLPass ()
std::unique_ptr<::mlir::Pass > mlir::impl::createConvertAMDGPUToROCDLPass (ConvertAMDGPUToROCDLPassOptions options)
std::unique_ptr<::mlir::Pass > mlir::createConvertAMDGPUToROCDLPass ()
std::unique_ptr<::mlir::Pass > mlir::createConvertAMDGPUToROCDLPass (ConvertAMDGPUToROCDLPassOptions options)
static Value convertUnsignedToInt (ConversionPatternRewriter &rewriter, Location loc, Value val, unsigned width)
 Zero-extend or truncate the unsigned number val to width bits.
static Value convertUnsignedToI32 (ConversionPatternRewriter &rewriter, Location loc, Value val)
 Convert an unsigned number val to i32.
static Value createI32Constant (ConversionPatternRewriter &rewriter, Location loc, int32_t value)
static Value convertUnsignedToI64 (ConversionPatternRewriter &rewriter, Location loc, Value val)
 Convert an unsigned number val to i64.
static Value createI64Constant (ConversionPatternRewriter &rewriter, Location loc, int64_t value)
static Value getLinearIndexI32 (ConversionPatternRewriter &rewriter, Location loc, MemRefDescriptor &memRefDescriptor, ValueRange indices, ArrayRef< int64_t > strides)
 Returns the linear index used to access an element in the memref.
static Value getNumRecords (ConversionPatternRewriter &rewriter, Location loc, MemRefType memrefType, MemRefDescriptor &memrefDescriptor, ArrayRef< int64_t > strides, int64_t elementByteWidth, unsigned numRecordsWidth, bool boundsCheck)
 Compute the contents of the num_records field for a given memref descriptor - that is, the number of bytes that's one element past the greatest possible valid index into the memref.
static Value makeBufferRsrc (ConversionPatternRewriter &rewriter, Location loc, Value basePointer, Value numRecords, bool boundsCheck, const ROCDL::TargetInfo &target, unsigned numRecordsWidth, Value cacheSwizzleStride=nullptr, unsigned addressSpace=8)
static Value packSmallFloatVectorOperand (ConversionPatternRewriter &rewriter, Location loc, Value input, bool allowBf16=true)
 Pack small float vector operands (fp4/fp6/fp8/bf16) into the format expected by scaled matrix multiply intrinsics (MFMA/WMMA).
static Value convertPackedVectorOperand (ConversionPatternRewriter &rewriter, Location loc, Value input, bool allowBf16=true)
 Converts packed vector operands to the expected ROCDL types.
static Value castScaleOperand (ConversionPatternRewriter &rewriter, Location loc, Value input)
 Converts the scaled MFMA/WMMA operands, scalesA and scalesB, from MLIR AMDGPU dialect convention to ROCDL and LLVM AMDGPU intrinsics convention.
static std::optional< ROCDL::WMMAMatrixScaleFormat > getWmmaScaleFormat (Type elemType)
 Maps f8 scale element types to WMMA scale format codes.
static std::optional< StringRef > getScaledWmmaIntrinsicName (int64_t m, int64_t n, int64_t k, bool isScale16)
 Determines the ROCDL intrinsic name for scaled WMMA based on dimensions and scale block size (16 or 32).
static void wmmaPushInputOperand (ConversionPatternRewriter &rewriter, Location loc, const TypeConverter *typeConverter, bool isUnsigned, Value llvmInput, Value mlirInput, SmallVectorImpl< Value > &operands, SmallVectorImpl< NamedAttribute > &attrs, StringRef attrName)
 Push an input operand.
static void wmmaPushOutputOperand (ConversionPatternRewriter &rewriter, Location loc, const TypeConverter *typeConverter, Value output, int32_t subwordOffset, bool clamp, SmallVectorImpl< Value > &operands, SmallVectorImpl< NamedAttribute > &attrs)
 Push the output operand.
static bool typeIsExpectedBf8ForTarget (const ROCDL::TargetInfo &target, Type type)
 Return true if type is the E5M2 variant of an 8-bit float that is supported by the _bf8 instructions on target.
static bool typeIsExpectedFp8ForTarget (const ROCDL::TargetInfo &target, Type type)
 Return true if type is the E4M3FN variant of an 8-bit float that is supported by the _fp8 instructions on target.
static std::optional< StringRef > mfmaOpToIntrinsic (MFMAOp mfma, const ROCDL::TargetInfo &target)
 Return the rocdl intrinsic corresponding to a MFMA operation mfma if one exists.
static std::optional< ROCDL::MatrixFormat > smallFloatTypeToMatrixFormat (Type mlirElemType)
static std::optional< ScaledMFMAIntrinsic > mfmaOpToScaledIntrinsic (Type aType, Type bType, Type destType, uint32_t m, uint32_t n, uint32_t k, uint32_t b, const ROCDL::TargetInfo &target)
static std::optional< ScaledMFMAIntrinsic > mfmaOpToScaledIntrinsic (MFMAOp mfma, const ROCDL::TargetInfo &target)
static std::optional< ScaledMFMAIntrinsic > mfmaOpToScaledIntrinsic (ScaledMFMAOp smfma, const ROCDL::TargetInfo &target)
static std::optional< StringRef > wmmaOpToIntrinsicRDNA (Type elemSourceType, Type elemBSourceType, Type elemDestType, uint32_t k, bool isRDNA3)
 Returns the rocdl intrinsic corresponding to a WMMA operation wmma for RDNA3/4 architectures.
static std::optional< StringRef > wmmaOpToIntrinsicGfx1250 (Type elemSourceType, Type elemBSourceType, Type elemDestType, uint32_t k)
 Return the rocdl intrinsic corresponding to a WMMA operation wmma for the gfx1250 architecture.
static std::optional< StringRef > smfmacOpToIntrinsic (SparseMFMAOp op, const ROCDL::TargetInfo &target)
 Returns the rocdl intrinsic corresponding to a SparseMFMA (smfmac) operation if one exists.
static std::optional< StringRef > wmmaOpToIntrinsic (WMMAOp wmma, const ROCDL::TargetInfo &target)
 Returns the rocdl intrinsic corresponding to a WMMA operation wmma if one exists.
static std::optional< SparseWMMAOpInfo > sparseWMMAOpToIntrinsic (SparseWMMAOp swmmac, const ROCDL::TargetInfo &target)

Macro Definition Documentation

◆ GEN_PASS_DEF_CONVERTAMDGPUTOROCDLPASS

#define GEN_PASS_DEF_CONVERTAMDGPUTOROCDLPASS

Definition at line 39 of file AMDGPUToROCDL.cpp.

Typedef Documentation

◆ ScaledMFMAIntrinsic

Initial value:
std::tuple<StringRef, ROCDL::MatrixFormat, ROCDL::MatrixFormat>

If there is a scaled MFMA instruction for the input element types aType and bType, output type destType, problem size M, N, K, and B (number of blocks) on the given chipset, return a tuple consisting of the OperationName of the intrinsic and the type codes that need to be passed to that intrinsic.

Note that this is also used to implement some un-scaled MFMAs, since the compiler represents the ordinary instruction as a "scaled" MFMA with a scale of 0.

Definition at line 1049 of file AMDGPUToROCDL.cpp.

Function Documentation

◆ castScaleOperand()

Value castScaleOperand ( ConversionPatternRewriter & rewriter,
Location loc,
Value input )
static

Converts the scaled MFMA/WMMA operands, scalesA and scalesB, from MLIR AMDGPU dialect convention to ROCDL and LLVM AMDGPU intrinsics convention.

Specifically:

  1. If input is a i8 value, zero extend it to i32
  2. If input is a vector of length 4 or 8 and type i8, cast it to i32

Note that the type of input has already been LLVM type converted: therefore 8-bit and smaller floats are represented as their corresponding iN integers.

Definition at line 733 of file AMDGPUToROCDL.cpp.

References mlir::Value::getType().

◆ convertPackedVectorOperand()

Value convertPackedVectorOperand ( ConversionPatternRewriter & rewriter,
Location loc,
Value input,
bool allowBf16 = true )
static

Converts packed vector operands to the expected ROCDL types.

Definition at line 701 of file AMDGPUToROCDL.cpp.

References mlir::Value::getType().

◆ convertUnsignedToI32()

Value convertUnsignedToI32 ( ConversionPatternRewriter & rewriter,
Location loc,
Value val )
static

Convert an unsigned number val to i32.

Definition at line 60 of file AMDGPUToROCDL.cpp.

References convertUnsignedToInt().

Referenced by getLinearIndexI32().

◆ convertUnsignedToI64()

Value convertUnsignedToI64 ( ConversionPatternRewriter & rewriter,
Location loc,
Value val )
static

Convert an unsigned number val to i64.

Definition at line 71 of file AMDGPUToROCDL.cpp.

References convertUnsignedToInt().

Referenced by getNumRecords().

◆ convertUnsignedToInt()

Value convertUnsignedToInt ( ConversionPatternRewriter & rewriter,
Location loc,
Value val,
unsigned width )
static

Zero-extend or truncate the unsigned number val to width bits.

Definition at line 47 of file AMDGPUToROCDL.cpp.

References mlir::Value::getType().

Referenced by convertUnsignedToI32(), convertUnsignedToI64(), and makeBufferRsrc().

◆ createI32Constant()

Value createI32Constant ( ConversionPatternRewriter & rewriter,
Location loc,
int32_t value )
static

◆ createI64Constant()

Value createI64Constant ( ConversionPatternRewriter & rewriter,
Location loc,
int64_t value )
static

Definition at line 76 of file AMDGPUToROCDL.cpp.

Referenced by getNumRecords().

◆ getLinearIndexI32()

Value getLinearIndexI32 ( ConversionPatternRewriter & rewriter,
Location loc,
MemRefDescriptor & memRefDescriptor,
ValueRange indices,
ArrayRef< int64_t > strides )
static

Returns the linear index used to access an element in the memref.

Definition at line 82 of file AMDGPUToROCDL.cpp.

References convertUnsignedToI32(), createI32Constant(), indices, and mlir::MemRefDescriptor::stride().

◆ getNumRecords()

Value getNumRecords ( ConversionPatternRewriter & rewriter,
Location loc,
MemRefType memrefType,
MemRefDescriptor & memrefDescriptor,
ArrayRef< int64_t > strides,
int64_t elementByteWidth,
unsigned numRecordsWidth,
bool boundsCheck )
static

Compute the contents of the num_records field for a given memref descriptor - that is, the number of bytes that's one element past the greatest possible valid index into the memref.

Definition at line 105 of file AMDGPUToROCDL.cpp.

References convertUnsignedToI64(), createI64Constant(), mlir::MemRefDescriptor::size(), and mlir::MemRefDescriptor::stride().

◆ getScaledWmmaIntrinsicName()

std::optional< StringRef > getScaledWmmaIntrinsicName ( int64_t m,
int64_t n,
int64_t k,
bool isScale16 )
static

Determines the ROCDL intrinsic name for scaled WMMA based on dimensions and scale block size (16 or 32).

Definition at line 765 of file AMDGPUToROCDL.cpp.

◆ getWmmaScaleFormat()

std::optional< ROCDL::WMMAMatrixScaleFormat > getWmmaScaleFormat ( Type elemType)
static

Maps f8 scale element types to WMMA scale format codes.

Definition at line 755 of file AMDGPUToROCDL.cpp.

◆ makeBufferRsrc()

Value makeBufferRsrc ( ConversionPatternRewriter & rewriter,
Location loc,
Value basePointer,
Value numRecords,
bool boundsCheck,
const ROCDL::TargetInfo & target,
unsigned numRecordsWidth,
Value cacheSwizzleStride = nullptr,
unsigned addressSpace = 8 )
static

Definition at line 135 of file AMDGPUToROCDL.cpp.

References convertUnsignedToInt(), createI32Constant(), and target.

◆ mfmaOpToIntrinsic()

std::optional< StringRef > mfmaOpToIntrinsic ( MFMAOp mfma,
const ROCDL::TargetInfo & target )
static

Return the rocdl intrinsic corresponding to a MFMA operation mfma if one exists.

This includes checking to ensure the intrinsic is supported on the architecture you are compiling for.

Definition at line 882 of file AMDGPUToROCDL.cpp.

References b, mlir::getElementTypeOrSelf(), mlir::Type::isBF16(), mlir::Type::isF16(), mlir::Type::isF32(), mlir::Type::isF64(), mlir::Type::isInteger(), target, typeIsExpectedBf8ForTarget(), and typeIsExpectedFp8ForTarget().

◆ mfmaOpToScaledIntrinsic() [1/3]

std::optional< ScaledMFMAIntrinsic > mfmaOpToScaledIntrinsic ( MFMAOp mfma,
const ROCDL::TargetInfo & target )
static

Definition at line 1084 of file AMDGPUToROCDL.cpp.

References mfmaOpToScaledIntrinsic(), and target.

◆ mfmaOpToScaledIntrinsic() [2/3]

std::optional< ScaledMFMAIntrinsic > mfmaOpToScaledIntrinsic ( ScaledMFMAOp smfma,
const ROCDL::TargetInfo & target )
static

Definition at line 1092 of file AMDGPUToROCDL.cpp.

References mfmaOpToScaledIntrinsic(), and target.

◆ mfmaOpToScaledIntrinsic() [3/3]

std::optional< ScaledMFMAIntrinsic > mfmaOpToScaledIntrinsic ( Type aType,
Type bType,
Type destType,
uint32_t m,
uint32_t n,
uint32_t k,
uint32_t b,
const ROCDL::TargetInfo & target )
static

◆ packSmallFloatVectorOperand()

Value packSmallFloatVectorOperand ( ConversionPatternRewriter & rewriter,
Location loc,
Value input,
bool allowBf16 = true )
static

Pack small float vector operands (fp4/fp6/fp8/bf16) into the format expected by scaled matrix multiply intrinsics (MFMA/WMMA).

Specifically:

  1. If the element type is bfloat16, bitcast it to i16 unless rocdl intrinsic allows bf16. Newer MFMAs support bf16 types on operand, check IntrinsicsAMDGPU.td file for reference.
  2. If instead we have a more than 64-bit quantity, use a <N / 4 x i32> instead, which is what the f8f6f4 intrinsics use.
  3. If input is a vector of N <= 8 bytes, bitcast it to a (N * 8)-bit integer.

Note that the type of input has already been LLVM type converted: therefore 8-bit and smaller floats are represented as their corresponding iN integers.

Definition at line 674 of file AMDGPUToROCDL.cpp.

References mlir::Value::getType().

◆ smallFloatTypeToMatrixFormat()

std::optional< ROCDL::MatrixFormat > smallFloatTypeToMatrixFormat ( Type mlirElemType)
static

Definition at line 1031 of file AMDGPUToROCDL.cpp.

Referenced by mfmaOpToScaledIntrinsic().

◆ smfmacOpToIntrinsic()

std::optional< StringRef > smfmacOpToIntrinsic ( SparseMFMAOp op,
const ROCDL::TargetInfo & target )
static

Returns the rocdl intrinsic corresponding to a SparseMFMA (smfmac) operation if one exists.

This includes checking to ensure the intrinsic is supported on the architecture you are compiling for.

Definition at line 1252 of file AMDGPUToROCDL.cpp.

References mlir::getElementTypeOrSelf(), mlir::Type::isBF16(), mlir::Type::isF16(), mlir::Type::isF32(), mlir::Type::isInteger(), target, typeIsExpectedBf8ForTarget(), and typeIsExpectedFp8ForTarget().

◆ sparseWMMAOpToIntrinsic()

std::optional< SparseWMMAOpInfo > sparseWMMAOpToIntrinsic ( SparseWMMAOp swmmac,
const ROCDL::TargetInfo & target )
static

◆ typeIsExpectedBf8ForTarget()

bool typeIsExpectedBf8ForTarget ( const ROCDL::TargetInfo & target,
Type type )
static

Return true if type is the E5M2 variant of an 8-bit float that is supported by the _bf8 instructions on target.

Definition at line 864 of file AMDGPUToROCDL.cpp.

References target.

Referenced by mfmaOpToIntrinsic(), and smfmacOpToIntrinsic().

◆ typeIsExpectedFp8ForTarget()

bool typeIsExpectedFp8ForTarget ( const ROCDL::TargetInfo & target,
Type type )
static

Return true if type is the E4M3FN variant of an 8-bit float that is supported by the _fp8 instructions on target.

Definition at line 872 of file AMDGPUToROCDL.cpp.

References target.

Referenced by mfmaOpToIntrinsic(), and smfmacOpToIntrinsic().

◆ wmmaOpToIntrinsic()

std::optional< StringRef > wmmaOpToIntrinsic ( WMMAOp wmma,
const ROCDL::TargetInfo & target )
static

Returns the rocdl intrinsic corresponding to a WMMA operation wmma if one exists.

This includes checking to ensure the intrinsic is supported on the architecture you are compiling for.

Definition at line 1351 of file AMDGPUToROCDL.cpp.

References target, wmmaOpToIntrinsicGfx1250(), and wmmaOpToIntrinsicRDNA().

◆ wmmaOpToIntrinsicGfx1250()

std::optional< StringRef > wmmaOpToIntrinsicGfx1250 ( Type elemSourceType,
Type elemBSourceType,
Type elemDestType,
uint32_t k )
static

Return the rocdl intrinsic corresponding to a WMMA operation wmma for the gfx1250 architecture.

Definition at line 1158 of file AMDGPUToROCDL.cpp.

References mlir::Type::isBF16(), mlir::Type::isF16(), mlir::Type::isF32(), and mlir::Type::isInteger().

Referenced by wmmaOpToIntrinsic().

◆ wmmaOpToIntrinsicRDNA()

std::optional< StringRef > wmmaOpToIntrinsicRDNA ( Type elemSourceType,
Type elemBSourceType,
Type elemDestType,
uint32_t k,
bool isRDNA3 )
static

Returns the rocdl intrinsic corresponding to a WMMA operation wmma for RDNA3/4 architectures.

Definition at line 1102 of file AMDGPUToROCDL.cpp.

References mlir::Type::isBF16(), mlir::Type::isF16(), mlir::Type::isF32(), and mlir::Type::isInteger().

Referenced by wmmaOpToIntrinsic().

◆ wmmaPushInputOperand()

void wmmaPushInputOperand ( ConversionPatternRewriter & rewriter,
Location loc,
const TypeConverter * typeConverter,
bool isUnsigned,
Value llvmInput,
Value mlirInput,
SmallVectorImpl< Value > & operands,
SmallVectorImpl< NamedAttribute > & attrs,
StringRef attrName )
static

Push an input operand.

If it is a float type, nothing to do. If it is an integer type, then we need to also push its signdness (1 for signed, 0 for unsigned) and we need to pack the input 16xi8 vector into a 4xi32 vector (or the 8xi8 vector into a 2xi32 one for gfx12+). We also need to convert bfloat inputs to i16 to account for the bfloat intrinsics having been defined before the AMD backend supported bfloat. We similarly need to pack 8-bit float types into integers as if they were i8 (which they are for the backend's purposes).

Definition at line 786 of file AMDGPUToROCDL.cpp.

References mlir::Type::getIntOrFloatBitWidth(), mlir::Value::getType(), mlir::Type::isSignedInteger(), and mlir::Type::isUnsignedInteger().

◆ wmmaPushOutputOperand()

void wmmaPushOutputOperand ( ConversionPatternRewriter & rewriter,
Location loc,
const TypeConverter * typeConverter,
Value output,
int32_t subwordOffset,
bool clamp,
SmallVectorImpl< Value > & operands,
SmallVectorImpl< NamedAttribute > & attrs )
static

Push the output operand.

For many cases this is only pushing the output in the operand list. But when we have f16 -> f16 or bf16 -> bf16 intrinsics, since the same numbers of VGPRs is used, we need to decide if to store the result in the upper 16 bits of the VGPRs or in the lower part. To store the result in the lower 16 bits, set subwordOffset to 1, otherwise result will be stored it in the upper part. The subwordOffset must not be set for gfx12, as the instructions have been changed to return fewer registers instead.

Definition at line 844 of file AMDGPUToROCDL.cpp.

References clamp(), mlir::Value::getType(), mlir::Type::isBF16(), mlir::Type::isF16(), and mlir::Type::isInteger().