#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/Dialect/Utils/StaticValueUtils.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/AsmParser/Parser.h"
#include "llvm/IR/Attributes.h"
#include "llvm/IR/Function.h"
#include "llvm/IR/IRBuilder.h"
#include "llvm/IR/IntrinsicsNVPTX.h"
#include "llvm/IR/Type.h"
#include "llvm/Support/Casting.h"
#include "llvm/Support/FormatVariadic.h"
#include "llvm/Support/SourceMgr.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 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)

Macro Definition Documentation

◆ _none

#define _none

Definition at line 1451 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:1402

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:1590

Definition at line 1405 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 1402 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 1292 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 1453 of file NVVMDialect.cpp.

◆ GET_ATTRDEF_CLASSES

#define GET_ATTRDEF_CLASSES

Definition at line 1963 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 1539 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 1409 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 1295 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:1453

Definition at line 1457 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 1523 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 1478 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 1497 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 1493 of file NVVMDialect.cpp.

◆ GET_OP_CLASSES

#define GET_OP_CLASSES

Definition at line 1960 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:1598

Definition at line 1602 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 1636 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 1598 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:1629

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

Function Documentation

◆ cpAsyncBulkTensorCommonVerifier()

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

static

Definition at line 62 of file NVVMDialect.cpp.

References mlir::emitError().

◆ getAllowedSizeK()

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

Definition at line 835 of file NVVMDialect.cpp.

◆ getAsPackedI32()

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

static

Definition at line 1724 of file NVVMDialect.cpp.

◆ 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 709 of file NVVMDialect.cpp.

References mlir::NVVM::inferMMAType().

◆ isAllowedSizeN()

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

Definition at line 891 of file NVVMDialect.cpp.

◆ isAllowedWGMMADataType()

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

Definition at line 849 of file NVVMDialect.cpp.

◆ isInt4PtxType()

static bool isInt4PtxType ( MMATypes type )

static

Definition at line 229 of file NVVMDialect.cpp.

Referenced by isIntegerPtxType().

◆ isInt8PtxType()

static bool isInt8PtxType ( MMATypes type )

static

Definition at line 233 of file NVVMDialect.cpp.

Referenced by isIntegerPtxType().

◆ isIntegerPtxType()

static bool isIntegerPtxType ( MMATypes type )

static

Definition at line 237 of file NVVMDialect.cpp.

References isInt4PtxType(), and isInt8PtxType().

◆ isValidVectorLength()

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

static

Definition at line 1670 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 1715 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 1247 of file NVVMDialect.cpp.

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

◆ inferMMATypeFromMNK()

◆ isAllowedSizeN()

◆ isAllowedWGMMADataType()

◆ isInt4PtxType()

◆ isInt8PtxType()

◆ isIntegerPtxType()

◆ isValidVectorLength()

◆ nvvmInferResultRanges()

◆ packValInto64Bits()