#include "mlir/Dialect/LLVMIR/NVVMDialect.h"
#include "mlir/Conversion/ConvertToLLVM/ToLLVMInterface.h"
#include "mlir/Dialect/GPU/IR/CompilationInterfaces.h"
#include "mlir/Dialect/GPU/IR/GPUDialect.h"
#include "mlir/IR/Builders.h"
#include "mlir/IR/BuiltinAttributes.h"
#include "mlir/IR/BuiltinTypes.h"
#include "mlir/IR/Diagnostics.h"
#include "mlir/IR/DialectImplementation.h"
#include "mlir/IR/MLIRContext.h"
#include "mlir/IR/Operation.h"
#include "mlir/IR/OperationSupport.h"
#include "mlir/IR/Types.h"
#include "llvm/ADT/STLExtras.h"
#include "llvm/ADT/TypeSwitch.h"
#include "llvm/IR/IRBuilder.h"
#include "llvm/Support/Casting.h"
#include "llvm/Support/FormatVariadic.h"
#include "llvm/Support/NVPTXAddrSpace.h"
#include "llvm/Support/raw_ostream.h"
#include <cassert>
#include <optional>
#include <string>
#include "mlir/Dialect/LLVMIR/NVVMOpsDialect.cpp.inc"
#include "mlir/Dialect/LLVMIR/NVVMOpsEnums.cpp.inc"
#include "mlir/Dialect/LLVMIR/NVVMOps.cpp.inc"
#include "mlir/Dialect/LLVMIR/NVVMOpsAttributes.cpp.inc"

Macros
#define	CP_ASYNC_ID_IMPL(mod, size, suffix) llvm::Intrinsic::nvvm_cp_async_##mod##_shared_global_##size##suffix

#define	GET_CP_ASYNC_ID(mod, size, has_cpsize) has_cpsize ? CP_ASYNC_ID_IMPL(mod, size, _s) : CP_ASYNC_ID_IMPL(mod, size, )

#define	CP_ASYNC_BULK_TENSOR_REDUCE_MODE(op, dim, mode) llvm::Intrinsic::nvvm_cp_async_bulk_tensor_##op##_##mode##_##dim##d

#define	CP_ASYNC_BULK_TENSOR_REDUCE(op, dim, is_im2col)

#define	GET_CP_ASYNC_BULK_TENSOR_ID(op, dims, is_im2col)

#define	_none

#define	CVT_F2TF32_ID_IMPL(rnd, relu, sf)

#define	GET_CVT_F2TF32_ID(rnd, relu, sf)

#define	GET_F32x2_TO_F6x2_ID(type, has_relu)

#define	GET_F32x2_TO_F8X2_US_ID(rnd, has_satf)

#define	GET_F32x2_TO_F8X2_S_ID(type, has_relu)

#define	GET_F16x2_TO_F8X2_ID(type, has_relu)

#define	GET_BF16X2_TO_F8X2_ID(rnd, has_satf)

#define	TCGEN05_COMMIT_IMPL(cg, is_shared, mc)

#define	GET_TCGEN05_COMMIT_ID(cta_group, is_shared, has_mc)

#define	TCGEN05_CP_IMPL(shape_mc, src_fmt, cg) llvm::Intrinsic::nvvm_tcgen05_cp##shape_mc##src_fmt##cg

#define	TCGEN05_CP_2CTA(shape_mc, src_fmt, is_2cta)

#define	GET_TCGEN05_CP_ID(shape_mc, src_fmt, is_2cta)

#define	GET_OP_LIST

#define	GET_ATTRDEF_LIST

#define	GET_OP_CLASSES

#define	GET_ATTRDEF_CLASSES

Functions
static LogicalResult	cpAsyncBulkTensorCommonVerifier (size_t tensorDims, bool isIm2Col, size_t numIm2ColOffsets, Location loc)

static LogicalResult	verifyTMALoadParams (size_t tensorDims, size_t numIm2colOff, TMALoadMode mode, Location loc)

static bool	isInt4PtxType (MMATypes type)

static bool	isInt8PtxType (MMATypes type)

static bool	isIntegerPtxType (MMATypes type)

