Skip to content

Commit

Permalink
[MLIR][NVVM]: Update setmaxregister NVVM Op (llvm#77594)
Browse files Browse the repository at this point in the history
This patch updates the setmaxregister NVVM Op to use the
intrinsics instead of inline-ptx.

* The interface remains same (as expected).
* Tests are added to verify the lowered intrinsics in
Target/LLVMIR/nvvmir.mlir.

Signed-off-by: Durgadoss R <durgadossr@nvidia.com>
  • Loading branch information
durga4github authored Jan 10, 2024
1 parent e203968 commit 6a075a9
Show file tree
Hide file tree
Showing 3 changed files with 20 additions and 10 deletions.
16 changes: 8 additions & 8 deletions mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
Original file line number Diff line number Diff line change
Expand Up @@ -463,17 +463,17 @@ def SetMaxRegisterAction : I32EnumAttr<"SetMaxRegisterAction", "NVVM set max reg
}
def SetMaxRegisterActionAttr : EnumAttr<NVVM_Dialect, SetMaxRegisterAction, "action">;

def NVVM_SetMaxRegisterOp : NVVM_PTXBuilder_Op<"setmaxregister"> {
def NVVM_SetMaxRegisterOp : NVVM_Op<"setmaxregister"> {
let arguments = (ins I32Attr:$regCount, SetMaxRegisterActionAttr:$action);
let assemblyFormat = "$action $regCount attr-dict";
let extraClassDefinition = [{
std::string $cppClass::getPtx() {
if(getAction() == NVVM::SetMaxRegisterAction::increase)
return std::string("setmaxnreg.inc.sync.aligned.u32 %0;");
return std::string("setmaxnreg.dec.sync.aligned.u32 %0;");
}
}];
let hasVerifier = 1;
string llvmBuilder = [{
auto intId = (op.getAction() == NVVM::SetMaxRegisterAction::increase) ?
llvm::Intrinsic::nvvm_setmaxnreg_inc_sync_aligned_u32 :
llvm::Intrinsic::nvvm_setmaxnreg_dec_sync_aligned_u32;

createIntrinsicCall(builder, intId, builder.getInt32($regCount));
}];
}

def NVVM_FenceMbarrierInitOp : NVVM_PTXBuilder_Op<"fence.mbarrier.init"> {
Expand Down
5 changes: 3 additions & 2 deletions mlir/test/Conversion/NVVMToLLVM/nvvm-to-llvm.mlir
Original file line number Diff line number Diff line change
Expand Up @@ -628,9 +628,10 @@ llvm.func @init_mbarrier_arrive_expect_tx(%desc : !llvm.ptr, %pred : i1) {
// -----

func.func @set_max_register() {
//CHECK: llvm.inline_asm has_side_effects asm_dialect = att "setmaxnreg.inc.sync.aligned.u32 $0;", "n"
// CHECK: nvvm.setmaxregister increase 232
nvvm.setmaxregister increase 232
//CHECK: llvm.inline_asm has_side_effects asm_dialect = att "setmaxnreg.dec.sync.aligned.u32 $0;", "n"

// CHECK: nvvm.setmaxregister decrease 40
nvvm.setmaxregister decrease 40
func.return
}
Expand Down
9 changes: 9 additions & 0 deletions mlir/test/Target/LLVMIR/nvvmir.mlir
Original file line number Diff line number Diff line change
Expand Up @@ -369,6 +369,15 @@ llvm.func @cp_async_mbarrier_arrive(%bar_shared: !llvm.ptr<3>, %bar_gen: !llvm.p
llvm.return
}

// CHECK-LABEL: @llvm_nvvm_setmaxregister
llvm.func @llvm_nvvm_setmaxregister() {
// CHECK-LLVM: call void @llvm.nvvm.setmaxnreg.inc.sync.aligned.u32(i32 256)
nvvm.setmaxregister increase 256
// CHECK-LLVM: call void @llvm.nvvm.setmaxnreg.dec.sync.aligned.u32(i32 24)
nvvm.setmaxregister decrease 24
llvm.return
}

// CHECK-LABEL: @ld_matrix
llvm.func @ld_matrix(%arg0: !llvm.ptr<3>) {
// CHECK: call i32 @llvm.nvvm.ldmatrix.sync.aligned.m8n8.x1.b16.p3(ptr addrspace(3) %{{.*}})
Expand Down

0 comments on commit 6a075a9

Please sign in to comment.