NVVMDialect.cpp File Reference

#include "mlir/Dialect/LLVMIR/NVVMDialect.h"
#include "mlir/Conversion/ConvertToLLVM/ToLLVMInterface.h"
#include "mlir/Dialect/GPU/IR/CompilationInterfaces.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/Type.h"
#include "llvm/Support/Casting.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"

Go to the source code of this file.

Macros
#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	GET_OP_LIST
#define	GET_ATTRDEF_LIST
#define	GET_OP_CLASSES
#define	GET_ATTRDEF_CLASSES

Functions
static void	printNVVMIntrinsicOp (OpAsmPrinter &p, Operation *op)
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)

Macro Definition Documentation

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)

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

GET_ATTRDEF_CLASSES

#define GET_ATTRDEF_CLASSES

Definition at line 1278 of file NVVMDialect.cpp.

GET_ATTRDEF_LIST

#define GET_ATTRDEF_LIST

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

GET_OP_CLASSES

#define GET_OP_CLASSES

Definition at line 1275 of file NVVMDialect.cpp.

GET_OP_LIST

#define GET_OP_LIST

Function Documentation

CpAsyncBulkTensorCommonVerifier()

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

static

Definition at line 82 of file NVVMDialect.cpp.

References mlir::emitError().

getAllowedSizeK()

FailureOr<int> getAllowedSizeK

(

NVVM::WGMMATypes

typeA

)

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

isAllowedSizeN()

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

Definition at line 828 of file NVVMDialect.cpp.

isAllowedWGMMADataType()

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

Definition at line 786 of file NVVMDialect.cpp.

isInt4PtxType()

static bool isInt4PtxType ( MMATypes  type )

static

Definition at line 171 of file NVVMDialect.cpp.

Referenced by isIntegerPtxType().

isInt8PtxType()

static bool isInt8PtxType ( MMATypes  type )

static

Definition at line 175 of file NVVMDialect.cpp.

Referenced by isIntegerPtxType().

isIntegerPtxType()

static bool isIntegerPtxType ( MMATypes  type )

static

Definition at line 179 of file NVVMDialect.cpp.

References isInt4PtxType(), and isInt8PtxType().

printNVVMIntrinsicOp()

static void printNVVMIntrinsicOp ( OpAsmPrinter &  p, Operation *  op  )

static

Definition at line 54 of file NVVMDialect.cpp.

References mlir::Operation::getNumResults(), mlir::Operation::getOperands(), and mlir::Operation::getResultTypes().