static std::pair< mlir::Type, unsigned >	inferMMATypeFromMNK (NVVM::MMATypes type, NVVM::MMAFrag frag, int m, int n, int k, MLIRContext *context)

FailureOr< int >	getAllowedSizeK (NVVM::WGMMATypes typeA)

LogicalResult	isAllowedWGMMADataType (NVVM::WGMMATypes typeD, NVVM::WGMMATypes typeA, NVVM::WGMMATypes typeB)

LogicalResult	isAllowedSizeN (int sizeN, NVVM::WGMMATypes typeA)

static llvm::Value *	packValInto64Bits (llvm::IRBuilderBase &builder, llvm::Value result, llvm::Value field, unsigned sizeInBits, unsigned start)
	Packs the given `field` into the `result`. More...

static unsigned	isValidVectorLength (NVVM::Tcgen05LdStShape shape, unsigned vecLen)

static void	nvvmInferResultRanges (Operation *op, Value result, ArrayRef<::mlir::ConstantIntRanges > argRanges, SetIntRangeFn setResultRanges)
	Infer the result ranges for the NVVM SpecialRangeableRegisterOp that might have ConstantRangeAttr. More...

static llvm::Value *	getAsPackedI32 (llvm::Value *arg, llvm::IRBuilderBase &builder)

static llvm::Value *	getParamCastedAddr (llvm::Value *addr, llvm::IRBuilderBase &builder)

Macro Definition Documentation

◆ _none

#define _none

Definition at line 1686 of file NVVMDialect.cpp.

◆ CP_ASYNC_BULK_TENSOR_REDUCE

#define CP_ASYNC_BULK_TENSOR_REDUCE	(	op,
		dim,
		is_im2col
	)

Value:

is_im2col ? CP_ASYNC_BULK_TENSOR_REDUCE_MODE(op, dim, im2col) \

: CP_ASYNC_BULK_TENSOR_REDUCE_MODE(op, dim, tile)

CP_ASYNC_BULK_TENSOR_REDUCE_MODE

#define CP_ASYNC_BULK_TENSOR_REDUCE_MODE(op, dim, mode)

Definition: NVVMDialect.cpp:1637

mlir::affine::tile

SmallVector< SmallVector< AffineForOp, 8 >, 8 > tile(ArrayRef< AffineForOp > forOps, ArrayRef< uint64_t > sizes, ArrayRef< AffineForOp > targets)

Performs tiling fo imperfectly nested loops (with interchange) by strip-mining the forOps by sizes an...

Definition: LoopUtils.cpp:1584

Definition at line 1640 of file NVVMDialect.cpp.

◆ CP_ASYNC_BULK_TENSOR_REDUCE_MODE

#define CP_ASYNC_BULK_TENSOR_REDUCE_MODE	(	op,
		dim,
		mode
	)	llvm::Intrinsic::nvvm_cp_async_bulk_tensor_##op##_##mode##_##dim##d

Definition at line 1637 of file NVVMDialect.cpp.

◆ CP_ASYNC_ID_IMPL

#define CP_ASYNC_ID_IMPL	(	mod,
		size,
		suffix
	)	llvm::Intrinsic::nvvm_cp_async_##mod##_shared_global_##size##suffix

Definition at line 1453 of file NVVMDialect.cpp.

◆ CVT_F2TF32_ID_IMPL

#define CVT_F2TF32_ID_IMPL	(	rnd,
		relu,
		sf
	)

Value:

hasRelu ? llvm::Intrinsic::nvvm_f2tf32_##rnd##relu##sf \

: llvm::Intrinsic::nvvm_f2tf32_##rnd##sf

Definition at line 1688 of file NVVMDialect.cpp.

◆ GET_ATTRDEF_CLASSES

#define GET_ATTRDEF_CLASSES

Definition at line 2251 of file NVVMDialect.cpp.

◆ GET_ATTRDEF_LIST

#define GET_ATTRDEF_LIST

◆ GET_BF16X2_TO_F8X2_ID

#define GET_BF16X2_TO_F8X2_ID	(	rnd,
		has_satf
	)

Value:

has_satf ? llvm::Intrinsic::nvvm_bf16x2_to_ue8m0x2_##rnd##_satfinite \

