This is an automated email from the ASF dual-hosted git repository. tqchen pushed a commit to branch tqchen/printer-doc-translation in repository https://gitbox.apache.org/repos/asf/tvm.git
commit 35d4c0b6530c43ebc9b3350e2561e3ee99884392 Author: Tianqi Chen <[email protected]> AuthorDate: Wed Sep 30 20:44:24 2026 +0000 [REFACTOR][CUDA] Reuse inference for fixed intrinsic result types --- src/backend/cuda/op/target_builtin.cc | 272 +++++----------------------------- 1 file changed, 36 insertions(+), 236 deletions(-) diff --git a/src/backend/cuda/op/target_builtin.cc b/src/backend/cuda/op/target_builtin.cc index 93d636423a..71cfe281be 100644 --- a/src/backend/cuda/op/target_builtin.cc +++ b/src/backend/cuda/op/target_builtin.cc @@ -40,51 +40,11 @@ namespace builtin { namespace { void RegisterDeviceIntrinsicAliases(); -Type InferTypeCudaSync(const CallNode* call) { - TVM_FFI_CHECK_EQ(call->args.size(), 0U, ValueError) - << "CUDA synchronization intrinsic expects no arguments"; - return PrimType::Void(); -} +Type InferTypeVoid(const CallNode*) { return PrimType::Void(); } -Type InferTypeCudaElectSync(const CallNode* call) { - TVM_FFI_CHECK_EQ(call->args.size(), 0U, ValueError) << "cuda.elect_sync expects no arguments"; - return PrimType::UInt(32); -} +Type InferTypeInt32(const CallNode*) { return PrimType::Int(32); } -Type InferTypeNvshmemFence(const CallNode* call) { - TVM_FFI_CHECK_EQ(call->args.size(), 0U, ValueError) << "nvshmem.fence expects no arguments"; - return PrimType::Void(); -} - -Type InferTypeCudaPrintf(const CallNode* call) { - TVM_FFI_CHECK_GE(call->args.size(), 1U, ValueError) << "cuda.printf expects a format string"; - TVM_FFI_CHECK(call->args[0].as<StringImmNode>(), ValueError) - << "cuda.printf expects a literal format string"; - for (size_t i = 1; i < call->args.size(); ++i) { - PrimType arg = call->args[i]->ty.as_or_throw<PrimType>(); - bool integer = arg.MatchesCode(DLDataTypeCode::kDLInt, DLDataTypeCode::kDLUInt) && - (arg.bits() == 8 || arg.bits() == 16 || arg.bits() == 32 || arg.bits() == 64); - bool floating = - arg.MatchesCode(DLDataTypeCode::kDLFloat) && (arg.bits() == 32 || arg.bits() == 64); - TVM_FFI_CHECK(arg.IsScalar() && - (integer || floating || arg.MatchesElementType(DLDataTypeCode::kDLBool, 8)), - ValueError) - << "cuda.printf expects supported scalar numeric arguments"; - } - return PrimType::Void(); -} - -Type InferTypeCudaSyncthreadsPredicate(const CallNode* call) { - TVM_FFI_CHECK_EQ(call->args.size(), 1U, ValueError) - << "CUDA predicate synchronization expects one argument"; - PrimType predicate = call->args[0]->ty.as_or_throw<PrimType>(); - TVM_FFI_CHECK( - predicate.IsScalar() && predicate.MatchesCode(DLDataTypeCode::kDLBool, DLDataTypeCode::kDLInt, - DLDataTypeCode::kDLUInt), - ValueError) - << "CUDA predicate synchronization expects a scalar integral predicate"; - return PrimType::Int(32); -} +Type InferTypeUInt32(const CallNode*) { return PrimType::UInt(32); } void CheckScalarInteger(const Expr& arg) { PrimType type = arg->ty.as_or_throw<PrimType>(); @@ -111,70 +71,6 @@ void CheckPointerElement(const Expr& arg, PrimType expected) { << "Device intrinsic pointer element type does not match its value type"; } -void CheckStringChoice(const Expr& arg, std::initializer_list<const char*> choices) { - const auto* value = arg.as<StringImmNode>(); - TVM_FFI_CHECK(value, ValueError) << "Device intrinsic expects a literal string argument"; - for (const char* choice : choices) { - if (value->value == choice) return; - } - TVM_FFI_THROW(ValueError) << "Unsupported device intrinsic string argument: " << value->value; -} - -Type InferTypeCudaMatrixDescriptor(const CallNode* call) { - TVM_FFI_CHECK_EQ(call->args.size(), 5U, ValueError) - << "CUDA matrix descriptor expects five arguments"; - CheckPointerArgument(call->args[0]); - CheckPointerArgument(call->args[1]); - for (size_t i = 2; i < 5; ++i) CheckScalarInteger(call->args[i]); - return PrimType::Void(); -} - -template <bool block_scaled> -Type InferTypeCudaInstructionDescriptor(const CallNode* call) { - constexpr size_t size = block_scaled ? 17 : 14; - constexpr size_t dtype_end = block_scaled ? 6 : 4; - constexpr size_t fields_begin = block_scaled ? 8 : 4; - constexpr size_t cta_group = block_scaled ? 13 : 9; - TVM_FFI_CHECK_EQ(call->args.size(), size, ValueError) - << "CUDA instruction descriptor has an invalid argument count"; - CheckPointerArgument(call->args[0]); - for (size_t i = 1; i < dtype_end; ++i) { - TVM_FFI_CHECK(call->args[i].as<StringImmNode>(), ValueError) - << "CUDA instruction descriptor expects literal dtype names"; - } - if constexpr (block_scaled) { - for (size_t i : {6U, 7U}) { - if (call->args[i]->ty.as<PointerTypeNode>()) { - CheckPointerArgument(call->args[i]); - } else { - CheckScalarInteger(call->args[i]); - } - } - } - for (size_t i = fields_begin; i < size; ++i) CheckScalarInteger(call->args[i]); - if (const auto* value = call->args[cta_group].as<IntImmNode>()) { - TVM_FFI_CHECK(value->value == 1 || value->value == 2, ValueError) - << "CUDA instruction descriptor expects one or two CTA groups"; - } - return PrimType::Void(); -} - -Type InferTypeCudaWgmmaNoopBarrier(const CallNode* call) { - TVM_FFI_CHECK_EQ(call->args.size(), 1U, ValueError) << "WGMMA noop barrier expects one argument"; - PrimType type = call->args[0]->ty.as_or_throw<PrimType>(); - TVM_FFI_CHECK( - type.IsScalar() && !type.IsVoid() && !type.MatchesCode(DLDataTypeCode::kDLOpaqueHandle), - ValueError) - << "WGMMA noop barrier expects a primitive scalar register"; - return PrimType::Void(); -} - -Type InferTypeCudaNanoSleep(const CallNode* call) { - TVM_FFI_CHECK_EQ(call->args.size(), 1U, ValueError) << "cuda.nano_sleep expects one argument"; - CheckScalarInteger(call->args[0]); - return PrimType::Void(); -} - Type InferTypeCudaLdg(const CallNode* call) { TVM_FFI_CHECK_EQ(call->args.size(), 2U, ValueError) << "Scalar cuda.ldg expects two arguments"; CheckPointerArgument(call->args[0]); @@ -245,85 +141,6 @@ Type InferTypeCudaAtomicCas(const CallNode* call) { return compare; } -Type InferTypeNvshmemQuery(const CallNode* call) { - TVM_FFI_CHECK_EQ(call->args.size(), 0U, ValueError) << "NVSHMEM query expects no arguments"; - return PrimType::Int(32); -} - -Type InferTypeNvshmemSync(const CallNode* call) { - TVM_FFI_CHECK_EQ(call->args.size(), 0U, ValueError) - << "NVSHMEM synchronization expects no arguments"; - return PrimType::Void(); -} - -void CheckNvshmemRmaArguments(const CallNode* call) { - CheckPointerArgument(call->args[0]); - CheckPointerArgument(call->args[1]); - CheckScalarInteger(call->args[2]); -} - -Type InferTypeNvshmemRma(const CallNode* call) { - TVM_FFI_CHECK_EQ(call->args.size(), 4U, ValueError) << "NVSHMEM RMA expects four arguments"; - CheckNvshmemRmaArguments(call); - CheckScalarInteger(call->args[3]); - return PrimType::Void(); -} - -void CheckNvshmemSignalArguments(const CallNode* call, size_t offset) { - CheckPointerElement(call->args[offset], PrimType::UInt(64)); - CheckScalarInteger(call->args[offset + 1]); - CheckStringChoice(call->args[offset + 2], {"set", "add"}); - CheckScalarInteger(call->args[offset + 3]); -} - -Type InferTypeNvshmemSignal(const CallNode* call) { - TVM_FFI_CHECK_EQ(call->args.size(), 4U, ValueError) << "nvshmem.signal_op expects four arguments"; - CheckNvshmemSignalArguments(call, 0); - return PrimType::Void(); -} - -Type InferTypeNvshmemWaitUntil(const CallNode* call) { - TVM_FFI_CHECK_EQ(call->args.size(), 4U, ValueError) - << "nvshmem.wait_until expects four arguments"; - CheckPointerElement(call->args[0], PrimType::UInt(64)); - CheckStringChoice(call->args[1], {"eq", "ne", "gt", "ge", "lt", "le"}); - CheckScalarInteger(call->args[2]); - CheckStringChoice(call->args[3], {"uint64_t", "uint64"}); - return PrimType::Void(); -} - -Type InferTypeNvshmemPutmemSignal(const CallNode* call) { - TVM_FFI_CHECK_EQ(call->args.size(), 7U, ValueError) - << "NVSHMEM putmem signal expects seven arguments"; - CheckNvshmemRmaArguments(call); - CheckNvshmemSignalArguments(call, 3); - return PrimType::Void(); -} - -Type InferTypeCudaMBarrierWait(const CallNode* call) { - TVM_FFI_CHECK_EQ(call->args.size(), 2U, ValueError) << "cuda.mbarrier_wait expects two arguments"; - const Type& barrier = call->args[0]->ty; - if (auto integer = barrier.as<PrimType>()) { - const auto* literal = call->args[0].as<IntImmNode>(); - bool null_pointer = literal && literal->value == 0 && - integer.value().MatchesElementType(DLDataTypeCode::kDLInt, 32); - TVM_FFI_CHECK( - integer.value().IsScalar() && - (integer.value().MatchesElementType(DLDataTypeCode::kDLUInt, 32) || null_pointer), - ValueError) - << "cuda.mbarrier_wait expects a uint32 shared address or pointer"; - } else { - TVM_FFI_CHECK(barrier.as<PointerTypeNode>(), ValueError) - << "cuda.mbarrier_wait expects a uint32 shared address or pointer"; - } - PrimType phase = call->args[1]->ty.as_or_throw<PrimType>(); - TVM_FFI_CHECK( - phase.IsScalar() && phase.MatchesCode(DLDataTypeCode::kDLBool, DLDataTypeCode::kDLInt, - DLDataTypeCode::kDLUInt), - ValueError) - << "cuda.mbarrier_wait expects a scalar phase"; - return PrimType::Void(); -} } // namespace void RegisterCudaTargetBuiltins() { @@ -695,38 +512,42 @@ void RegisterDeviceIntrinsicAliases() { sig::arg("b_offset"), sig::arg("acc_ptr"), sig::arg("c_offset"), sig::arg("saturate"), sig::var_args("args")); - for (const char* name : - {"tirx.cuda.tcgen05_encode_matrix_descriptor", "tirx.cuda.wgmma_encode_matrix_descriptor"}) { - OpDef(name).set_attr<FInferType>("FInferType", - FInferType::FromNative<&InferTypeCudaMatrixDescriptor>()); + for (const char* name : { + "tirx.cuda.tcgen05_encode_matrix_descriptor", + "tirx.cuda.wgmma_encode_matrix_descriptor", + "tirx.cuda.tcgen05_encode_instr_descriptor", + "tirx.cuda.tcgen05_encode_instr_descriptor_block_scaled", + "tirx.cuda.wgmma_noop_barrier", + "tirx.cuda.cluster_sync", + "tirx.cuda.cta_sync", + "tirx.cuda.grid_sync", + "tirx.cuda.warp_sync", + "tirx.cuda.thread_fence", + "tirx.cuda.mbarrier_wait", + "tirx.cuda.printf", + "tirx.cuda.nano_sleep", + "tirx.nvshmem.fence", + "tirx.nvshmem.quiet", + "tirx.nvshmem.barrier_all", + "tirx.nvshmem.getmem_nbi", + "tirx.nvshmem.getmem_nbi_warp", + "tirx.nvshmem.putmem_nbi", + "tirx.nvshmem.putmem_nbi_warp", + "tirx.nvshmem.putmem_nbi_block", + "tirx.nvshmem.signal_op", + "tirx.nvshmem.wait_until", + "tirx.nvshmem.putmem_signal_nbi", + "tirx.nvshmem.putmem_signal_nbi_warp", + "tirx.nvshmem.putmem_signal_nbi_block", + }) { + OpDef(name).set_attr<FInferType>("FInferType", FInferType::FromNative<&InferTypeVoid>()); } - OpDef("tirx.cuda.tcgen05_encode_instr_descriptor") - .set_attr<FInferType>("FInferType", - FInferType::FromNative<&InferTypeCudaInstructionDescriptor<false>>()); - OpDef("tirx.cuda.tcgen05_encode_instr_descriptor_block_scaled") - .set_attr<FInferType>("FInferType", - FInferType::FromNative<&InferTypeCudaInstructionDescriptor<true>>()); - OpDef("tirx.cuda.wgmma_noop_barrier") - .set_attr<FInferType>("FInferType", FInferType::FromNative<&InferTypeCudaWgmmaNoopBarrier>()); - - for (const char* name : {"tirx.cuda.cluster_sync", "tirx.cuda.cta_sync", "tirx.cuda.grid_sync", - "tirx.cuda.warp_sync", "tirx.cuda.thread_fence"}) { - OpDef(name).set_attr<FInferType>("FInferType", FInferType::FromNative<&InferTypeCudaSync>()); + for (const char* name : {"tirx.cuda.syncthreads_and", "tirx.cuda.syncthreads_or", + "tirx.nvshmem.my_pe", "tirx.nvshmem.n_pes"}) { + OpDef(name).set_attr<FInferType>("FInferType", FInferType::FromNative<&InferTypeInt32>()); } OpDef("tirx.cuda.elect_sync") - .set_attr<FInferType>("FInferType", FInferType::FromNative<&InferTypeCudaElectSync>()); - OpDef("tirx.cuda.mbarrier_wait") - .set_attr<FInferType>("FInferType", FInferType::FromNative<&InferTypeCudaMBarrierWait>()); - OpDef("tirx.nvshmem.fence") - .set_attr<FInferType>("FInferType", FInferType::FromNative<&InferTypeNvshmemFence>()); - OpDef("tirx.cuda.printf") - .set_attr<FInferType>("FInferType", FInferType::FromNative<&InferTypeCudaPrintf>()); - for (const char* name : {"tirx.cuda.syncthreads_and", "tirx.cuda.syncthreads_or"}) { - OpDef(name).set_attr<FInferType>("FInferType", - FInferType::FromNative<&InferTypeCudaSyncthreadsPredicate>()); - } - OpDef("tirx.cuda.nano_sleep") - .set_attr<FInferType>("FInferType", FInferType::FromNative<&InferTypeCudaNanoSleep>()); + .set_attr<FInferType>("FInferType", FInferType::FromNative<&InferTypeUInt32>()); OpDef("tirx.cuda.ldg") .set_attr<FInferType>("FInferType", FInferType::FromNative<&InferTypeCudaLdg>()); OpDef("tirx.s_tir.ldg32") @@ -735,27 +556,6 @@ void RegisterDeviceIntrinsicAliases() { .set_attr<FInferType>("FInferType", FInferType::FromNative<&InferTypeCudaAtomicAdd>()); OpDef("tirx.cuda.atomic_cas") .set_attr<FInferType>("FInferType", FInferType::FromNative<&InferTypeCudaAtomicCas>()); - for (const char* name : {"tirx.nvshmem.my_pe", "tirx.nvshmem.n_pes"}) { - OpDef(name).set_attr<FInferType>("FInferType", - FInferType::FromNative<&InferTypeNvshmemQuery>()); - } - for (const char* name : {"tirx.nvshmem.quiet", "tirx.nvshmem.barrier_all"}) { - OpDef(name).set_attr<FInferType>("FInferType", FInferType::FromNative<&InferTypeNvshmemSync>()); - } - for (const char* name : - {"tirx.nvshmem.getmem_nbi", "tirx.nvshmem.getmem_nbi_warp", "tirx.nvshmem.putmem_nbi", - "tirx.nvshmem.putmem_nbi_warp", "tirx.nvshmem.putmem_nbi_block"}) { - OpDef(name).set_attr<FInferType>("FInferType", FInferType::FromNative<&InferTypeNvshmemRma>()); - } - OpDef("tirx.nvshmem.signal_op") - .set_attr<FInferType>("FInferType", FInferType::FromNative<&InferTypeNvshmemSignal>()); - OpDef("tirx.nvshmem.wait_until") - .set_attr<FInferType>("FInferType", FInferType::FromNative<&InferTypeNvshmemWaitUntil>()); - for (const char* name : {"tirx.nvshmem.putmem_signal_nbi", "tirx.nvshmem.putmem_signal_nbi_warp", - "tirx.nvshmem.putmem_signal_nbi_block"}) { - OpDef(name).set_attr<FInferType>("FInferType", - FInferType::FromNative<&InferTypeNvshmemPutmemSignal>()); - } } } // namespace
