src/tirx/transform/bind_target.cc - tvm - Git at Google

 /*
  * Licensed to the Apache Software Foundation (ASF) under one
  * or more contributor license agreements.  See the NOTICE file
  * distributed with this work for additional information
  * regarding copyright ownership.  The ASF licenses this file
  * to you under the Apache License, Version 2.0 (the
  * "License"); you may not use this file except in compliance
  * with the License.  You may obtain a copy of the License at
  *
  *   http://www.apache.org/licenses/LICENSE-2.0
  *
  * Unless required by applicable law or agreed to in writing,
  * software distributed under the License is distributed on an
  * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
  * KIND, either express or implied.  See the License for the
  * specific language governing permissions and limitations
  * under the License.
  */

 /*!
  * \file bind_target.cc
  * \brief Pass to bind target to primfunc for heterogeneous compilation.
  *
  * This pass analyzes function call patterns in an IRModule and binds appropriate
  * targets (host/device) to each PrimFunc based on where they are called from.
  *
  * The pass handles the following scenarios:
  * 1. Functions called from host code (CPU)
  * 2. Functions called from device code (GPU/accelerator)
  * 3. Functions called from both host and device
  * 4. Externally exposed functions (entry points)
  *
  * For functions called from both host and device, the pass creates duplicates
  * with appropriate targets and updates call sites accordingly.
  */

 #include <tvm/ffi/cast.h>
 #include <tvm/ffi/reflection/registry.h>
 #include <tvm/ir/unique_name_supply.h>
 #include <tvm/s_tir/stmt.h>
 #include <tvm/s_tir/transform.h>
 #include <tvm/tirx/stmt_functor.h>
 #include <tvm/tirx/transform.h>

 #include <unordered_set>

 #include "tvm/ir/attrs.h"

 namespace tvm {
 namespace tirx {

 /*!
  * \brief Visitor class to classify function calls as host or device calls.
  *
  * This visitor traverses the IRModule to identify which functions are called
  * from host code vs device code. It tracks GPU scopes (thread binding loops
  * and thread extent attributes) to determine the calling context.
  */
 class FunctionClassifierVisitor : public StmtExprVisitor {
  public:
   /*!
    * \brief Analyze function call patterns in the IRModule.
    * \param mod The IRModule to analyze
    * \return A tuple containing:
    *         - Set of GlobalVarNodes called from host code
    *         - Set of GlobalVarNodes called from device code
    * \note A single function can be called by both host and device contexts.
    */
   static std::tuple<std::unordered_set<const GlobalVarNode*>,
                     std::unordered_set<const GlobalVarNode*>>
   GetFunctionCallers(const IRModule& mod) {
     FunctionClassifierVisitor visitor;

     // Only analyze externally exposed functions as potential callers
     // since they represent the entry points where host/device calls originate
     for (const auto& [gvar, func] : mod->functions) {
       bool is_externally_exposed = func->GetAttr<ffi::String>(tvm::attr::kGlobalSymbol).has_value();
       const auto* prim_func = func.as<PrimFuncNode>();

       if (is_externally_exposed && prim_func != nullptr) {
         visitor.VisitStmt(prim_func->body);
       }
     }

     return std::make_tuple(visitor.host_called_global_vars_, visitor.device_called_global_vars_);
   }

  private:
   using StmtExprVisitor::VisitStmt_;

   void VisitExpr_(const CallNode* op) final {
     const auto* global_var = op->op.as<GlobalVarNode>();
     if (global_var != nullptr) {
       // Classify the call based on current scope
       if (is_under_gpu_scope_) {
         device_called_global_vars_.insert(global_var);
       } else {
         host_called_global_vars_.insert(global_var);
       }
     }
     StmtExprVisitor::VisitExpr_(op);
   }

   void VisitStmt_(const ForNode* op) final {
     if (op->kind == ForKind::kThreadBinding) {
       // Enter GPU scope for thread binding loops
       bool last_is_under_gpu_scope = is_under_gpu_scope_;
       is_under_gpu_scope_ = true;
       StmtExprVisitor::VisitStmt_(op);
       is_under_gpu_scope_ = last_is_under_gpu_scope;
     } else {
       StmtExprVisitor::VisitStmt_(op);
     }
   }

   void VisitStmt_(const AttrStmtNode* op) final {
     if (op->attr_key == attr::thread_extent || op->attr_key == s_tir::attr::virtual_thread ||
         op->attr_key == attr::kDeviceEntry) {
       // Enter GPU scope for thread extent and virtual thread attributes
       bool last_is_under_gpu_scope = is_under_gpu_scope_;
       is_under_gpu_scope_ = true;
       StmtExprVisitor::VisitStmt_(op);
       is_under_gpu_scope_ = last_is_under_gpu_scope;
     } else {
       StmtExprVisitor::VisitStmt_(op);
     }
   }

  private:
   /*! \brief Whether the current statement is under a GPU scope */
   bool is_under_gpu_scope_ = false;
   /*! \brief Set of functions called from host code */
   std::unordered_set<const GlobalVarNode*> host_called_global_vars_;
   /*! \brief Set of functions called from device code */
   std::unordered_set<const GlobalVarNode*> device_called_global_vars_;
 };

 /*!
  * \brief Mutator class to substitute function calls in host contexts.
  *
  * This mutator replaces calls to functions that have been duplicated for
  * host/device contexts. It only performs substitutions when not under
  * GPU scope to ensure device calls remain unchanged.
  */
 class CallSubstitutor : public StmtExprMutator {
  public:
   /*!
    * \brief Constructor with function replacement mapping.
    * \param replacements Map from original GlobalVar to host-specific GlobalVar
    */
   explicit CallSubstitutor(const ffi::Map<GlobalVar, GlobalVar>& replacements)
       : replacements_(replacements) {}

   /*!
    * \brief Substitute function calls in a PrimFunc.
    * \param func The PrimFunc to process
    * \return The modified PrimFunc with updated calls
    */
   PrimFunc Substitute(PrimFunc func) {
     auto f = func.CopyOnWrite();
     auto body = VisitStmt(f->body);

     // Only update if the body actually changed
     if (body.same_as(func->body)) {
       return func;
     }

     f->body = std::move(body);
     return func;
   }

  private:
   using StmtExprMutator::VisitStmt_;

   Expr VisitExpr_(const CallNode* op) final {
     auto call = StmtExprMutator::VisitExpr_(op).as_or_throw<Call>();

     // Only substitute calls when not under GPU scope
     if (!is_under_gpu_scope_) {
       if (auto old_gvar = call->op.as<GlobalVar>()) {
         if (auto new_gvar = replacements_.Get(old_gvar.value())) {
           call.CopyOnWrite()->op = new_gvar.value();
         }
       }
     }
     return call;
   }

   Stmt VisitStmt_(const ForNode* op) final {
     if (op->kind == ForKind::kThreadBinding) {
       // Enter GPU scope for thread binding loops
       bool last_is_under_gpu_scope = is_under_gpu_scope_;
       is_under_gpu_scope_ = true;
       auto stmt = StmtExprMutator::VisitStmt_(op);
       is_under_gpu_scope_ = last_is_under_gpu_scope;
       return stmt;
     } else {
       return StmtExprMutator::VisitStmt_(op);
     }
   }

   Stmt VisitStmt_(const AttrStmtNode* op) final {
     if (op->attr_key == attr::thread_extent || op->attr_key == s_tir::attr::virtual_thread ||
         op->attr_key == attr::kDeviceEntry) {
       // Enter GPU scope for thread extent and virtual thread attributes
       bool last_is_under_gpu_scope = is_under_gpu_scope_;
       is_under_gpu_scope_ = true;
       auto stmt = StmtExprMutator::VisitStmt_(op);
       is_under_gpu_scope_ = last_is_under_gpu_scope;
       return stmt;
     } else {
       return StmtExprMutator::VisitStmt_(op);
     }
   }

  private:
   /*! \brief Whether the current statement is under a GPU scope */
   bool is_under_gpu_scope_ = false;
   /*! \brief Mapping from original functions to host-specific duplicates */
   ffi::Map<GlobalVar, GlobalVar> replacements_;
 };

 /*!
  * \brief Bind appropriate targets to functions in an IRModule.
  *
  * This function analyzes the call patterns in the module and binds appropriate
  * targets to each PrimFunc based on where they are called from. The binding
  * follows these rules:
  *
  * 1. Externally exposed functions (with global symbol) get the full target
  * 2. Functions called only from host get the host target
  * 3. Functions called only from device get the device target
  * 4. Functions called from both contexts get the device target, and a duplicate
  *    is created with the host target for host callers
  *
  * \param mod The IRModule to process
  * \param target The target to bind (should include both host and device)
  * \return The modified IRModule with targets bound to functions
  */
 IRModule BindTarget(IRModule mod, const Target& target) {
   // Extract host and device targets
   auto target_host = target->host.value_or(Target("llvm")).as_or_throw<Target>();
   auto target_without_host = target.WithoutHost();

   auto mod_copy_on_write = mod.CopyOnWrite();
   auto new_mod = ffi::GetRef<IRModule>(mod_copy_on_write);

   // Step 1: Analyze function call patterns
   auto [host_called_global_vars, device_called_global_vars] =
       FunctionClassifierVisitor::GetFunctionCallers(mod);

   // Step 2: Bind target to functions with the following rules:
   //  1. If the function has a target, and the target has a host, and the function does not have a
   //     host, then add the host to the function target
   //  2. If the function is marked as host function, bind the host target to the function
   //  3. If the function is externally exposed (with global symbol), bind the full target
   //  4. If the function is not externally exposed:
   //    2.1 If the function is called by both host and device, bind the device target to the current
   //        function and duplicate the function with the host target.
   //    2.2 If the function is called by host only, bind the host target to the current function
   //    2.3 If the function is called by device only, bind the device target to the current function
   //    2.4 If the function is not called by any host or device, skip binding

   // Track duplicated functions for call replacement
   ffi::Map<GlobalVar, GlobalVar> host_function_replacements;
   UniqueNameSupply global_names(new_mod->functions.begin(), new_mod->functions.end(),
                                 [](const auto& kv) { return kv.first->name_hint; });

   for (auto [gvar, func] : mod->functions) {
     const auto* prim_func_node = func.as<PrimFuncNode>();
     if (prim_func_node == nullptr) {
       // Skip non-PrimFunc entries
       continue;
     }
     auto prim_func = ffi::GetRef<PrimFunc>(prim_func_node);

     bool is_externally_exposed =
         prim_func->GetAttr<ffi::String>(tvm::attr::kGlobalSymbol).has_value();

     if (auto func_target = func->GetAttr<Target>(tvm::attr::kTarget)) {
       // Rule 1: If the function has a target, and the target has a host, and the function does not
       // have a host, then add the host to the function target
       auto func_target_host = func_target.value()->GetHost();
       auto target_host = target->GetHost();

       if (target_host && !func_target_host && is_externally_exposed) {
         auto new_target = Target::WithHost(func_target.value(), target_host.value());
         new_mod->Update(gvar, WithAttr(std::move(prim_func), tvm::attr::kTarget, new_target));
       }
       continue;
     }

     if (prim_func->HasNonzeroAttr(tvm::tirx::attr::kIsHostFunc)) {
       // Rule 2: If the function is marked as host function, bind the host target to the function
       prim_func = WithAttr(std::move(prim_func), tvm::attr::kTarget,
                            Target::WithHost(target_host, target_host));
       new_mod->Update(gvar, WithoutAttr(std::move(prim_func), tvm::tirx::attr::kIsHostFunc));
       continue;
     }

     if (is_externally_exposed) {
       // Rule 3: Externally exposed functions get the full target
       new_mod->Update(gvar, WithAttr(std::move(prim_func), tvm::attr::kTarget, target));
     } else {
       const auto* gvar_node = gvar.get();
       bool called_by_host = host_called_global_vars.count(gvar_node);
       bool called_by_device = device_called_global_vars.count(gvar_node);

       if (called_by_host && called_by_device) {
         // Rule 4.1: Called by both host and device
         // Bind device target to current function
         PrimFunc host_func = s_tir::RenewDefs(prim_func);
         new_mod->Update(gvar,
                         WithAttr(std::move(prim_func), tvm::attr::kTarget, target_without_host));

         // Create duplicate with host target for host callers
         host_func = WithAttr(std::move(host_func), tvm::attr::kTarget, target_host);
         ffi::String host_func_name = gvar->name_hint + "_host";
         GlobalVar host_gvar = GlobalVar(global_names->FreshName(host_func_name, false));

         new_mod->Add(host_gvar, host_func);
         host_function_replacements.Set(gvar, host_gvar);

       } else if (called_by_host) {
         // Rule 4.2: Called by host only
         new_mod->Update(gvar, WithAttr(std::move(prim_func), tvm::attr::kTarget, target_host));
       } else if (called_by_device) {
         // Rule 4.3: Called by device only
         new_mod->Update(gvar,
                         WithAttr(std::move(prim_func), tvm::attr::kTarget, target_without_host));
       } else {
         // Rule 4.4: Not called by any context
         // NOTE: To keep the current behavior, we bind the target to the full target, but it needs
         // further check
         new_mod->Update(gvar,
                         WithAttr(std::move(prim_func), tvm::attr::kTarget, target_without_host));
       }
     }
   }

   // Step 3: Update call sites in externally exposed functions
   if (!host_function_replacements.empty()) {
     CallSubstitutor substitutor(host_function_replacements);

     for (auto [gvar, func] : mod->functions) {
       const auto* prim_func = func.as<PrimFuncNode>();
       if (prim_func == nullptr) {
         continue;
       }

       bool is_externally_exposed =
           prim_func->GetAttr<ffi::String>(tvm::attr::kGlobalSymbol).has_value();
       if (is_externally_exposed) {
         // Update calls in externally exposed functions to use host duplicates
         PrimFunc new_func = substitutor.Substitute(func.as_or_throw<PrimFunc>());
         new_mod->Update(gvar, new_func);
       }
     }
   }

   return new_mod;
 }

 namespace transform {

 /*!
  * \brief Create a pass that binds targets to functions in an IRModule.
  *
  * This pass analyzes the call patterns in the module and binds appropriate
  * targets (host/device) to each PrimFunc based on where they are called from.
  *
  * \param target The target to bind (should include both host and device)
  * \return A transform pass that performs target binding
  */
 transform::Pass BindTarget(Target target) {
   auto fpass = [target](IRModule mod, transform::PassContext ctx) {
     return tvm::tirx::BindTarget(mod, target);
   };
   return tirx::transform::CreateModulePass(fpass, 0, "tirx.BindTarget", {});
 }

 TVM_FFI_STATIC_INIT_BLOCK() {
   namespace refl = tvm::ffi::reflection;
   refl::GlobalDef().def("tirx.transform.BindTarget", BindTarget);
 }

 }  // namespace transform
 }  // namespace tirx
 }  // namespace tvm
	/*
	* Licensed to the Apache Software Foundation (ASF) under one
	* or more contributor license agreements. See the NOTICE file
	* distributed with this work for additional information
	* regarding copyright ownership. The ASF licenses this file
	* to you under the Apache License, Version 2.0 (the
	* "License"); you may not use this file except in compliance
	* with the License. You may obtain a copy of the License at
	*
	* http://www.apache.org/licenses/LICENSE-2.0
	*
	* Unless required by applicable law or agreed to in writing,
	* software distributed under the License is distributed on an
	* "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
	* KIND, either express or implied. See the License for the
	* specific language governing permissions and limitations
	* under the License.
	*/

	/*!
	* \file bind_target.cc
	* \brief Pass to bind target to primfunc for heterogeneous compilation.
	*
	* This pass analyzes function call patterns in an IRModule and binds appropriate
	* targets (host/device) to each PrimFunc based on where they are called from.
	*
	* The pass handles the following scenarios:
	* 1. Functions called from host code (CPU)
	* 2. Functions called from device code (GPU/accelerator)
	* 3. Functions called from both host and device
	* 4. Externally exposed functions (entry points)
	*
	* For functions called from both host and device, the pass creates duplicates
	* with appropriate targets and updates call sites accordingly.
	*/

	#include <tvm/ffi/cast.h>
	#include <tvm/ffi/reflection/registry.h>
	#include <tvm/ir/unique_name_supply.h>
	#include <tvm/s_tir/stmt.h>
	#include <tvm/s_tir/transform.h>
	#include <tvm/tirx/stmt_functor.h>
	#include <tvm/tirx/transform.h>

	#include <unordered_set>

	#include "tvm/ir/attrs.h"

	namespace tvm {
	namespace tirx {

	/*!
	* \brief Visitor class to classify function calls as host or device calls.
	*
	* This visitor traverses the IRModule to identify which functions are called
	* from host code vs device code. It tracks GPU scopes (thread binding loops
	* and thread extent attributes) to determine the calling context.
	*/
	class FunctionClassifierVisitor : public StmtExprVisitor {
	public:
	/*!
	* \brief Analyze function call patterns in the IRModule.
	* \param mod The IRModule to analyze
	* \return A tuple containing:
	* - Set of GlobalVarNodes called from host code
	* - Set of GlobalVarNodes called from device code
	* \note A single function can be called by both host and device contexts.
	*/
	static std::tuple<std::unordered_set<const GlobalVarNode*>,
	std::unordered_set<const GlobalVarNode*>>
	GetFunctionCallers(const IRModule& mod) {
	FunctionClassifierVisitor visitor;

	// Only analyze externally exposed functions as potential callers
	// since they represent the entry points where host/device calls originate
	for (const auto& [gvar, func] : mod->functions) {
	bool is_externally_exposed = func->GetAttr<ffi::String>(tvm::attr::kGlobalSymbol).has_value();
	const auto* prim_func = func.as<PrimFuncNode>();

	if (is_externally_exposed && prim_func != nullptr) {
	visitor.VisitStmt(prim_func->body);
	}
	}

	return std::make_tuple(visitor.host_called_global_vars_, visitor.device_called_global_vars_);
	}

	private:
	using StmtExprVisitor::VisitStmt_;

	void VisitExpr_(const CallNode* op) final {
	const auto* global_var = op->op.as<GlobalVarNode>();
	if (global_var != nullptr) {
	// Classify the call based on current scope
	if (is_under_gpu_scope_) {
	device_called_global_vars_.insert(global_var);
	} else {
	host_called_global_vars_.insert(global_var);
	}
	}
	StmtExprVisitor::VisitExpr_(op);
	}

	void VisitStmt_(const ForNode* op) final {
	if (op->kind == ForKind::kThreadBinding) {
	// Enter GPU scope for thread binding loops
	bool last_is_under_gpu_scope = is_under_gpu_scope_;
	is_under_gpu_scope_ = true;
	StmtExprVisitor::VisitStmt_(op);
	is_under_gpu_scope_ = last_is_under_gpu_scope;
	} else {
	StmtExprVisitor::VisitStmt_(op);
	}
	}

	void VisitStmt_(const AttrStmtNode* op) final {
	if (op->attr_key == attr::thread_extent \|\| op->attr_key == s_tir::attr::virtual_thread \|\|
	op->attr_key == attr::kDeviceEntry) {
	// Enter GPU scope for thread extent and virtual thread attributes
	bool last_is_under_gpu_scope = is_under_gpu_scope_;
	is_under_gpu_scope_ = true;
	StmtExprVisitor::VisitStmt_(op);
	is_under_gpu_scope_ = last_is_under_gpu_scope;
	} else {
	StmtExprVisitor::VisitStmt_(op);
	}
	}

	private:
	/! \brief Whether the current statement is under a GPU scope /
	bool is_under_gpu_scope_ = false;
	/! \brief Set of functions called from host code /
	std::unordered_set<const GlobalVarNode*> host_called_global_vars_;
	/! \brief Set of functions called from device code /
	std::unordered_set<const GlobalVarNode*> device_called_global_vars_;
	};

	/*!
	* \brief Mutator class to substitute function calls in host contexts.
	*
	* This mutator replaces calls to functions that have been duplicated for
	* host/device contexts. It only performs substitutions when not under
	* GPU scope to ensure device calls remain unchanged.
	*/
	class CallSubstitutor : public StmtExprMutator {
	public:
	/*!
	* \brief Constructor with function replacement mapping.
	* \param replacements Map from original GlobalVar to host-specific GlobalVar
	*/
	explicit CallSubstitutor(const ffi::Map<GlobalVar, GlobalVar>& replacements)
	: replacements_(replacements) {}

	/*!
	* \brief Substitute function calls in a PrimFunc.
	* \param func The PrimFunc to process
	* \return The modified PrimFunc with updated calls
	*/
	PrimFunc Substitute(PrimFunc func) {
	auto f = func.CopyOnWrite();
	auto body = VisitStmt(f->body);

	// Only update if the body actually changed
	if (body.same_as(func->body)) {
	return func;
	}

	f->body = std::move(body);
	return func;
	}

	private:
	using StmtExprMutator::VisitStmt_;

	Expr VisitExpr_(const CallNode* op) final {
	auto call = StmtExprMutator::VisitExpr_(op).as_or_throw<Call>();

	// Only substitute calls when not under GPU scope
	if (!is_under_gpu_scope_) {
	if (auto old_gvar = call->op.as<GlobalVar>()) {
	if (auto new_gvar = replacements_.Get(old_gvar.value())) {
	call.CopyOnWrite()->op = new_gvar.value();
	}
	}
	}
	return call;
	}

	Stmt VisitStmt_(const ForNode* op) final {
	if (op->kind == ForKind::kThreadBinding) {
	// Enter GPU scope for thread binding loops
	bool last_is_under_gpu_scope = is_under_gpu_scope_;
	is_under_gpu_scope_ = true;
	auto stmt = StmtExprMutator::VisitStmt_(op);
	is_under_gpu_scope_ = last_is_under_gpu_scope;
	return stmt;
	} else {
	return StmtExprMutator::VisitStmt_(op);
	}
	}

	Stmt VisitStmt_(const AttrStmtNode* op) final {
	if (op->attr_key == attr::thread_extent \|\| op->attr_key == s_tir::attr::virtual_thread \|\|
	op->attr_key == attr::kDeviceEntry) {
	// Enter GPU scope for thread extent and virtual thread attributes
	bool last_is_under_gpu_scope = is_under_gpu_scope_;
	is_under_gpu_scope_ = true;
	auto stmt = StmtExprMutator::VisitStmt_(op);
	is_under_gpu_scope_ = last_is_under_gpu_scope;
	return stmt;
	} else {
	return StmtExprMutator::VisitStmt_(op);
	}
	}

	private:
	/! \brief Whether the current statement is under a GPU scope /
	bool is_under_gpu_scope_ = false;
	/! \brief Mapping from original functions to host-specific duplicates /
	ffi::Map<GlobalVar, GlobalVar> replacements_;
	};

	/*!
	* \brief Bind appropriate targets to functions in an IRModule.
	*
	* This function analyzes the call patterns in the module and binds appropriate
	* targets to each PrimFunc based on where they are called from. The binding
	* follows these rules:
	*
	* 1. Externally exposed functions (with global symbol) get the full target
	* 2. Functions called only from host get the host target
	* 3. Functions called only from device get the device target
	* 4. Functions called from both contexts get the device target, and a duplicate
	* is created with the host target for host callers
	*
	* \param mod The IRModule to process
	* \param target The target to bind (should include both host and device)
	* \return The modified IRModule with targets bound to functions
	*/
	IRModule BindTarget(IRModule mod, const Target& target) {
	// Extract host and device targets
	auto target_host = target->host.value_or(Target("llvm")).as_or_throw<Target>();
	auto target_without_host = target.WithoutHost();

	auto mod_copy_on_write = mod.CopyOnWrite();
	auto new_mod = ffi::GetRef<IRModule>(mod_copy_on_write);

	// Step 1: Analyze function call patterns
	auto [host_called_global_vars, device_called_global_vars] =
	FunctionClassifierVisitor::GetFunctionCallers(mod);

	// Step 2: Bind target to functions with the following rules:
	// 1. If the function has a target, and the target has a host, and the function does not have a
	// host, then add the host to the function target
	// 2. If the function is marked as host function, bind the host target to the function
	// 3. If the function is externally exposed (with global symbol), bind the full target
	// 4. If the function is not externally exposed:
	// 2.1 If the function is called by both host and device, bind the device target to the current
	// function and duplicate the function with the host target.
	// 2.2 If the function is called by host only, bind the host target to the current function
	// 2.3 If the function is called by device only, bind the device target to the current function
	// 2.4 If the function is not called by any host or device, skip binding

	// Track duplicated functions for call replacement
	ffi::Map<GlobalVar, GlobalVar> host_function_replacements;
	UniqueNameSupply global_names(new_mod->functions.begin(), new_mod->functions.end(),
	[](const auto& kv) { return kv.first->name_hint; });

	for (auto [gvar, func] : mod->functions) {
	const auto* prim_func_node = func.as<PrimFuncNode>();
	if (prim_func_node == nullptr) {
	// Skip non-PrimFunc entries
	continue;
	}
	auto prim_func = ffi::GetRef<PrimFunc>(prim_func_node);

	bool is_externally_exposed =
	prim_func->GetAttr<ffi::String>(tvm::attr::kGlobalSymbol).has_value();

	if (auto func_target = func->GetAttr<Target>(tvm::attr::kTarget)) {
	// Rule 1: If the function has a target, and the target has a host, and the function does not
	// have a host, then add the host to the function target
	auto func_target_host = func_target.value()->GetHost();
	auto target_host = target->GetHost();

	if (target_host && !func_target_host && is_externally_exposed) {
	auto new_target = Target::WithHost(func_target.value(), target_host.value());
	new_mod->Update(gvar, WithAttr(std::move(prim_func), tvm::attr::kTarget, new_target));
	}
	continue;
	}

	if (prim_func->HasNonzeroAttr(tvm::tirx::attr::kIsHostFunc)) {
	// Rule 2: If the function is marked as host function, bind the host target to the function
	prim_func = WithAttr(std::move(prim_func), tvm::attr::kTarget,
	Target::WithHost(target_host, target_host));
	new_mod->Update(gvar, WithoutAttr(std::move(prim_func), tvm::tirx::attr::kIsHostFunc));
	continue;
	}

	if (is_externally_exposed) {
	// Rule 3: Externally exposed functions get the full target
	new_mod->Update(gvar, WithAttr(std::move(prim_func), tvm::attr::kTarget, target));
	} else {
	const auto* gvar_node = gvar.get();
	bool called_by_host = host_called_global_vars.count(gvar_node);
	bool called_by_device = device_called_global_vars.count(gvar_node);

	if (called_by_host && called_by_device) {
	// Rule 4.1: Called by both host and device
	// Bind device target to current function
	PrimFunc host_func = s_tir::RenewDefs(prim_func);
	new_mod->Update(gvar,
	WithAttr(std::move(prim_func), tvm::attr::kTarget, target_without_host));

	// Create duplicate with host target for host callers
	host_func = WithAttr(std::move(host_func), tvm::attr::kTarget, target_host);
	ffi::String host_func_name = gvar->name_hint + "_host";
	GlobalVar host_gvar = GlobalVar(global_names->FreshName(host_func_name, false));

	new_mod->Add(host_gvar, host_func);
	host_function_replacements.Set(gvar, host_gvar);

	} else if (called_by_host) {
	// Rule 4.2: Called by host only
	new_mod->Update(gvar, WithAttr(std::move(prim_func), tvm::attr::kTarget, target_host));
	} else if (called_by_device) {
	// Rule 4.3: Called by device only
	new_mod->Update(gvar,
	WithAttr(std::move(prim_func), tvm::attr::kTarget, target_without_host));
	} else {
	// Rule 4.4: Not called by any context
	// NOTE: To keep the current behavior, we bind the target to the full target, but it needs
	// further check
	new_mod->Update(gvar,
	WithAttr(std::move(prim_func), tvm::attr::kTarget, target_without_host));
	}
	}
	}

	// Step 3: Update call sites in externally exposed functions
	if (!host_function_replacements.empty()) {
	CallSubstitutor substitutor(host_function_replacements);

	for (auto [gvar, func] : mod->functions) {
	const auto* prim_func = func.as<PrimFuncNode>();
	if (prim_func == nullptr) {
	continue;
	}

	bool is_externally_exposed =
	prim_func->GetAttr<ffi::String>(tvm::attr::kGlobalSymbol).has_value();
	if (is_externally_exposed) {
	// Update calls in externally exposed functions to use host duplicates
	PrimFunc new_func = substitutor.Substitute(func.as_or_throw<PrimFunc>());
	new_mod->Update(gvar, new_func);
	}
	}
	}

	return new_mod;
	}

	namespace transform {

	/*!
	* \brief Create a pass that binds targets to functions in an IRModule.
	*
	* This pass analyzes the call patterns in the module and binds appropriate
	* targets (host/device) to each PrimFunc based on where they are called from.
	*
	* \param target The target to bind (should include both host and device)
	* \return A transform pass that performs target binding
	*/
	transform::Pass BindTarget(Target target) {
	auto fpass = [target](IRModule mod, transform::PassContext ctx) {
	return tvm::tirx::BindTarget(mod, target);
	};
	return tirx::transform::CreateModulePass(fpass, 0, "tirx.BindTarget", {});
	}

	TVM_FFI_STATIC_INIT_BLOCK() {
	namespace refl = tvm::ffi::reflection;
	refl::GlobalDef().def("tirx.transform.BindTarget", BindTarget);
	}

	} // namespace transform
	} // namespace tirx
	} // namespace tvm