: llvm::Intrinsic::nvvm_bf16x2_to_ue8m0x2_##rnd

Definition at line 1774 of file NVVMDialect.cpp.

◆ GET_CP_ASYNC_BULK_TENSOR_ID

#define GET_CP_ASYNC_BULK_TENSOR_ID	(	op,
		dims,
		is_im2col
	)

Value:

  [&]() -> auto {                                                              \
    switch (dims) {                                                            \
    case 1:                                                                    \
      return CP_ASYNC_BULK_TENSOR_REDUCE_MODE(op, 1, tile);                    \
    case 2:                                                                    \
      return CP_ASYNC_BULK_TENSOR_REDUCE_MODE(op, 2, tile);                    \
    case 3:                                                                    \
      return CP_ASYNC_BULK_TENSOR_REDUCE(op, 3, is_im2col);                    \
    case 4:                                                                    \
      return CP_ASYNC_BULK_TENSOR_REDUCE(op, 4, is_im2col);                    \
    case 5:                                                                    \
      return CP_ASYNC_BULK_TENSOR_REDUCE(op, 5, is_im2col);                    \
    default:                                                                   \
      llvm_unreachable("Invalid TensorDim in CpAsyncBulkTensorReduceOp.");     \
    }                                                                          \
  }()

Definition at line 1644 of file NVVMDialect.cpp.

◆ GET_CP_ASYNC_ID

#define GET_CP_ASYNC_ID	(	mod,
		size,
		has_cpsize
	)	has_cpsize ? CP_ASYNC_ID_IMPL(mod, size, _s) : CP_ASYNC_ID_IMPL(mod, size, )

Definition at line 1456 of file NVVMDialect.cpp.

◆ GET_CVT_F2TF32_ID

#define GET_CVT_F2TF32_ID	(	rnd,
		relu,
		sf
	)

Value:

hasSatFinite ? CVT_F2TF32_ID_IMPL(rnd, relu, sf) \

: CVT_F2TF32_ID_IMPL(rnd, relu, )

CVT_F2TF32_ID_IMPL

#define CVT_F2TF32_ID_IMPL(rnd, relu, sf)

Definition: NVVMDialect.cpp:1688

Definition at line 1692 of file NVVMDialect.cpp.

◆ GET_F16x2_TO_F8X2_ID

#define GET_F16x2_TO_F8X2_ID	(	type,
		has_relu
	)

Value:

has_relu ? llvm::Intrinsic::nvvm_f16x2_to_##type##_rn_relu \

: llvm::Intrinsic::nvvm_f16x2_to_##type##_rn

Definition at line 1758 of file NVVMDialect.cpp.

◆ GET_F32x2_TO_F6x2_ID

#define GET_F32x2_TO_F6x2_ID	(	type,
		has_relu
	)

Value:

has_relu ? llvm::Intrinsic::nvvm_ff_to_##type##_rn_relu_satfinite \

: llvm::Intrinsic::nvvm_ff_to_##type##_rn_satfinite

Definition at line 1713 of file NVVMDialect.cpp.

◆ GET_F32x2_TO_F8X2_S_ID

#define GET_F32x2_TO_F8X2_S_ID	(	type,
		has_relu
	)

Value:

has_relu ? llvm::Intrinsic::nvvm_ff_to_##type##_rn_relu \

: llvm::Intrinsic::nvvm_ff_to_##type##_rn

Definition at line 1732 of file NVVMDialect.cpp.

◆ GET_F32x2_TO_F8X2_US_ID

#define GET_F32x2_TO_F8X2_US_ID	(	rnd,
		has_satf
	)

Value:

has_satf ? llvm::Intrinsic::nvvm_ff_to_ue8m0x2_##rnd##_satfinite \

: llvm::Intrinsic::nvvm_ff_to_ue8m0x2_##rnd

Definition at line 1728 of file NVVMDialect.cpp.

◆ GET_OP_CLASSES

#define GET_OP_CLASSES

Definition at line 2248 of file NVVMDialect.cpp.

◆ GET_OP_LIST

#define GET_OP_LIST

◆ GET_TCGEN05_COMMIT_ID

