gemini-code-assist[bot] commented on code in PR #19603:
URL: https://github.com/apache/tvm/pull/19603#discussion_r3299728365
##########
src/relax/transform/dead_code_elimination.cc:
##########
@@ -33,19 +33,77 @@
*/
#include <tvm/ffi/reflection/registry.h>
-#include <tvm/ir/analysis.h>
#include <tvm/relax/analysis.h>
#include <tvm/relax/expr.h>
#include <tvm/relax/expr_functor.h>
#include <tvm/relax/transform.h>
+#include <tvm/tirx/expr_functor.h>
+#include <tvm/tirx/function.h>
+#include <tvm/tirx/stmt_functor.h>
+
+#include <unordered_set>
#include "utils.h"
namespace tvm {
namespace relax {
+namespace {
+
+struct RelaxCalleeCollector : relax::ExprVisitor {
+ std::vector<GlobalVar>* callees;
+ explicit RelaxCalleeCollector(std::vector<GlobalVar>* out) : callees(out) {}
+ void VisitExpr_(const GlobalVarNode* node) final {
+ callees->push_back(ffi::GetRef<GlobalVar>(node));
+ }
+};
+
+struct TIRxCalleeCollector : tirx::StmtExprVisitor {
+ std::vector<GlobalVar>* callees;
+ explicit TIRxCalleeCollector(std::vector<GlobalVar>* out) : callees(out) {}
+ void VisitExpr_(const tirx::CallNode* node) final {
+ tirx::StmtExprVisitor::VisitExpr_(node);
+ if (auto opt_gvar = node->op.as<GlobalVar>()) {
+ callees->push_back(opt_gvar.value());
+ }
+ }
+};
+
+// Collect the GlobalVars directly called by `func`. Dedups while
+// preserving first-encounter order (same semantics the old
+// support::OrderedSet path provided).
+ffi::Array<GlobalVar> CollectCallees(const BaseFunc& func) {
+ std::vector<GlobalVar> raw;
+ if (auto opt = func.as<relax::Function>()) {
+ RelaxCalleeCollector visitor(&raw);
+ visitor(opt.value());
+ } else if (func.as<relax::ExternFuncNode>()) {
Review Comment:
For consistency with `relax::Function` and `tirx::PrimFunc` (which are
`ObjectRef` subclasses), we should use `relax::ExternFunc` instead of
`relax::ExternFuncNode`.
```suggestion
} else if (func.as<relax::ExternFunc>()) {
```
--
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]
