gemini-code-assist[bot] commented on code in PR #19663:
URL: https://github.com/apache/tvm/pull/19663#discussion_r3349780819
##########
src/tirx/transform/split_host_device.cc:
##########
@@ -152,6 +197,448 @@ PrimFunc SplitHostDevice(PrimFunc func, IRModule*
device_mod,
return func;
}
+// Device kernel launch lowering
+
+namespace {
+
+struct KernelInfo {
+ // The device on which the PrimFunc runs.
+ Target target;
+
+ // The externally visible symbol which may refer to the PrimFunc
+ // when launching a device kernel.
+ ffi::String global_symbol;
+
+ // The parameters accepted by the PrimFunc. Used to rewrite
+ // `launch_args` to be in terms of the calling scope.
+ ffi::Array<Var> params;
+
+ // The launch parameters that should annotate the PrimFunc, if the
+ // kernel is ever called from the host.
+ ffi::Array<ffi::String> launch_params;
+
+ // Additional arguments which must be provided to the host-side
+ // ffi::Function. These may be in terms of the function's parameters
+ // (e.g. a function that computes the average of `N` elements, and
+ // which must be launched with `N` CUDA threads).
+ ffi::Array<PrimExpr> launch_args;
+};
+
+/*!
+ * \brief Visitor class to collect device-side program information.
+ */
+class DeviceInfoCollector : public StmtVisitor {
+ public:
+ static KernelInfo Collect(const GlobalVar& gvar, const PrimFunc& func) {
+ DeviceInfoCollector collector;
+ collector.info_.target =
func->GetAttr<Target>(tvm::attr::kTarget).value().WithoutHost();
+ collector.info_.params = func->params;
+
+ collector(func->body);
+
+ // The dynamic shared memory is required to be the last of the
+ // kernel launch parameters.
+ if (collector.dyn_shmem_size) {
+ collector.info_.launch_params.push_back(
+ tvm::runtime::launch_param::kUseDynamicSharedMemoryTag);
+ }
+
+ collector.info_.global_symbol =
+
func->GetAttr<ffi::String>(tvm::attr::kGlobalSymbol).value_or(gvar->name_hint);
+
+ collector.info_.launch_args = collector.info_.launch_params.Map(
+ [&](const auto& param) { return collector.GetArgument(param); });
+
+ return collector.info_;
+ }
+
+ private:
+ PrimExpr GetArgument(const ffi::String& launch_param) const {
+ if (launch_param ==
tvm::runtime::launch_param::kUseDynamicSharedMemoryTag) {
+ TVM_FFI_ICHECK(dyn_shmem_size.defined())
+ << "Compute kernel requires launch parameter \"" << launch_param
+ << "\", but PrimFunc did not contain AllocBuffer node with shared
dynamic scope.";
+ return dyn_shmem_size.value();
+ }
+
+ auto extent = thread_extent.Get(launch_param);
+ TVM_FFI_ICHECK(extent) << "Compute kernel requires launch parameter \"" <<
launch_param
+ << "\", but PrimFunc does not contain AttrStmt \""
<< attr::thread_extent
+ << "\" defining this thread extent";
+ return extent.value();
+ }
+
+ void VisitStmt_(const BindNode* op) final {
+ // Track Bind definitions so that thread_extent values and
+ // dyn_shmem_size expressions that reference locally-bound
+ // variables (e.g. CSE variables) can be inlined back to
+ // expressions over function parameters. Substitute earlier
+ // bindings into the value to handle chains (cse_v2 = f(cse_v1)).
+ PrimExpr value = bind_map_.size() ? Substitute(op->value, bind_map_) :
op->value;
+ bind_map_.Set(op->var, value);
+ StmtVisitor::VisitStmt_(op);
+ }
+
+ void VisitStmt_(const AttrStmtNode* op) final {
+ if (op->attr_key == attr::thread_extent) {
+ ffi::String thread_tag;
+ if (auto iv = op->node.as<IterVar>()) {
+ thread_tag = iv.value()->thread_tag;
+ TVM_FFI_ICHECK_NE(thread_tag.length(), 0U);
+ } else if (auto var = op->node.as<Var>()) {
+ thread_tag = var.value()->name_hint;
Review Comment:
In TVM, it is more idiomatic and standard to perform conditional downcasting
using the underlying node types (`IterVarNode` and `VarNode`) rather than the
wrapper classes (`IterVar` and `Var`). This avoids the need for `.value()`
calls and is more robust across different TVM FFI/Object configurations.
```suggestion
if (const auto* iv = op->node.as<IterVarNode>()) {
thread_tag = iv->thread_tag;
TVM_FFI_ICHECK_NE(thread_tag.length(), 0U);
} else if (const auto* var = op->node.as<VarNode>()) {
thread_tag = var->name_hint;
```
--
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]
