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 c44e384d6ad90c7e84d3f75ad3b4a79702f98d35 Author: Tianqi Chen <[email protected]> AuthorDate: Wed Sep 30 21:02:57 2026 +0000 [FIX][CUDA] Preserve intrinsic operand checks in validators --- src/backend/cuda/op/target_builtin.cc | 170 ++++++++++++++++++++++++++++++++++ 1 file changed, 170 insertions(+) diff --git a/src/backend/cuda/op/target_builtin.cc b/src/backend/cuda/op/target_builtin.cc index 71cfe281be..c0931f173b 100644 --- a/src/backend/cuda/op/target_builtin.cc +++ b/src/backend/cuda/op/target_builtin.cc @@ -71,6 +71,138 @@ 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; +} + +void CheckNvshmemRmaArguments(const CallNode* call) { + CheckPointerArgument(call->args[0]); + CheckPointerArgument(call->args[1]); + CheckScalarInteger(call->args[2]); +} + +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]); +} + +void ValidateCudaPrintf(const CallNode* call) { + 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"; + } +} + +void ValidateCudaSyncthreadsPredicate(const CallNode* call) { + 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"; +} + +void ValidateCudaMatrixDescriptor(const CallNode* call) { + CheckPointerArgument(call->args[0]); + CheckPointerArgument(call->args[1]); + for (size_t i = 2; i < 5; ++i) CheckScalarInteger(call->args[i]); +} + +template <bool block_scaled> +void ValidateCudaInstructionDescriptor(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; + 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"; + } +} + +void ValidateCudaWgmmaNoopBarrier(const CallNode* call) { + 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"; +} + +void ValidateCudaNanoSleep(const CallNode* call) { CheckScalarInteger(call->args[0]); } + +void ValidateNvshmemRma(const CallNode* call) { + CheckNvshmemRmaArguments(call); + CheckScalarInteger(call->args[3]); +} + +void ValidateNvshmemSignal(const CallNode* call) { CheckNvshmemSignalArguments(call, 0); } + +void ValidateNvshmemWaitUntil(const CallNode* call) { + 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"}); +} + +void ValidateNvshmemPutmemSignal(const CallNode* call) { + CheckNvshmemRmaArguments(call); + CheckNvshmemSignalArguments(call, 3); +} + +void ValidateCudaMBarrierWait(const CallNode* call) { + 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"; +} + Type InferTypeCudaLdg(const CallNode* call) { TVM_FFI_CHECK_EQ(call->args.size(), 2U, ValueError) << "Scalar cuda.ldg expects two arguments"; CheckPointerArgument(call->args[0]); @@ -356,6 +488,17 @@ void RegisterDeviceIntrinsic(const char* op_name, const char* op_namespace, RegisterDeviceIntrinsicAttrs(def, op_namespace, effect_kind, names.printer); } +void AddDeviceIntrinsicValidator(const char* name, void (*validate)(const CallNode*)) { + OpDef def(name); + auto signature = + def.op()->validator.cast<ffi::reflection::NativeFunction<void(const CallNode*)>>(); + ffi::TypedFunction<void(const CallNode*)> validator([signature, validate](const CallNode* call) { + signature(call); + validate(call); + }); + def.set_validator(validator, true); +} + void RegisterDeviceIntrinsicAliases() { RegisterDeviceIntrinsic("cuda_any_sync", "cuda", CallEffectKind::kPure, sig::arg("mask"), sig::arg("pred")); @@ -512,6 +655,33 @@ 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"}) { + AddDeviceIntrinsicValidator(name, ValidateCudaMatrixDescriptor); + } + AddDeviceIntrinsicValidator("tirx.cuda.tcgen05_encode_instr_descriptor", + ValidateCudaInstructionDescriptor<false>); + AddDeviceIntrinsicValidator("tirx.cuda.tcgen05_encode_instr_descriptor_block_scaled", + ValidateCudaInstructionDescriptor<true>); + AddDeviceIntrinsicValidator("tirx.cuda.wgmma_noop_barrier", ValidateCudaWgmmaNoopBarrier); + AddDeviceIntrinsicValidator("tirx.cuda.printf", ValidateCudaPrintf); + for (const char* name : {"tirx.cuda.syncthreads_and", "tirx.cuda.syncthreads_or"}) { + AddDeviceIntrinsicValidator(name, ValidateCudaSyncthreadsPredicate); + } + AddDeviceIntrinsicValidator("tirx.cuda.nano_sleep", ValidateCudaNanoSleep); + AddDeviceIntrinsicValidator("tirx.cuda.mbarrier_wait", ValidateCudaMBarrierWait); + 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"}) { + AddDeviceIntrinsicValidator(name, ValidateNvshmemRma); + } + AddDeviceIntrinsicValidator("tirx.nvshmem.signal_op", ValidateNvshmemSignal); + AddDeviceIntrinsicValidator("tirx.nvshmem.wait_until", ValidateNvshmemWaitUntil); + for (const char* name : {"tirx.nvshmem.putmem_signal_nbi", "tirx.nvshmem.putmem_signal_nbi_warp", + "tirx.nvshmem.putmem_signal_nbi_block"}) { + AddDeviceIntrinsicValidator(name, ValidateNvshmemPutmemSignal); + } + for (const char* name : { "tirx.cuda.tcgen05_encode_matrix_descriptor", "tirx.cuda.wgmma_encode_matrix_descriptor",
