From d950a164a86e5ed591d8b9e185d745d0c5087274 Mon Sep 17 00:00:00 2001 From: Srinivasa Ravi Date: Mon, 1 Sep 2025 19:03:00 +0530 Subject: [PATCH 1/6] [MLIR][NVVM] Add clusterlaunchcontrol Ops This change adds the `clusterlaunchcontrol.try.cancel` and `clusterlaunchcontrol.query.cancel` Ops to the NVVM dialect. Tests are added in `clusterlaunchcontrol.mlir`. PTX Reference: https://docs.nvidia.com/cuda/parallel-thread-execution/#parallel-synchronization-and-communication-instructions-clusterlaunchcontrol-try-cancel --- mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td | 113 ++++++++++++++++++ mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp | 63 ++++++++++ .../LLVMIR/nvvm/clusterlaunchcontrol.mlir | 64 ++++++++++ mlir/test/Target/LLVMIR/nvvmir-invalid.mlir | 16 +++ 4 files changed, 256 insertions(+) create mode 100644 mlir/test/Target/LLVMIR/nvvm/clusterlaunchcontrol.mlir diff --git a/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td b/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td index 9d93b4efe7a5b..da6ed6adb2456 100644 --- a/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td +++ b/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td @@ -4035,6 +4035,119 @@ def NVVM_DotAccumulate2WayOp : NVVM_Op<"dot.accumulate.2way"> { }]; } +//===----------------------------------------------------------------------===// +// NVVM clusterlaunchcontrol Ops. +//===----------------------------------------------------------------------===// + +def NVVM_ClusterLaunchControlTryCancelOp + : NVVM_Op<"clusterlaunchcontrol.try.cancel", [NVVMRequiresSM<100>]> { + let summary = "Request atomically canceling the launch of a cluster that has not started running yet"; + let description = [{ + `clusterlaunchcontrol.try.cancel` requests atomically canceling the launch + of a cluster that has not started running yet. It asynchronously writes an + opaque response to shared memory indicating whether the operation succeeded + or failed. + + Operand `addr` specifies the naturally aligned address of the 16-byte wide + shared memory location where the request's response is written. + + Operand `mbar` specifies the mbarrier object used to track the completion + of the asynchronous operation. + + If `multicast` is specified, the response is asynchronously written to the + corresponding local shared memory location (specifed by `addr`) of each CTA + in the requesting cluster. + + [For more information, see PTX ISA](https://docs.nvidia.com/cuda/parallel-thread-execution/#parallel-synchronization-and-communication-instructions-clusterlaunchcontrol-try-cancel) + }]; + + let arguments = (ins UnitAttr:$multicast, + LLVM_PointerShared: $addr, + LLVM_PointerShared: $mbar); + + let assemblyFormat = "(`multicast` $multicast^ `,`)? $addr `,` $mbar attr-dict"; + + let extraClassDeclaration = [{ + static mlir::NVVM::IDArgPair + getIntrinsicIDAndArgs(Operation &op, LLVM::ModuleTranslation &mt, + llvm::IRBuilderBase &builder); + }]; + + string llvmBuilder = [{ + auto [id, args] = + NVVM::ClusterLaunchControlTryCancelOp::getIntrinsicIDAndArgs( + *op, moduleTranslation, builder); + createIntrinsicCall(builder, id, args); + }]; +} + +def ClusterLaunchControlIsCanceled + : I32EnumCase<"IS_CANCELED", 0, "is_canceled">; +def ClusterLaunchControlGetFirstCTAIDX + : I32EnumCase<"GET_FIRST_CTA_ID_X", 1, "get_first_cta_id_x">; +def ClusterLaunchControlGetFirstCTAIDY + : I32EnumCase<"GET_FIRST_CTA_ID_Y", 2, "get_first_cta_id_y">; +def ClusterLaunchControlGetFirstCTAIDZ + : I32EnumCase<"GET_FIRST_CTA_ID_Z", 3, "get_first_cta_id_z">; + +def ClusterLaunchControlQueryType + : I32Enum<"ClusterLaunchControlQueryType", + "NVVM ClusterLaunchControlQueryType", + [ClusterLaunchControlIsCanceled, ClusterLaunchControlGetFirstCTAIDX, + ClusterLaunchControlGetFirstCTAIDY, ClusterLaunchControlGetFirstCTAIDZ]> { + let cppNamespace = "::mlir::NVVM"; +} + +def ClusterLaunchControlQueryTypeAttr + : EnumAttr { + let assemblyFormat = "$value"; +} + +def NVVM_ClusterLaunchControlQueryCancelOp + : NVVM_Op<"clusterlaunchcontrol.query.cancel", [NVVMRequiresSM<100>]> { + let summary = "Query the response of a clusterlaunchcontrol.try.cancel operation"; + let description = [{ + `clusterlaunchcontrol.query.cancel` queries the response of a + `clusterlaunchcontrol.try.cancel` operation. + + Operand `try_cancel_response` specifies the response of the + `clusterlaunchcontrol.try.cancel` operation to be queried. + + Operand `query_type` specifies the type of query to perform and can be one + of the following: + - `is_canceled` : Returns true if the try cancel request succeeded, + otherwise returns false. + - `get_first_cta_id_{x/y/z}` : Behaviour is defined only if the try cancel + request succeeded. Returns the x, y, or z coordinate of the first CTA in + the canceled cluster. + + [For more information, see PTX ISA](https://docs.nvidia.com/cuda/parallel-thread-execution/#parallel-synchronization-and-communication-instructions-clusterlaunchcontrol-query-cancel) + }]; + + let arguments = (ins DefaultValuedAttr:$query_type, + I128:$try_cancel_response); + let results = (outs AnyTypeOf<[I1, I32]>:$res); + + let assemblyFormat = "(`query` `=` $query_type^ `,`)? $try_cancel_response attr-dict `:` type($res)"; + + let hasVerifier = 1; + + let extraClassDeclaration = [{ + static mlir::NVVM::IDArgPair + getIntrinsicIDAndArgs(Operation &op, LLVM::ModuleTranslation &mt, + llvm::IRBuilderBase &builder); + }]; + + string llvmBuilder = [{ + auto [id, args] = + NVVM::ClusterLaunchControlQueryCancelOp::getIntrinsicIDAndArgs( + *op, moduleTranslation, builder); + $res = createIntrinsicCall(builder, id, args); + }]; +} + //===----------------------------------------------------------------------===// // NVVM target attribute. //===----------------------------------------------------------------------===// diff --git a/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp b/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp index 376e3c3e1fcbe..69c635a972c79 100644 --- a/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp +++ b/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp @@ -1402,6 +1402,24 @@ LogicalResult NVVM::PrefetchOp::verify() { return success(); } +LogicalResult NVVM::ClusterLaunchControlQueryCancelOp::verify() { + switch (getQueryType()) { + case NVVM::ClusterLaunchControlQueryType::IS_CANCELED: + if (!getType().isInteger(1)) + return emitOpError("is_canceled query type returns an i1"); + break; + case NVVM::ClusterLaunchControlQueryType::GET_FIRST_CTA_ID_X: + case NVVM::ClusterLaunchControlQueryType::GET_FIRST_CTA_ID_Y: + case NVVM::ClusterLaunchControlQueryType::GET_FIRST_CTA_ID_Z: + if (!getType().isInteger(32)) { + return emitOpError("get_first_cta_id_x, get_first_cta_id_y, " + "get_first_cta_id_z query types return an i32"); + } + break; + } + return success(); +} + /// Packs the given `field` into the `result`. /// The `result` is 64-bits and each `field` can be 32-bits or narrower. static llvm::Value * @@ -2088,6 +2106,51 @@ bool NVVM::InlinePtxOp::getAsmValues( return false; // No manual mapping needed } +NVVM::IDArgPair ClusterLaunchControlTryCancelOp::getIntrinsicIDAndArgs( + Operation &op, LLVM::ModuleTranslation &mt, llvm::IRBuilderBase &builder) { + auto curOp = cast(op); + llvm::SmallVector args; + args.push_back(mt.lookupValue(curOp.getAddr())); + args.push_back(mt.lookupValue(curOp.getMbar())); + + return curOp.getMulticast() + ? NVVM::IDArgPair( + {llvm::Intrinsic:: + nvvm_clusterlaunchcontrol_try_cancel_async_multicast_shared, + args}) + : NVVM::IDArgPair( + {llvm::Intrinsic:: + nvvm_clusterlaunchcontrol_try_cancel_async_shared, + args}); +} + +NVVM::IDArgPair ClusterLaunchControlQueryCancelOp::getIntrinsicIDAndArgs( + Operation &op, LLVM::ModuleTranslation &mt, llvm::IRBuilderBase &builder) { + auto curOp = cast(op); + llvm::SmallVector args; + args.push_back(mt.lookupValue(curOp.getTryCancelResponse())); + + switch (curOp.getQueryType()) { + case NVVM::ClusterLaunchControlQueryType::IS_CANCELED: + return {llvm::Intrinsic::nvvm_clusterlaunchcontrol_query_cancel_is_canceled, + args}; + case NVVM::ClusterLaunchControlQueryType::GET_FIRST_CTA_ID_X: + return {llvm::Intrinsic:: + nvvm_clusterlaunchcontrol_query_cancel_get_first_ctaid_x, + args}; + case NVVM::ClusterLaunchControlQueryType::GET_FIRST_CTA_ID_Y: + return {llvm::Intrinsic:: + nvvm_clusterlaunchcontrol_query_cancel_get_first_ctaid_y, + args}; + case NVVM::ClusterLaunchControlQueryType::GET_FIRST_CTA_ID_Z: + return {llvm::Intrinsic:: + nvvm_clusterlaunchcontrol_query_cancel_get_first_ctaid_z, + args}; + default: + llvm_unreachable("Invalid query type"); + } +} + //===----------------------------------------------------------------------===// // NVVMDialect initialization, type parsing, and registration. //===----------------------------------------------------------------------===// diff --git a/mlir/test/Target/LLVMIR/nvvm/clusterlaunchcontrol.mlir b/mlir/test/Target/LLVMIR/nvvm/clusterlaunchcontrol.mlir new file mode 100644 index 0000000000000..3100231e0de2f --- /dev/null +++ b/mlir/test/Target/LLVMIR/nvvm/clusterlaunchcontrol.mlir @@ -0,0 +1,64 @@ +// RUN: mlir-translate -mlir-to-llvmir %s | FileCheck %s + +llvm.func @clusterlaunchcontrol_try_cancel(%addr: !llvm.ptr<3>, %mbar: !llvm.ptr<3>) { + // CHECK-LABEL: define void @clusterlaunchcontrol_try_cancel(ptr addrspace(3) %0, ptr addrspace(3) %1) { + // CHECK-NEXT: call void @llvm.nvvm.clusterlaunchcontrol.try_cancel.async.shared(ptr addrspace(3) %0, ptr addrspace(3) %1) + // CHECK-NEXT: ret void + // CHECK-NEXT: } + nvvm.clusterlaunchcontrol.try.cancel %addr, %mbar + llvm.return +} + +llvm.func @clusterlaunchcontrol_try_cancel_multicast(%addr: !llvm.ptr<3>, %mbar: !llvm.ptr<3>) { + // CHECK-LABEL: define void @clusterlaunchcontrol_try_cancel_multicast(ptr addrspace(3) %0, ptr addrspace(3) %1) { + // CHECK-NEXT: call void @llvm.nvvm.clusterlaunchcontrol.try_cancel.async.multicast.shared(ptr addrspace(3) %0, ptr addrspace(3) %1) + // CHECK-NEXT: ret void + // CHECK-NEXT: } + nvvm.clusterlaunchcontrol.try.cancel multicast, %addr, %mbar + llvm.return +} + +llvm.func @clusterlaunchcontrol_query_cancel(%try_cancel_response: i128) { + // CHECK-LABEL: define void @clusterlaunchcontrol_query_cancel(i128 %0) { + // CHECK-NEXT: %2 = call i1 @llvm.nvvm.clusterlaunchcontrol.query_cancel.is_canceled(i128 %0) + // CHECK-NEXT: ret void + // CHECK-NEXT: } + nvvm.clusterlaunchcontrol.query.cancel %try_cancel_response : i1 + llvm.return +} + +llvm.func @clusterlaunchcontrol_query_cancel_is_canceled(%try_cancel_response: i128) { + // CHECK-LABEL: define void @clusterlaunchcontrol_query_cancel_is_canceled(i128 %0) { + // CHECK-NEXT: %2 = call i1 @llvm.nvvm.clusterlaunchcontrol.query_cancel.is_canceled(i128 %0) + // CHECK-NEXT: ret void + // CHECK-NEXT: } + %res = nvvm.clusterlaunchcontrol.query.cancel query = is_canceled, %try_cancel_response : i1 + llvm.return +} + +llvm.func @clusterlaunchcontrol_query_cancel_get_first_cta_id_x(%try_cancel_response: i128) { + // CHECK-LABEL: define void @clusterlaunchcontrol_query_cancel_get_first_cta_id_x(i128 %0) { + // CHECK-NEXT: %2 = call i32 @llvm.nvvm.clusterlaunchcontrol.query_cancel.get_first_ctaid.x(i128 %0) + // CHECK-NEXT: ret void + // CHECK-NEXT: } + %res = nvvm.clusterlaunchcontrol.query.cancel query = get_first_cta_id_x, %try_cancel_response : i32 + llvm.return +} + +llvm.func @clusterlaunchcontrol_query_cancel_get_first_cta_id_y(%try_cancel_response: i128) { + // CHECK-LABEL: define void @clusterlaunchcontrol_query_cancel_get_first_cta_id_y(i128 %0) { + // CHECK-NEXT: %2 = call i32 @llvm.nvvm.clusterlaunchcontrol.query_cancel.get_first_ctaid.y(i128 %0) + // CHECK-NEXT: ret void + // CHECK-NEXT: } + %res = nvvm.clusterlaunchcontrol.query.cancel query = get_first_cta_id_y, %try_cancel_response : i32 + llvm.return +} + +llvm.func @clusterlaunchcontrol_query_cancel_get_first_cta_id_z(%try_cancel_response: i128) { + // CHECK-LABEL: define void @clusterlaunchcontrol_query_cancel_get_first_cta_id_z(i128 %0) { + // CHECK-NEXT: %2 = call i32 @llvm.nvvm.clusterlaunchcontrol.query_cancel.get_first_ctaid.z(i128 %0) + // CHECK-NEXT: ret void + // CHECK-NEXT: } + %res = nvvm.clusterlaunchcontrol.query.cancel query = get_first_cta_id_z, %try_cancel_response : i32 + llvm.return +} diff --git a/mlir/test/Target/LLVMIR/nvvmir-invalid.mlir b/mlir/test/Target/LLVMIR/nvvmir-invalid.mlir index b35a6dbcca286..383f4829f3287 100644 --- a/mlir/test/Target/LLVMIR/nvvmir-invalid.mlir +++ b/mlir/test/Target/LLVMIR/nvvmir-invalid.mlir @@ -535,3 +535,19 @@ llvm.func @nanosleep() { nvvm.nanosleep 100000000000000 llvm.return } + +// ----- + +llvm.func @clusterlaunchcontrol_query_cancel_is_canceled_invalid_return_type(%try_cancel_response: i128) { + // expected-error@+1 {{'nvvm.clusterlaunchcontrol.query.cancel' op is_canceled query type returns an i1}} + %res = nvvm.clusterlaunchcontrol.query.cancel query = is_canceled, %try_cancel_response : i32 + llvm.return +} + +// ----- + +llvm.func @clusterlaunchcontrol_query_cancel_get_first_cta_id_invalid_return_type(%try_cancel_response: i128) { + // expected-error@+1 {{'nvvm.clusterlaunchcontrol.query.cancel' op get_first_cta_id_x, get_first_cta_id_y, get_first_cta_id_z query types return an i32}} + %res = nvvm.clusterlaunchcontrol.query.cancel query = get_first_cta_id_x, %try_cancel_response : i1 + llvm.return +} From b89266ed9ea6e9ab7d2159a55c0482e1400ca813 Mon Sep 17 00:00:00 2001 From: Srinivasa Ravi Date: Wed, 3 Sep 2025 11:28:53 +0530 Subject: [PATCH 2/6] remove default case from switch statement as all cases are covered --- mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp | 2 -- 1 file changed, 2 deletions(-) diff --git a/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp b/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp index 69c635a972c79..3da0d6440d13b 100644 --- a/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp +++ b/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp @@ -2146,8 +2146,6 @@ NVVM::IDArgPair ClusterLaunchControlQueryCancelOp::getIntrinsicIDAndArgs( return {llvm::Intrinsic:: nvvm_clusterlaunchcontrol_query_cancel_get_first_ctaid_z, args}; - default: - llvm_unreachable("Invalid query type"); } } From b0e96ab254a16c6f7d0a5db7166d501a3e39d0ee Mon Sep 17 00:00:00 2001 From: Srinivasa Ravi Date: Thu, 4 Sep 2025 10:02:36 +0530 Subject: [PATCH 3/6] rename arguments --- mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td | 4 ++-- mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp | 4 ++-- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td b/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td index da6ed6adb2456..ce8ebfa9d23e4 100644 --- a/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td +++ b/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td @@ -4062,8 +4062,8 @@ def NVVM_ClusterLaunchControlTryCancelOp }]; let arguments = (ins UnitAttr:$multicast, - LLVM_PointerShared: $addr, - LLVM_PointerShared: $mbar); + LLVM_PointerShared: $smemAddress, + LLVM_PointerShared: $mbarrier); let assemblyFormat = "(`multicast` $multicast^ `,`)? $addr `,` $mbar attr-dict"; diff --git a/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp b/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp index 3da0d6440d13b..e500ab83985ac 100644 --- a/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp +++ b/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp @@ -2110,8 +2110,8 @@ NVVM::IDArgPair ClusterLaunchControlTryCancelOp::getIntrinsicIDAndArgs( Operation &op, LLVM::ModuleTranslation &mt, llvm::IRBuilderBase &builder) { auto curOp = cast(op); llvm::SmallVector args; - args.push_back(mt.lookupValue(curOp.getAddr())); - args.push_back(mt.lookupValue(curOp.getMbar())); + args.push_back(mt.lookupValue(curOp.getSmemAddress())); + args.push_back(mt.lookupValue(curOp.getMbarrrier())); return curOp.getMulticast() ? NVVM::IDArgPair( From 6dfb5045d754b7adfef9e2e329a55bb09f969367 Mon Sep 17 00:00:00 2001 From: Srinivasa Ravi Date: Thu, 4 Sep 2025 12:05:46 +0530 Subject: [PATCH 4/6] fix assembly format --- mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td b/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td index ce8ebfa9d23e4..5e941c4ca005b 100644 --- a/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td +++ b/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td @@ -4065,7 +4065,7 @@ def NVVM_ClusterLaunchControlTryCancelOp LLVM_PointerShared: $smemAddress, LLVM_PointerShared: $mbarrier); - let assemblyFormat = "(`multicast` $multicast^ `,`)? $addr `,` $mbar attr-dict"; + let assemblyFormat = "(`multicast` $multicast^ `,`)? $smemAddress `,` $mbarrier attr-dict"; let extraClassDeclaration = [{ static mlir::NVVM::IDArgPair From ebd0052caad65f30d3eb758770b2b7b055ca76b1 Mon Sep 17 00:00:00 2001 From: Srinivasa Ravi Date: Thu, 4 Sep 2025 14:21:46 +0530 Subject: [PATCH 5/6] fix typo --- mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp b/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp index e500ab83985ac..582f737dbfcce 100644 --- a/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp +++ b/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp @@ -2111,7 +2111,7 @@ NVVM::IDArgPair ClusterLaunchControlTryCancelOp::getIntrinsicIDAndArgs( auto curOp = cast(op); llvm::SmallVector args; args.push_back(mt.lookupValue(curOp.getSmemAddress())); - args.push_back(mt.lookupValue(curOp.getMbarrrier())); + args.push_back(mt.lookupValue(curOp.getMbarrier())); return curOp.getMulticast() ? NVVM::IDArgPair( From 65eda892aa139e457ca0457cb2389ee35c751adb Mon Sep 17 00:00:00 2001 From: Srinivasa Ravi Date: Tue, 9 Sep 2025 17:42:55 +0530 Subject: [PATCH 6/6] address comments --- mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td | 27 ++++++------ mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp | 42 ++++++++++--------- .../LLVMIR/nvvm/clusterlaunchcontrol.mlir | 9 ---- 3 files changed, 34 insertions(+), 44 deletions(-) diff --git a/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td b/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td index 5e941c4ca005b..5011e77b70405 100644 --- a/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td +++ b/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td @@ -4048,11 +4048,11 @@ def NVVM_ClusterLaunchControlTryCancelOp opaque response to shared memory indicating whether the operation succeeded or failed. - Operand `addr` specifies the naturally aligned address of the 16-byte wide - shared memory location where the request's response is written. + Operand `smemAddress` specifies the naturally aligned address of the + 16-byte wide shared memory location where the request's response is written. - Operand `mbar` specifies the mbarrier object used to track the completion - of the asynchronous operation. + Operand `mbarrier` specifies the mbarrier object used to track the + completion of the asynchronous operation. If `multicast` is specified, the response is asynchronously written to the corresponding local shared memory location (specifed by `addr`) of each CTA @@ -4109,28 +4109,25 @@ def NVVM_ClusterLaunchControlQueryCancelOp let summary = "Query the response of a clusterlaunchcontrol.try.cancel operation"; let description = [{ `clusterlaunchcontrol.query.cancel` queries the response of a - `clusterlaunchcontrol.try.cancel` operation. - - Operand `try_cancel_response` specifies the response of the - `clusterlaunchcontrol.try.cancel` operation to be queried. + `clusterlaunchcontrol.try.cancel` operation specified by operand + `try_cancel_response`. Operand `query_type` specifies the type of query to perform and can be one of the following: - `is_canceled` : Returns true if the try cancel request succeeded, - otherwise returns false. - - `get_first_cta_id_{x/y/z}` : Behaviour is defined only if the try cancel - request succeeded. Returns the x, y, or z coordinate of the first CTA in - the canceled cluster. + and false otherwise. + - `get_first_cta_id_{x/y/z}` : Returns the x, y, or z coordinate of the + first CTA in the canceled cluster. Behaviour is defined only if the try + cancel request succeeded. [For more information, see PTX ISA](https://docs.nvidia.com/cuda/parallel-thread-execution/#parallel-synchronization-and-communication-instructions-clusterlaunchcontrol-query-cancel) }]; - let arguments = (ins DefaultValuedAttr:$query_type, + let arguments = (ins ClusterLaunchControlQueryTypeAttr:$query_type, I128:$try_cancel_response); let results = (outs AnyTypeOf<[I1, I32]>:$res); - let assemblyFormat = "(`query` `=` $query_type^ `,`)? $try_cancel_response attr-dict `:` type($res)"; + let assemblyFormat = "`query` `=` $query_type `,` $try_cancel_response attr-dict `:` type($res)"; let hasVerifier = 1; diff --git a/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp b/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp index 582f737dbfcce..82b14ac8cfc97 100644 --- a/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp +++ b/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp @@ -2113,15 +2113,13 @@ NVVM::IDArgPair ClusterLaunchControlTryCancelOp::getIntrinsicIDAndArgs( args.push_back(mt.lookupValue(curOp.getSmemAddress())); args.push_back(mt.lookupValue(curOp.getMbarrier())); - return curOp.getMulticast() - ? NVVM::IDArgPair( - {llvm::Intrinsic:: - nvvm_clusterlaunchcontrol_try_cancel_async_multicast_shared, - args}) - : NVVM::IDArgPair( - {llvm::Intrinsic:: - nvvm_clusterlaunchcontrol_try_cancel_async_shared, - args}); + llvm::Intrinsic::ID intrinsicID = + curOp.getMulticast() + ? llvm::Intrinsic:: + nvvm_clusterlaunchcontrol_try_cancel_async_multicast_shared + : llvm::Intrinsic::nvvm_clusterlaunchcontrol_try_cancel_async_shared; + + return {intrinsicID, args}; } NVVM::IDArgPair ClusterLaunchControlQueryCancelOp::getIntrinsicIDAndArgs( @@ -2130,23 +2128,27 @@ NVVM::IDArgPair ClusterLaunchControlQueryCancelOp::getIntrinsicIDAndArgs( llvm::SmallVector args; args.push_back(mt.lookupValue(curOp.getTryCancelResponse())); + llvm::Intrinsic::ID intrinsicID; + switch (curOp.getQueryType()) { case NVVM::ClusterLaunchControlQueryType::IS_CANCELED: - return {llvm::Intrinsic::nvvm_clusterlaunchcontrol_query_cancel_is_canceled, - args}; + intrinsicID = + llvm::Intrinsic::nvvm_clusterlaunchcontrol_query_cancel_is_canceled; + break; case NVVM::ClusterLaunchControlQueryType::GET_FIRST_CTA_ID_X: - return {llvm::Intrinsic:: - nvvm_clusterlaunchcontrol_query_cancel_get_first_ctaid_x, - args}; + intrinsicID = llvm::Intrinsic:: + nvvm_clusterlaunchcontrol_query_cancel_get_first_ctaid_x; + break; case NVVM::ClusterLaunchControlQueryType::GET_FIRST_CTA_ID_Y: - return {llvm::Intrinsic:: - nvvm_clusterlaunchcontrol_query_cancel_get_first_ctaid_y, - args}; + intrinsicID = llvm::Intrinsic:: + nvvm_clusterlaunchcontrol_query_cancel_get_first_ctaid_y; + break; case NVVM::ClusterLaunchControlQueryType::GET_FIRST_CTA_ID_Z: - return {llvm::Intrinsic:: - nvvm_clusterlaunchcontrol_query_cancel_get_first_ctaid_z, - args}; + intrinsicID = llvm::Intrinsic:: + nvvm_clusterlaunchcontrol_query_cancel_get_first_ctaid_z; + break; } + return {intrinsicID, args}; } //===----------------------------------------------------------------------===// diff --git a/mlir/test/Target/LLVMIR/nvvm/clusterlaunchcontrol.mlir b/mlir/test/Target/LLVMIR/nvvm/clusterlaunchcontrol.mlir index 3100231e0de2f..6ff85fa38a0a0 100644 --- a/mlir/test/Target/LLVMIR/nvvm/clusterlaunchcontrol.mlir +++ b/mlir/test/Target/LLVMIR/nvvm/clusterlaunchcontrol.mlir @@ -18,15 +18,6 @@ llvm.func @clusterlaunchcontrol_try_cancel_multicast(%addr: !llvm.ptr<3>, %mbar: llvm.return } -llvm.func @clusterlaunchcontrol_query_cancel(%try_cancel_response: i128) { - // CHECK-LABEL: define void @clusterlaunchcontrol_query_cancel(i128 %0) { - // CHECK-NEXT: %2 = call i1 @llvm.nvvm.clusterlaunchcontrol.query_cancel.is_canceled(i128 %0) - // CHECK-NEXT: ret void - // CHECK-NEXT: } - nvvm.clusterlaunchcontrol.query.cancel %try_cancel_response : i1 - llvm.return -} - llvm.func @clusterlaunchcontrol_query_cancel_is_canceled(%try_cancel_response: i128) { // CHECK-LABEL: define void @clusterlaunchcontrol_query_cancel_is_canceled(i128 %0) { // CHECK-NEXT: %2 = call i1 @llvm.nvvm.clusterlaunchcontrol.query_cancel.is_canceled(i128 %0)