|
MLIR 24.0.0git
|
#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 "mlir/Transforms/InliningUtils.h"#include "llvm/ADT/STLExtras.h"#include "llvm/ADT/TypeSwitch.h"#include "llvm/IR/IRBuilder.h"#include "llvm/IR/NVVMIntrinsicUtils.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"Go to the source code of this file.
Classes | |
| struct | ConvertFsubToFnegFadd |
Macros | |
| #define | CP_ASYNC_ID_IMPL(mod, size, suffix) |
| #define | GET_CP_ASYNC_ID(mod, size, has_cpsize) |
| #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 | 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) |
| #define | TCGEN05_CP_2CTA(shape_mc, src_fmt, is_2cta) |
| #define | GET_TCGEN05_CP_ID(shape_mc, src_fmt, is_2cta) |
| #define | TCGEN05LDRED(SHAPE, NUM, TYPE) |
| #define | GET_OP_LIST |
| #define | GET_ATTRDEF_LIST |
| #define | GET_OP_CLASSES |
| #define | GET_ATTRDEF_CLASSES |
Functions | |
| static bool | isPtrInAddrSpace (mlir::Value ptr, NVVMMemorySpace targetAS) |
| static bool | isPtrInGenericSpace (mlir::Value ptr) |
| static bool | isPtrInSharedCTASpace (mlir::Value ptr) |
| static bool | isPtrInSharedClusterSpace (mlir::Value ptr) |
| static llvm::Value * | castPtrToAddrSpace (llvm::IRBuilderBase &builder, llvm::Value *ptr, NVVMMemorySpace targetAS) |
| static llvm::nvvm::CTAGroupKind | getNVVMCtaGroupKind (NVVM::CTAGroupKind ctaGroup) |
| 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 LogicalResult | verifyMBarrierArriveLikeOp (Operation *op, Value addr, NVVM::MemScopeKind scope, Value retVal=nullptr) |
| static LogicalResult | inferMBarrierArriveResultTypes (MLIRContext *context, Value addr, SmallVectorImpl< Type > &inferredReturnTypes) |
| Only shared_cluster (ptr<7>) produces zero results; all other address spaces (including generic) return i64. | |
| static bool | isCompatibleReturnTypesOptionalResult (TypeRange inferred, TypeRange actual) |
| For ops with optional results, allow the user to omit the result even when inference would produce one. | |
| static LogicalResult | verifyConvertF32x2ToFP16x2Op (Twine dstType, FPRoundingMode rnd, bool hasRandomBits, Operation *op) |
| static bool | isInt4PtxType (MMATypes type) |
| static bool | isInt8PtxType (MMATypes type) |
| static bool | isIntegerPtxType (MMATypes type) |
| static void | printOperandList (OpAsmPrinter &p, StringRef name, ArrayRef< Value > operands) |
| static LogicalResult | parseMmaOperand (OpAsmParser &parser, StringRef operandName, SmallVectorImpl< OpAsmParser::UnresolvedOperand > ®s) |
| template<typename Op> | |
| static void | processOperandFragments (Op &op, std::array< MMAOperandFragment, 3 > &frags, SmallVectorImpl< Type > ®Types, SmallVectorImpl< StringRef > &ignoreAttrNames) |
| static LogicalResult | parseMmaTypeSignature (OpAsmParser &parser, SmallVectorImpl< Type > &operandTypes) |
| static void | inferAndSetMultiplicandTypes (MLIRContext *ctx, NamedAttrList &attrs, const SmallVectorImpl< Type > &operandTypes) |
| template<typename OpType> | |
| static void | addBlockScaleProperties (OpBuilder &builder, OperationState &result, ArrayRef< int64_t > shape, ScaleVecSize scaleVecSize, BlockScaleFormat blockScaleFormat, MMABlockScaleKind kind) |
| static void | addInferredMultiplicandTypes (MLIRContext *ctx, OperationState &result, ValueRange operandA, ValueRange operandB, std::optional< std::array< MMATypes, 2 > > multiplicandPtxTypes) |
| template<typename OpTy> | |
| static MMATypes | inferPtxTypeFromResult (OpTy op) |
| 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) |
| template<typename OpType> | |
| static LogicalResult | verifyAddSubFOp (OpType op) |
| static llvm::Value * | packValInto64Bits (llvm::IRBuilderBase &builder, llvm::Value *result, llvm::Value *field, unsigned sizeInBits, unsigned start) |
| Packs the given field into the result. | |
| static llvm::Intrinsic::ID | getBarrierSyncIntrinsic (bool aligned, bool hasCount) |
| Maps the (aligned, hasCount) pair to the @llvm.nvvm.barrier.cta.sync.
| |
| static llvm::Intrinsic::ID | getBarrierReductionIntrinsic (bool aligned, NVVM::BarrierReduction kind) |
| Maps the (aligned, kind) pair to the @llvm.nvvm.barrier.cta.red.
| |
| 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. | |
| static LogicalResult | verifyConstantRangeAttr (Operation *op, std::optional< LLVM::ConstantRangeAttr > rangeAttr) |
| Verify the range attribute satisfies LLVM ConstantRange constructor requirements for NVVM SpecialRangeableRegisterOp. | |
| static llvm::Value * | getAsPackedI32 (llvm::Value *arg, llvm::IRBuilderBase &builder) |
| static llvm::Value * | getParamCastedAddr (llvm::Value *addr, llvm::IRBuilderBase &builder) |
| static LogicalResult | verifyTcgen05MMAOp (bool isATensor, mlir::Value disableOutputLane, NVVM::CTAGroupKind ctaGroup, bool hasAShift, NVVM::Tcgen05MMACollectorOp collectorOp, Location loc) |
| static LogicalResult | verifyTcgen05MMABlockScaleOp (NVVM::Tcgen05MMACollectorOp collectorOp, NVVM::Tcgen05MMAKind kind, NVVM::Tcgen05MMABlockScale blockScale, Location loc) |
Variables | |
| static constexpr unsigned | notIntrinsic = llvm::Intrinsic::not_intrinsic |
| #define _none |
Definition at line 4555 of file NVVMDialect.cpp.
| #define CP_ASYNC_ID_IMPL | ( | mod, | |
| size, | |||
| suffix ) |
Definition at line 4063 of file NVVMDialect.cpp.
| #define CVT_F2TF32_ID_IMPL | ( | rnd, | |
| relu, | |||
| sf ) |
Definition at line 4557 of file NVVMDialect.cpp.
| #define GET_ATTRDEF_CLASSES |
Definition at line 6573 of file NVVMDialect.cpp.
| #define GET_ATTRDEF_LIST |
| #define GET_CP_ASYNC_ID | ( | mod, | |
| size, | |||
| has_cpsize ) |
Definition at line 4066 of file NVVMDialect.cpp.
| #define GET_CVT_F2TF32_ID | ( | rnd, | |
| relu, | |||
| sf ) |
Definition at line 4561 of file NVVMDialect.cpp.
| #define GET_F16x2_TO_F8X2_ID | ( | type, | |
| has_relu ) |
Definition at line 4728 of file NVVMDialect.cpp.
| #define GET_F32x2_TO_F6x2_ID | ( | type, | |
| has_relu ) |
Definition at line 4599 of file NVVMDialect.cpp.
| #define GET_F32x2_TO_F8X2_S_ID | ( | type, | |
| has_relu ) |
Definition at line 4696 of file NVVMDialect.cpp.
| #define GET_F32x2_TO_F8X2_US_ID | ( | rnd, | |
| has_satf ) |
Definition at line 4692 of file NVVMDialect.cpp.
| #define GET_OP_CLASSES |
Definition at line 6570 of file NVVMDialect.cpp.
| #define GET_OP_LIST |
| #define GET_TCGEN05_COMMIT_ID | ( | cta_group, | |
| is_shared, | |||
| has_mc ) |
Definition at line 5110 of file NVVMDialect.cpp.
| #define GET_TCGEN05_CP_ID | ( | shape_mc, | |
| src_fmt, | |||
| is_2cta ) |
Definition at line 5144 of file NVVMDialect.cpp.
| #define TCGEN05_COMMIT_IMPL | ( | cg, | |
| is_shared, | |||
| mc ) |
Definition at line 5106 of file NVVMDialect.cpp.
| #define TCGEN05_CP_2CTA | ( | shape_mc, | |
| src_fmt, | |||
| is_2cta ) |
Definition at line 5140 of file NVVMDialect.cpp.
| #define TCGEN05_CP_IMPL | ( | shape_mc, | |
| src_fmt, | |||
| cg ) |
Definition at line 5137 of file NVVMDialect.cpp.
| #define TCGEN05LDRED | ( | SHAPE, | |
| NUM, | |||
| TYPE ) |
Definition at line 6256 of file NVVMDialect.cpp.
|
static |
Definition at line 1757 of file NVVMDialect.cpp.
References mlir::Builder::getAttr(), mlir::Builder::getContext(), and result.
|
static |
Definition at line 1773 of file NVVMDialect.cpp.
References mlir::getType(), and result.
|
static |
Definition at line 73 of file NVVMDialect.cpp.
|
static |
Definition at line 100 of file NVVMDialect.cpp.
References mlir::emitError(), and success().
|
static |
Definition at line 2603 of file NVVMDialect.cpp.
|
static |
Definition at line 5412 of file NVVMDialect.cpp.
|
static |
Maps the (aligned, kind) pair to the @llvm.nvvm.barrier.cta.red.
* intrinsic ID.
Definition at line 3451 of file NVVMDialect.cpp.
Maps the (aligned, hasCount) pair to the @llvm.nvvm.barrier.cta.sync.
* intrinsic ID.
Definition at line 3438 of file NVVMDialect.cpp.
|
static |
Definition at line 83 of file NVVMDialect.cpp.
|
static |
Definition at line 5461 of file NVVMDialect.cpp.
|
static |
Definition at line 1739 of file NVVMDialect.cpp.
References mlir::NamedAttrList::get(), and mlir::NamedAttrList::set().
|
static |
Only shared_cluster (ptr<7>) produces zero results; all other address spaces (including generic) return i64.
Definition at line 313 of file NVVMDialect.cpp.
References isPtrInSharedClusterSpace(), and success().
|
static |
Definition at line 2396 of file NVVMDialect.cpp.
References mlir::NVVM::inferMMAType().
|
static |
Definition at line 1792 of file NVVMDialect.cpp.
|
static |
Definition at line 2659 of file NVVMDialect.cpp.
References success().
|
static |
Definition at line 2617 of file NVVMDialect.cpp.
References success().
For ops with optional results, allow the user to omit the result even when inference would produce one.
This preserves backward compatibility: the result can be silently discarded (e.g., for fire-and-forget arrive ops).
Definition at line 360 of file NVVMDialect.cpp.
|
static |
Definition at line 721 of file NVVMDialect.cpp.
Referenced by isIntegerPtxType().
|
static |
Definition at line 725 of file NVVMDialect.cpp.
Referenced by isIntegerPtxType().
|
static |
Definition at line 729 of file NVVMDialect.cpp.
References isInt4PtxType(), and isInt8PtxType().
|
static |
Definition at line 56 of file NVVMDialect.cpp.
Referenced by isPtrInGenericSpace(), isPtrInSharedClusterSpace(), and isPtrInSharedCTASpace().
|
static |
Definition at line 61 of file NVVMDialect.cpp.
References isPtrInAddrSpace().
|
static |
Definition at line 69 of file NVVMDialect.cpp.
References isPtrInAddrSpace().
Referenced by inferMBarrierArriveResultTypes(), and verifyMBarrierArriveLikeOp().
|
static |
Definition at line 65 of file NVVMDialect.cpp.
References isPtrInAddrSpace().
Definition at line 5327 of file NVVMDialect.cpp.
|
static |
Infer the result ranges for the NVVM SpecialRangeableRegisterOp that might have ConstantRangeAttr.
Definition at line 5375 of file NVVMDialect.cpp.
References mlir::Operation::getAttrOfType(), mlir::IntegerValueRange::getMaxRange(), and result.
|
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 3335 of file NVVMDialect.cpp.
References result.
|
static |
Definition at line 1674 of file NVVMDialect.cpp.
References mlir::AsmParser::OptionalSquare, mlir::AsmParser::parseKeyword(), mlir::OpAsmParser::parseOperandList(), and success().
|
static |
Definition at line 1715 of file NVVMDialect.cpp.
References mlir::AsmParser::emitError(), mlir::AsmParser::getCurrentLocation(), mlir::AsmParser::parseColon(), mlir::AsmParser::parseCommaSeparatedList(), mlir::AsmParser::parseLParen(), mlir::AsmParser::parseRParen(), mlir::AsmParser::parseType(), and success().
|
static |
Definition at line 1665 of file NVVMDialect.cpp.
References mlir::OpAsmPrinter::printOperands().
|
static |
Definition at line 1688 of file NVVMDialect.cpp.
|
static |
Definition at line 3203 of file NVVMDialect.cpp.
References mlir::Type::isBF16(), mlir::Type::isF16(), mlir::Type::isF64(), and success().
|
static |
Verify the range attribute satisfies LLVM ConstantRange constructor requirements for NVVM SpecialRangeableRegisterOp.
Definition at line 5389 of file NVVMDialect.cpp.
References mlir::Operation::emitOpError(), and success().
|
static |
Definition at line 592 of file NVVMDialect.cpp.
References mlir::Operation::emitOpError(), and success().
|
static |
Definition at line 253 of file NVVMDialect.cpp.
References mlir::Operation::emitError(), isPtrInSharedClusterSpace(), and success().
|
static |
Definition at line 6068 of file NVVMDialect.cpp.
References mlir::emitError(), and success().
|
static |
Definition at line 5803 of file NVVMDialect.cpp.
References mlir::emitError(), mlir::Value::getType(), and success().
|
static |
Definition at line 159 of file NVVMDialect.cpp.
References mlir::emitError(), and success().
|
staticconstexpr |
Definition at line 50 of file NVVMDialect.cpp.