#define GET_TCGEN05_COMMIT_ID	(	cta_group,
		is_shared,
		has_mc
	)

Value:

has_mc ? TCGEN05_COMMIT_IMPL(cta_group, is_shared, _mc) \

: TCGEN05_COMMIT_IMPL(cta_group, is_shared, )

TCGEN05_COMMIT_IMPL

#define TCGEN05_COMMIT_IMPL(cg, is_shared, mc)

Definition: NVVMDialect.cpp:1833

Definition at line 1837 of file NVVMDialect.cpp.

◆ GET_TCGEN05_CP_ID

#define GET_TCGEN05_CP_ID	(	shape_mc,
		src_fmt,
		is_2cta
	)

Value:

  [&]() -> auto {                                                              \
    if ((src_fmt) == Tcgen05CpSrcFormat::B6x16_P32)                            \
      return TCGEN05_CP_2CTA(shape_mc, _b6x16_p32, is_2cta);                   \
    if ((src_fmt) == Tcgen05CpSrcFormat::B4x16_P64)                            \
      return TCGEN05_CP_2CTA(shape_mc, _b4x16_p64, is_2cta);                   \
    return TCGEN05_CP_2CTA(shape_mc, , is_2cta);                               \
  }()

Definition at line 1871 of file NVVMDialect.cpp.

◆ TCGEN05_COMMIT_IMPL

#define TCGEN05_COMMIT_IMPL	(	cg,
		is_shared,
		mc
	)

Value:

is_shared ? llvm::Intrinsic::nvvm_tcgen05_commit##mc##_shared##_##cg \

: llvm::Intrinsic::nvvm_tcgen05_commit##mc##_##cg

Definition at line 1833 of file NVVMDialect.cpp.

◆ TCGEN05_CP_2CTA

#define TCGEN05_CP_2CTA	(	shape_mc,
		src_fmt,
		is_2cta
	)

Value:

is_2cta ? TCGEN05_CP_IMPL(shape_mc, src_fmt, _cg2) \

: TCGEN05_CP_IMPL(shape_mc, src_fmt, _cg1)

TCGEN05_CP_IMPL

#define TCGEN05_CP_IMPL(shape_mc, src_fmt, cg)

Definition: NVVMDialect.cpp:1864

Definition at line 1867 of file NVVMDialect.cpp.

◆ TCGEN05_CP_IMPL

#define TCGEN05_CP_IMPL	(	shape_mc,
		src_fmt,
		cg
	)	llvm::Intrinsic::nvvm_tcgen05_cp##shape_mc##src_fmt##cg

Definition at line 1864 of file NVVMDialect.cpp.

Function Documentation

◆ cpAsyncBulkTensorCommonVerifier()

static LogicalResult cpAsyncBulkTensorCommonVerifier	(	size_t	tensorDims,
		bool	isIm2Col,
		size_t	numIm2ColOffsets,
		Location	loc
	)

static

Definition at line 55 of file NVVMDialect.cpp.

References mlir::emitError().

◆ getAllowedSizeK()

FailureOr<int> getAllowedSizeK ( NVVM::WGMMATypes typeA )

Definition at line 954 of file NVVMDialect.cpp.

◆ getAsPackedI32()

static llvm::Value* getAsPackedI32	(	llvm::Value *	arg,
		llvm::IRBuilderBase &	builder
	)

static

Definition at line 1959 of file NVVMDialect.cpp.

◆ getParamCastedAddr()

static llvm::Value* getParamCastedAddr	(	llvm::Value *	addr,
		llvm::IRBuilderBase &	builder
	)

static

Definition at line 2008 of file NVVMDialect.cpp.

References mlir::get().

◆ inferMMATypeFromMNK()

static std::pair<mlir::Type, unsigned> inferMMATypeFromMNK	(	NVVM::MMATypes	type,
		NVVM::MMAFrag	frag,
		int	m,
		int	n,
		int	k,
		MLIRContext *	context
	)

static

Definition at line 783 of file NVVMDialect.cpp.

References mlir::NVVM::inferMMAType().

◆ isAllowedSizeN()

