#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	_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)

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

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

static 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)

Variables
static constexpr unsigned	notIntrinsic = llvm::Intrinsic::not_intrinsic

Macro Definition Documentation

◆ _none

#define _none

Definition at line 1949 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 1497 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 1951 of file NVVMDialect.cpp.

◆ GET_ATTRDEF_CLASSES

#define GET_ATTRDEF_CLASSES

Definition at line 2619 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 2037 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 1500 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:1951

Definition at line 1955 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 2021 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 1976 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 1995 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 1991 of file NVVMDialect.cpp.

◆ GET_OP_CLASSES

#define GET_OP_CLASSES

Definition at line 2616 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:2096

Definition at line 2100 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 2134 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 2096 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:2127

Definition at line 2130 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 2127 of file NVVMDialect.cpp.

Function Documentation

◆ cpAsyncBulkTensorCommonVerifier()

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

static

Definition at line 57 of file NVVMDialect.cpp.

References mlir::emitError().

◆ getAllowedSizeK()

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

static

Definition at line 981 of file NVVMDialect.cpp.

◆ getAsPackedI32()

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

static

Definition at line 2222 of file NVVMDialect.cpp.

◆ getParamCastedAddr()

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

static

Definition at line 2271 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 810 of file NVVMDialect.cpp.

References mlir::NVVM::inferMMAType().

◆ isAllowedSizeN()

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

static

Definition at line 1037 of file NVVMDialect.cpp.

◆ isAllowedWGMMADataType()

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

static

Definition at line 995 of file NVVMDialect.cpp.

◆ isInt4PtxType()

static bool isInt4PtxType ( MMATypes type )

static

Definition at line 330 of file NVVMDialect.cpp.

Referenced by isIntegerPtxType().

◆ isInt8PtxType()

static bool isInt8PtxType ( MMATypes type )

static

Definition at line 334 of file NVVMDialect.cpp.

Referenced by isIntegerPtxType().

◆ isIntegerPtxType()

static bool isIntegerPtxType ( MMATypes type )

static

Definition at line 338 of file NVVMDialect.cpp.

References isInt4PtxType(), and isInt8PtxType().

◆ isValidVectorLength()

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

static

Definition at line 2168 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 2213 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 1452 of file NVVMDialect.cpp.

◆ verifyTMALoadParams()

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

static

Definition at line 116 of file NVVMDialect.cpp.

References mlir::emitError().

Variable Documentation

◆ notIntrinsic

constexpr unsigned notIntrinsic = llvm::Intrinsic::not_intrinsic

staticconstexpr

Definition at line 48 of file NVVMDialect.cpp.

Macros

Functions

Variables

Macro Definition Documentation

◆ _none

◆ CP_ASYNC_ID_IMPL

◆ CVT_F2TF32_ID_IMPL

◆ GET_ATTRDEF_CLASSES

◆ GET_ATTRDEF_LIST

◆ GET_BF16X2_TO_F8X2_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()

Variable Documentation

◆ notIntrinsic