gemini-code-assist[bot] commented on code in PR #19838:
URL: https://github.com/apache/tvm/pull/19838#discussion_r3440318611
##########
src/backend/cuda/op/target_builtin.cc:
##########
@@ -333,76 +115,27 @@ TIRX_DEFINE_BUILTIN_FUNC(timer_end_cuda)
TIRX_DEFINE_BUILTIN_FUNC(timer_finalize_cuda)
.set_attr<TCallEffectKind>("TCallEffectKind",
static_cast<int64_t>(CallEffectKind::kOpaque));
-TIRX_DEFINE_BUILTIN_FUNC(nvshmem_my_pe)
- .set_attr<TCallEffectKind>("TCallEffectKind",
static_cast<int64_t>(CallEffectKind::kOpaque));
-
-TIRX_DEFINE_BUILTIN_FUNC(nvshmem_n_pes)
- .set_attr<TCallEffectKind>("TCallEffectKind",
static_cast<int64_t>(CallEffectKind::kOpaque));
-
-TIRX_DEFINE_BUILTIN_FUNC(nvshmem_getmem_nbi)
- .set_attr<TCallEffectKind>("TCallEffectKind",
static_cast<int64_t>(CallEffectKind::kOpaque));
-
-TIRX_DEFINE_BUILTIN_FUNC(nvshmem_putmem_nbi)
- .set_attr<TCallEffectKind>("TCallEffectKind",
static_cast<int64_t>(CallEffectKind::kOpaque));
-
-TIRX_DEFINE_BUILTIN_FUNC(nvshmem_getmem_nbi_warp)
- .set_attr<TCallEffectKind>("TCallEffectKind",
static_cast<int64_t>(CallEffectKind::kOpaque));
-
-TIRX_DEFINE_BUILTIN_FUNC(nvshmem_putmem_nbi_warp)
- .set_attr<TCallEffectKind>("TCallEffectKind",
static_cast<int64_t>(CallEffectKind::kOpaque));
-
-TIRX_DEFINE_BUILTIN_FUNC(nvshmem_getmem_nbi_block)
- .set_attr<TCallEffectKind>("TCallEffectKind",
static_cast<int64_t>(CallEffectKind::kOpaque));
-
-TIRX_DEFINE_BUILTIN_FUNC(nvshmem_putmem_nbi_block)
- .set_attr<TCallEffectKind>("TCallEffectKind",
static_cast<int64_t>(CallEffectKind::kOpaque));
-
-TIRX_DEFINE_BUILTIN_FUNC(nvshmem_signal_op)
- .set_attr<TCallEffectKind>("TCallEffectKind",
static_cast<int64_t>(CallEffectKind::kOpaque));
-
-TIRX_DEFINE_BUILTIN_FUNC(nvshmem_wait_until)
- .set_attr<TCallEffectKind>("TCallEffectKind",
static_cast<int64_t>(CallEffectKind::kOpaque));
-
-TIRX_DEFINE_BUILTIN_FUNC(nvshmem_quiet)
- .set_attr<TCallEffectKind>("TCallEffectKind",
static_cast<int64_t>(CallEffectKind::kOpaque));
-
-TIRX_DEFINE_BUILTIN_FUNC(nvshmem_putmem_signal_nbi)
- .set_attr<TCallEffectKind>("TCallEffectKind",
static_cast<int64_t>(CallEffectKind::kOpaque));
-
-TIRX_DEFINE_BUILTIN_FUNC(nvshmem_putmem_signal_nbi_warp)
- .set_attr<TCallEffectKind>("TCallEffectKind",
static_cast<int64_t>(CallEffectKind::kOpaque));
-
-TIRX_DEFINE_BUILTIN_FUNC(nvshmem_putmem_signal_nbi_block)
- .set_attr<TCallEffectKind>("TCallEffectKind",
static_cast<int64_t>(CallEffectKind::kOpaque));
-
-TIRX_DEFINE_BUILTIN_FUNC(nvshmem_fence)
- .set_attr<TCallEffectKind>("TCallEffectKind",
static_cast<int64_t>(CallEffectKind::kOpaque));
-
-TIRX_DEFINE_BUILTIN_FUNC(nvshmem_barrier_all)
- .set_attr<TCallEffectKind>("TCallEffectKind",
static_cast<int64_t>(CallEffectKind::kOpaque));
-
RegisterDeviceIntrinsicAliases();
// clang-format on
}
namespace {
struct DeviceIntrinsicRegistration {
- const char* flat_name;
+ const char* name;
Review Comment:
Phasing out the flat macro registrations has caused a regression where the
`TScriptDtypePrintLocation` attribute is lost for several key PTX intrinsics
(such as `ptx.mma`, `ptx.mma_sp`, `ptx.ldmatrix`, `ptx.cp_async`, etc.). These
intrinsics previously had this attribute set to
`ScriptDtypePrintLocation::kFirst` to ensure correct TVMScript printing and
parsing.\n\nTo fix this, we should add a `dtype_print_location` field to
`DeviceIntrinsicRegistration` (defaulted to `0` or
`ScriptDtypePrintLocation::kNone` to avoid breaking existing
registrations):\n\n```cpp\nstruct DeviceIntrinsicRegistration {\n const char*
name;\n const char* namespace_name;\n CallEffectKind effect_kind;\n int
dtype_print_location = 0; // Default to 0 (no print location attribute)\n};\n```
##########
src/backend/cuda/op/target_builtin.cc:
##########
@@ -419,12 +152,11 @@ void RegisterDeviceIntrinsic(const
DeviceIntrinsicRegistration& reg) {
.set_attr<TScriptPrinterName>("TScriptPrinterName", printer_name,
/*plevel=*/15);
};
- register_one(flat_op_name);
register_one(canonical_op_name);
Review Comment:
To support restoring the lost `TScriptDtypePrintLocation` attribute for PTX
intrinsics, we should update the `register_one` lambda inside
`RegisterDeviceIntrinsic` to set the attribute when `reg.dtype_print_location`
is non-zero:\n\n```cpp\n auto register_one = [&](const std::string& op_name)
{\n auto entry = OpRegEntry::RegisterOrGet(op_name)\n .set_name()\n
.set_num_inputs(-1)\n .set_attr<TCallEffectKind>("TCallEffectKind",
static_cast<int64_t>(reg.effect_kind))\n
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("device_intrin"),
/*plevel=*/15)\n
.set_attr<TDeviceIntrinsicNamespace>("TDeviceIntrinsicNamespace",
namespace_attr, /*plevel=*/15)\n
.set_attr<TScriptPrinterName>("TScriptPrinterName", printer_name,
/*plevel=*/15);\n if (reg.dtype_print_location != 0) {\n
entry.set_attr<TScriptDtypePrintLocation>(\n
"TScriptDtypePrintLocation", static_cast<int64_t>(reg.dtype_print_location),
/*plevel=*/15);\n }\n
};\n```
##########
src/tirx/transform/tile_primitive_dispatch.cc:
##########
@@ -125,7 +125,7 @@ class ElectSyncFinder : public StmtExprVisitor {
void VisitExpr_(const CallNode* op) final {
auto is_canonical_elect_sync = [&]() {
- static const Op& ptx_elect_sync_op = Op::Get("tirx.ptx_elect_sync");
+ static const Op& ptx_elect_sync_op = Op::Get("tirx.ptx.elect_sync");
Review Comment:
The dual-name matching logic for `ptx_elect_sync` can be simplified now that
the flat alias has been removed. Since `ptx_elect_sync_op` is resolved directly
to the canonical `tirx.ptx.elect_sync` operator, the subsequent string
comparison check is redundant and can be removed, similar to the cleanup done
in `IsPtxElectSyncCall` in `src/tirx/analysis/filter_canonical.cc`.\n\nConsider
simplifying the lambda to:\n```cpp\n auto is_canonical_elect_sync = [&]()
{\n static const Op& ptx_elect_sync_op = Op::Get("tirx.ptx.elect_sync");\n
return op->op.same_as(ptx_elect_sync_op);\n };\n```
##########
src/backend/cuda/op/target_builtin.cc:
##########
@@ -419,12 +152,11 @@ void RegisterDeviceIntrinsic(const
DeviceIntrinsicRegistration& reg) {
.set_attr<TScriptPrinterName>("TScriptPrinterName", printer_name,
/*plevel=*/15);
};
- register_one(flat_op_name);
register_one(canonical_op_name);
}
#define TIRX_DEVICE_INTRIN_ALIAS(OpName, Namespace, EffectKind) \
- {#OpName, #Namespace, CallEffectKind::EffectKind}
+ { #OpName, #Namespace, CallEffectKind::EffectKind }
Review Comment:
To allow registering intrinsics with a specific `TScriptDtypePrintLocation`
(such as `ptx.mma`, `ptx.ldmatrix`, etc.), we can define a companion macro
`TIRX_DEVICE_INTRIN_ALIAS_WITH_DTYPE`:\n\n```cpp\n#define
TIRX_DEVICE_INTRIN_ALIAS_WITH_DTYPE(OpName, Namespace, EffectKind, DtypeLoc) {
#OpName, #Namespace, CallEffectKind::EffectKind,
static_cast<int>(ScriptDtypePrintLocation::DtypeLoc) }\n```\n\nThis can then be
used in `kDeviceIntrinsics` for the affected intrinsics to restore their
TVMScript printing behavior.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
