Skip to content

Commit

Permalink
Revert "[mlir][NVVM] Disallow results on kernel functions (llvm#96399)…
Browse files Browse the repository at this point in the history
…" (llvm#97074)

NVVM IR itself doesn't place any restriction that a function annotated
as `nvvm.kernel` actually has no results, so this is a mismatch at the
NVVMDialect level and NVVMIR. The GPU dialect might enforce that kernel
functions have no results, but it doesn't make sense to propagate this
constraint downstream.

Reverts llvm#96399
  • Loading branch information
Mogball authored and lravenclaw committed Jul 3, 2024
1 parent 6c9765b commit 9c63882
Show file tree
Hide file tree
Showing 2 changed files with 3 additions and 17 deletions.
13 changes: 3 additions & 10 deletions mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -214,8 +214,7 @@ void MmaOp::print(OpAsmPrinter &p) {
p.printOptionalAttrDict(this->getOperation()->getAttrs(), ignoreAttrNames);

// Print the types of the operands and result.
p << " : "
<< "(";
p << " : " << "(";
llvm::interleaveComma(SmallVector<Type, 3>{frags[0].regs[0].getType(),
frags[1].regs[0].getType(),
frags[2].regs[0].getType()},
Expand Down Expand Up @@ -956,9 +955,7 @@ std::string NVVM::WgmmaMmaAsyncOp::getPtx() {
ss << "},";
// Need to map read/write registers correctly.
regCnt = (regCnt * 2);
ss << " $" << (regCnt) << ","
<< " $" << (regCnt + 1) << ","
<< " p";
ss << " $" << (regCnt) << "," << " $" << (regCnt + 1) << "," << " p";
if (getTypeD() != WGMMATypes::s32) {
ss << ", $" << (regCnt + 3) << ", $" << (regCnt + 4);
}
Expand Down Expand Up @@ -1056,14 +1053,10 @@ LogicalResult NVVMDialect::verifyOperationAttribute(Operation *op,
StringAttr attrName = attr.getName();
// Kernel function attribute should be attached to functions.
if (attrName == NVVMDialect::getKernelFuncAttrName()) {
auto funcOp = dyn_cast<LLVM::LLVMFuncOp>(op);
if (!funcOp) {
if (!isa<LLVM::LLVMFuncOp>(op)) {
return op->emitError() << "'" << NVVMDialect::getKernelFuncAttrName()
<< "' attribute attached to unexpected op";
}
if (!funcOp.getResultTypes().empty()) {
return op->emitError() << "kernel function cannot have results";
}
}
// If maxntid and reqntid exist, it must be an array with max 3 dim
if (attrName == NVVMDialect::getMaxntidAttrName() ||
Expand Down
7 changes: 0 additions & 7 deletions mlir/test/Target/LLVMIR/nvvmir.mlir
Original file line number Diff line number Diff line change
Expand Up @@ -574,10 +574,3 @@ llvm.func @kernel_func(%arg0: !llvm.ptr {llvm.byval = i32, nvvm.grid_constant})
llvm.func @kernel_func(%arg0: !llvm.ptr {llvm.byval = i32, nvvm.grid_constant}, %arg1: f32, %arg2: !llvm.ptr {llvm.byval = f32, nvvm.grid_constant}) attributes {nvvm.kernel} {
llvm.return
}

// -----

// expected-error @below{{kernel function cannot have results}}
llvm.func @kernel_with_result(%i: i32) -> i32 attributes {nvvm.kernel} {
llvm.return %i : i32
}

0 comments on commit 9c63882

Please sign in to comment.