LogicalResult isAllowedSizeN	(	int	sizeN,
		NVVM::WGMMATypes	typeA
	)

Definition at line 1010 of file NVVMDialect.cpp.

◆ isAllowedWGMMADataType()

LogicalResult isAllowedWGMMADataType	(	NVVM::WGMMATypes	typeD,
		NVVM::WGMMATypes	typeA,
		NVVM::WGMMATypes	typeB
	)

Definition at line 968 of file NVVMDialect.cpp.

◆ isInt4PtxType()

static bool isInt4PtxType ( MMATypes type )

static

Definition at line 303 of file NVVMDialect.cpp.

Referenced by isIntegerPtxType().

◆ isInt8PtxType()

static bool isInt8PtxType ( MMATypes type )

static

Definition at line 307 of file NVVMDialect.cpp.

Referenced by isIntegerPtxType().

◆ isIntegerPtxType()

static bool isIntegerPtxType ( MMATypes type )

static

Definition at line 311 of file NVVMDialect.cpp.

References isInt4PtxType(), and isInt8PtxType().

◆ isValidVectorLength()

static unsigned isValidVectorLength	(	NVVM::Tcgen05LdStShape	shape,
		unsigned	vecLen
	)

static

Definition at line 1905 of file NVVMDialect.cpp.

◆ nvvmInferResultRanges()

static void nvvmInferResultRanges	(	Operation *	op,
		Value	result,
		ArrayRef<::mlir::ConstantIntRanges >	argRanges,
		SetIntRangeFn	setResultRanges
	)

static

Infer the result ranges for the NVVM SpecialRangeableRegisterOp that might have ConstantRangeAttr.

Definition at line 1950 of file NVVMDialect.cpp.

References mlir::Operation::getAttrOfType().

◆ packValInto64Bits()

static llvm::Value* packValInto64Bits	(	llvm::IRBuilderBase &	builder,
		llvm::Value *	result,
		llvm::Value *	field,
		unsigned	sizeInBits,
		unsigned	start
	)

static

Packs the given field into the result.

The result is 64-bits and each field can be 32-bits or narrower.

Definition at line 1408 of file NVVMDialect.cpp.

◆ verifyTMALoadParams()

static LogicalResult verifyTMALoadParams	(	size_t	tensorDims,
		size_t	numIm2colOff,
		TMALoadMode	mode,
		Location	loc
	)

static

Definition at line 121 of file NVVMDialect.cpp.

References mlir::emitError().

Macros

Functions

Macro Definition Documentation

◆ _none

◆ CP_ASYNC_BULK_TENSOR_REDUCE

◆ CP_ASYNC_BULK_TENSOR_REDUCE_MODE

◆ CP_ASYNC_ID_IMPL

◆ CVT_F2TF32_ID_IMPL

◆ GET_ATTRDEF_CLASSES

◆ GET_ATTRDEF_LIST

◆ GET_BF16X2_TO_F8X2_ID

◆ GET_CP_ASYNC_BULK_TENSOR_ID

◆ GET_CP_ASYNC_ID

◆ GET_CVT_F2TF32_ID

◆ GET_F16x2_TO_F8X2_ID

◆ GET_F32x2_TO_F6x2_ID

◆ GET_F32x2_TO_F8X2_S_ID

◆ GET_F32x2_TO_F8X2_US_ID

◆ GET_OP_CLASSES

◆ GET_OP_LIST

◆ GET_TCGEN05_COMMIT_ID

◆ GET_TCGEN05_CP_ID

◆ TCGEN05_COMMIT_IMPL

◆ TCGEN05_CP_2CTA

◆ TCGEN05_CP_IMPL

Function Documentation

◆ cpAsyncBulkTensorCommonVerifier()

◆ getAllowedSizeK()

◆ getAsPackedI32()

◆ getParamCastedAddr()

◆ inferMMATypeFromMNK()

◆ isAllowedSizeN()

◆ isAllowedWGMMADataType()

◆ isInt4PtxType()

◆ isInt8PtxType()

◆ isIntegerPtxType()

◆ isValidVectorLength()

◆ nvvmInferResultRanges()

◆ packValInto64Bits()

◆ verifyTMALoadParams()