e13d23bc6c
I am not sure about the meaning of Type in the name (was it meant be interpreted as Kind?), and given the importance and meaning of Type in the context of MLIR, its probably better to rename it. Given the comment in the source code, the suggestion in the GitHub issue and the final discussions in the review, this patch renames the OperandType to UnresolvedOperand. Fixes https://github.com/llvm/llvm-project/issues/54446 Differential Revision: https://reviews.llvm.org/D122142
281 lines
11 KiB
C++
281 lines
11 KiB
C++
//===- NVVMDialect.cpp - NVVM IR Ops and Dialect registration -------------===//
|
|
//
|
|
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
|
|
// See https://llvm.org/LICENSE.txt for license information.
|
|
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
|
|
//
|
|
//===----------------------------------------------------------------------===//
|
|
//
|
|
// This file defines the types and operation details for the NVVM IR dialect in
|
|
// MLIR, and the LLVM IR dialect. It also registers the dialect.
|
|
//
|
|
// The NVVM dialect only contains GPU specific additions on top of the general
|
|
// LLVM dialect.
|
|
//
|
|
//===----------------------------------------------------------------------===//
|
|
|
|
#include "mlir/Dialect/LLVMIR/NVVMDialect.h"
|
|
|
|
#include "mlir/IR/Builders.h"
|
|
#include "mlir/IR/BuiltinTypes.h"
|
|
#include "mlir/IR/DialectImplementation.h"
|
|
#include "mlir/IR/MLIRContext.h"
|
|
#include "mlir/IR/Operation.h"
|
|
#include "mlir/IR/OperationSupport.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/SourceMgr.h"
|
|
|
|
using namespace mlir;
|
|
using namespace NVVM;
|
|
|
|
#include "mlir/Dialect/LLVMIR/NVVMOpsDialect.cpp.inc"
|
|
#include "mlir/Dialect/LLVMIR/NVVMOpsEnums.cpp.inc"
|
|
|
|
//===----------------------------------------------------------------------===//
|
|
// Printing/parsing for NVVM ops
|
|
//===----------------------------------------------------------------------===//
|
|
|
|
static void printNVVMIntrinsicOp(OpAsmPrinter &p, Operation *op) {
|
|
p << " " << op->getOperands();
|
|
if (op->getNumResults() > 0)
|
|
p << " : " << op->getResultTypes();
|
|
}
|
|
|
|
// <operation> ::= `llvm.nvvm.vote.ballot.sync %mask, %pred` : result_type
|
|
ParseResult VoteBallotOp::parse(OpAsmParser &parser, OperationState &result) {
|
|
MLIRContext *context = parser.getContext();
|
|
auto int32Ty = IntegerType::get(context, 32);
|
|
auto int1Ty = IntegerType::get(context, 1);
|
|
|
|
SmallVector<OpAsmParser::UnresolvedOperand, 8> ops;
|
|
Type type;
|
|
return failure(parser.parseOperandList(ops) ||
|
|
parser.parseOptionalAttrDict(result.attributes) ||
|
|
parser.parseColonType(type) ||
|
|
parser.addTypeToList(type, result.types) ||
|
|
parser.resolveOperands(ops, {int32Ty, int1Ty},
|
|
parser.getNameLoc(), result.operands));
|
|
}
|
|
|
|
void VoteBallotOp::print(OpAsmPrinter &p) { printNVVMIntrinsicOp(p, *this); }
|
|
|
|
LogicalResult CpAsyncOp::verify() {
|
|
if (size() != 4 && size() != 8 && size() != 16)
|
|
return emitError("expected byte size to be either 4, 8 or 16.");
|
|
return success();
|
|
}
|
|
|
|
LogicalResult MmaOp::verify() {
|
|
MLIRContext *context = getContext();
|
|
auto f16Ty = Float16Type::get(context);
|
|
auto f16x2Ty = LLVM::getFixedVectorType(f16Ty, 2);
|
|
auto f32Ty = Float32Type::get(context);
|
|
auto f16x2x4StructTy = LLVM::LLVMStructType::getLiteral(
|
|
context, {f16x2Ty, f16x2Ty, f16x2Ty, f16x2Ty});
|
|
auto f32x8StructTy = LLVM::LLVMStructType::getLiteral(
|
|
context, {f32Ty, f32Ty, f32Ty, f32Ty, f32Ty, f32Ty, f32Ty, f32Ty});
|
|
|
|
auto operandTypes = getOperandTypes();
|
|
if (operandTypes != SmallVector<Type, 8>(8, f16x2Ty) &&
|
|
operandTypes != ArrayRef<Type>{f16x2Ty, f16x2Ty, f16x2Ty, f16x2Ty, f32Ty,
|
|
f32Ty, f32Ty, f32Ty, f32Ty, f32Ty, f32Ty,
|
|
f32Ty}) {
|
|
return emitOpError("expected operands to be 4 <halfx2>s followed by either "
|
|
"4 <halfx2>s or 8 floats");
|
|
}
|
|
if (getType() != f32x8StructTy && getType() != f16x2x4StructTy) {
|
|
return emitOpError("expected result type to be a struct of either 4 "
|
|
"<halfx2>s or 8 floats");
|
|
}
|
|
|
|
auto alayout = (*this)->getAttrOfType<StringAttr>("alayout");
|
|
auto blayout = (*this)->getAttrOfType<StringAttr>("blayout");
|
|
|
|
if (!(alayout && blayout) ||
|
|
!(alayout.getValue() == "row" || alayout.getValue() == "col") ||
|
|
!(blayout.getValue() == "row" || blayout.getValue() == "col")) {
|
|
return emitOpError("alayout and blayout attributes must be set to either "
|
|
"\"row\" or \"col\"");
|
|
}
|
|
|
|
if (operandTypes == ArrayRef<Type>{f16x2Ty, f16x2Ty, f16x2Ty, f16x2Ty, f32Ty,
|
|
f32Ty, f32Ty, f32Ty, f32Ty, f32Ty, f32Ty,
|
|
f32Ty} &&
|
|
getType() == f32x8StructTy && alayout.getValue() == "row" &&
|
|
blayout.getValue() == "col") {
|
|
return success();
|
|
}
|
|
return emitOpError("unimplemented mma.sync variant");
|
|
}
|
|
|
|
LogicalResult ShflOp::verify() {
|
|
if (!(*this)->getAttrOfType<UnitAttr>("return_value_and_is_valid"))
|
|
return success();
|
|
auto type = getType().dyn_cast<LLVM::LLVMStructType>();
|
|
auto elementType = (type && type.getBody().size() == 2)
|
|
? type.getBody()[1].dyn_cast<IntegerType>()
|
|
: nullptr;
|
|
if (!elementType || elementType.getWidth() != 1)
|
|
return emitError("expected return type to be a two-element struct with "
|
|
"i1 as the second element");
|
|
return success();
|
|
}
|
|
|
|
std::pair<mlir::Type, unsigned> NVVM::inferMMAType(NVVM::MMATypes type,
|
|
NVVM::MMAFrag frag,
|
|
MLIRContext *context) {
|
|
unsigned numberElements = 0;
|
|
Type elementType;
|
|
OpBuilder builder(context);
|
|
Type f16x2 = VectorType::get(2, builder.getF16Type());
|
|
if (type == NVVM::MMATypes::f16) {
|
|
elementType = f16x2;
|
|
if (frag == NVVM::MMAFrag::a || frag == NVVM::MMAFrag::b)
|
|
numberElements = 8;
|
|
else
|
|
numberElements = 4;
|
|
} else if (type == NVVM::MMATypes::f32) {
|
|
elementType = builder.getF32Type();
|
|
numberElements = 8;
|
|
} else if (type == NVVM::MMATypes::tf32) {
|
|
elementType = builder.getI32Type();
|
|
numberElements = 4;
|
|
}
|
|
assert(numberElements != 0 && elementType != nullptr);
|
|
return std::make_pair(elementType, numberElements);
|
|
}
|
|
|
|
LogicalResult NVVM::WMMALoadOp::verify() {
|
|
unsigned addressSpace =
|
|
ptr().getType().cast<LLVM::LLVMPointerType>().getAddressSpace();
|
|
if (addressSpace != 0 && addressSpace != 1 && addressSpace != 3)
|
|
return emitOpError("expected source pointer in memory "
|
|
"space 0, 1, 3");
|
|
|
|
if (NVVM::WMMALoadOp::getIntrinsicID(m(), n(), k(), layout(), eltype(),
|
|
frag()) == 0)
|
|
return emitOpError() << "invalid attribute combination";
|
|
std::pair<Type, unsigned> typeInfo =
|
|
inferMMAType(eltype(), frag(), getContext());
|
|
Type dstType = LLVM::LLVMStructType::getLiteral(
|
|
getContext(), SmallVector<Type, 8>(typeInfo.second, typeInfo.first));
|
|
if (getType() != dstType)
|
|
return emitOpError("expected destination type is a structure of ")
|
|
<< typeInfo.second << " elements of type " << typeInfo.first;
|
|
return success();
|
|
}
|
|
|
|
LogicalResult NVVM::WMMAStoreOp::verify() {
|
|
unsigned addressSpace =
|
|
ptr().getType().cast<LLVM::LLVMPointerType>().getAddressSpace();
|
|
if (addressSpace != 0 && addressSpace != 1 && addressSpace != 3)
|
|
return emitOpError("expected operands to be a source pointer in memory "
|
|
"space 0, 1, 3");
|
|
|
|
if (NVVM::WMMAStoreOp::getIntrinsicID(m(), n(), k(), layout(), eltype()) == 0)
|
|
return emitOpError() << "invalid attribute combination";
|
|
std::pair<Type, unsigned> typeInfo =
|
|
inferMMAType(eltype(), NVVM::MMAFrag::c, getContext());
|
|
if (args().size() != typeInfo.second)
|
|
return emitOpError() << "expected " << typeInfo.second << " data operands";
|
|
if (llvm::any_of(args(), [&typeInfo](Value operands) {
|
|
return operands.getType() != typeInfo.first;
|
|
}))
|
|
return emitOpError() << "expected data operands of type " << typeInfo.first;
|
|
return success();
|
|
}
|
|
|
|
LogicalResult NVVM::WMMAMmaOp::verify() {
|
|
if (NVVM::WMMAMmaOp::getIntrinsicID(m(), n(), k(), layoutA(), layoutB(),
|
|
eltypeA(), eltypeB()) == 0)
|
|
return emitOpError() << "invalid attribute combination";
|
|
std::pair<Type, unsigned> typeInfoA =
|
|
inferMMAType(eltypeA(), NVVM::MMAFrag::a, getContext());
|
|
std::pair<Type, unsigned> typeInfoB =
|
|
inferMMAType(eltypeA(), NVVM::MMAFrag::b, getContext());
|
|
std::pair<Type, unsigned> typeInfoC =
|
|
inferMMAType(eltypeB(), NVVM::MMAFrag::c, getContext());
|
|
SmallVector<Type, 32> arguments;
|
|
arguments.append(typeInfoA.second, typeInfoA.first);
|
|
arguments.append(typeInfoB.second, typeInfoB.first);
|
|
arguments.append(typeInfoC.second, typeInfoC.first);
|
|
unsigned numArgs = arguments.size();
|
|
if (args().size() != numArgs)
|
|
return emitOpError() << "expected " << numArgs << " arguments";
|
|
for (unsigned i = 0; i < numArgs; i++) {
|
|
if (args()[i].getType() != arguments[i])
|
|
return emitOpError() << "expected argument " << i << " to be of type "
|
|
<< arguments[i];
|
|
}
|
|
Type dstType = LLVM::LLVMStructType::getLiteral(
|
|
getContext(), SmallVector<Type, 8>(typeInfoC.second, typeInfoC.first));
|
|
if (getType() != dstType)
|
|
return emitOpError("expected destination type is a structure of ")
|
|
<< typeInfoC.second << " elements of type " << typeInfoC.first;
|
|
return success();
|
|
}
|
|
|
|
LogicalResult NVVM::LdMatrixOp::verify() {
|
|
unsigned addressSpace =
|
|
ptr().getType().cast<LLVM::LLVMPointerType>().getAddressSpace();
|
|
if (addressSpace != 3)
|
|
return emitOpError("expected source pointer in memory space 3");
|
|
|
|
if (num() != 1 && num() != 2 && num() != 4)
|
|
return emitOpError("expected num attribute to be 1, 2 or 4");
|
|
|
|
Type i32 = IntegerType::get(getContext(), 32);
|
|
if (num() == 1 && getType() != i32)
|
|
return emitOpError("expected destination type is i32");
|
|
if (num() == 2 || num() == 4) {
|
|
Type dstType = LLVM::LLVMStructType::getLiteral(
|
|
getContext(), SmallVector<Type>(num(), i32));
|
|
if (getType() != dstType)
|
|
return emitOpError("expected destination type is a structure of ")
|
|
<< num() << " elements of type i32";
|
|
}
|
|
return success();
|
|
}
|
|
|
|
//===----------------------------------------------------------------------===//
|
|
// NVVMDialect initialization, type parsing, and registration.
|
|
//===----------------------------------------------------------------------===//
|
|
|
|
// TODO: This should be the llvm.nvvm dialect once this is supported.
|
|
void NVVMDialect::initialize() {
|
|
addOperations<
|
|
#define GET_OP_LIST
|
|
#include "mlir/Dialect/LLVMIR/NVVMOps.cpp.inc"
|
|
>();
|
|
addAttributes<
|
|
#define GET_ATTRDEF_LIST
|
|
#include "mlir/Dialect/LLVMIR/NVVMOpsAttributes.cpp.inc"
|
|
>();
|
|
|
|
// Support unknown operations because not all NVVM operations are
|
|
// registered.
|
|
allowUnknownOperations();
|
|
}
|
|
|
|
LogicalResult NVVMDialect::verifyOperationAttribute(Operation *op,
|
|
NamedAttribute attr) {
|
|
// Kernel function attribute should be attached to functions.
|
|
if (attr.getName() == NVVMDialect::getKernelFuncAttrName()) {
|
|
if (!isa<LLVM::LLVMFuncOp>(op)) {
|
|
return op->emitError() << "'" << NVVMDialect::getKernelFuncAttrName()
|
|
<< "' attribute attached to unexpected op";
|
|
}
|
|
}
|
|
return success();
|
|
}
|
|
|
|
#define GET_OP_CLASSES
|
|
#include "mlir/Dialect/LLVMIR/NVVMOps.cpp.inc"
|
|
|
|
#define GET_ATTRDEF_CLASSES
|
|
#include "mlir/Dialect/LLVMIR/NVVMOpsAttributes.cpp.inc"
|