gemini-code-assist[bot] commented on code in PR #19598:
URL: https://github.com/apache/tvm/pull/19598#discussion_r3295779090
##########
python/tvm/tirx/functor.py:
##########
@@ -945,6 +1060,55 @@ def visit_string_imm_(self, op: StringImm) -> None:
_ffi_api.PyStmtExprVisitorDefaultVisitExpr(self._outer(), op) # type:
ignore
+def _analyzer_from_module(mod):
+ from tvm.arith.analyzer import Analyzer # pylint:
disable=import-outside-toplevel
+
+ analyzer = Analyzer.__new__(Analyzer)
+ analyzer._const_int_bound = mod("const_int_bound")
+ analyzer._const_int_bound_update = mod("const_int_bound_update")
+ analyzer._const_int_bound_is_bound = mod("const_int_bound_is_bound")
+ analyzer._bind = mod("bind")
+ analyzer._modular_set = mod("modular_set")
+ analyzer._simplify = mod("Simplify")
+ analyzer._rewrite_simplify = mod("rewrite_simplify")
+ analyzer._get_rewrite_simplify_stats = mod("get_rewrite_simplify_stats")
+ analyzer._reset_rewrite_simplify_stats =
mod("reset_rewrite_simplify_stats")
+ analyzer._canonical_simplify = mod("canonical_simplify")
+ analyzer._int_set = mod("int_set")
+ analyzer._enter_constraint_context = mod("enter_constraint_context")
+ analyzer._can_prove_equal = mod("can_prove_equal")
+ analyzer._can_prove = mod("can_prove")
+ analyzer._get_enabled_extensions = mod("get_enabled_extensions")
+ analyzer._set_enabled_extensions = mod("set_enabled_extensions")
+ return analyzer
Review Comment:
The `_analyzer_from_module` function uses a fragile approach to reconstruct
a Python `Analyzer` object by bypassing `__init__` and manually populating
internal attributes (prefixed with `_`). In standard TVM,
`tvm.arith.analyzer.Analyzer` is a `tvm.runtime.Object` that expects a valid
C++ handle initialized during `__init__`. Bypassing this means that any method
on the resulting `analyzer` object that invokes an FFI call (e.g., `can_prove`,
`simplify`) will likely fail because the underlying C++ handle is null.
Consider defining a dedicated Python wrapper class in TIRX that uses the
provided `mod` functions directly, or ensure that the
`tvm.arith.analyzer.Analyzer` class is specifically designed to support this
kind of mocking.
##########
src/tirx/ir/py_functor.cc:
##########
@@ -23,13 +23,131 @@
* StmtExprVisitor/StmtExprMutator.
*/
+#include <tvm/arith/analyzer.h>
#include <tvm/ffi/reflection/registry.h>
+#include <tvm/ir/scope_stack.h>
+#include <tvm/ir/with_context.h>
+#include <tvm/s_tir/stmt.h>
+#include <tvm/tirx/analysis.h>
+#include <tvm/tirx/builtin.h>
#include <tvm/tirx/expr_functor.h>
+#include <tvm/tirx/op.h>
#include <tvm/tirx/stmt_functor.h>
+#include <cstdint>
+#include <memory>
+#include <string>
+
namespace tvm {
namespace tirx {
+namespace {
+
+ffi::Function MakeAnalyzerModule(std::shared_ptr<arith::Analyzer> analyzer) {
+ using ffi::Function;
+ using ffi::TypedFunction;
+ auto f = [analyzer](std::string name) -> ffi::Function {
+ if (name == "const_int_bound") {
+ return ffi::Function([analyzer](ffi::PackedArgs args, ffi::Any* ret) {
+ *ret = analyzer->const_int_bound(args[0].cast<PrimExpr>());
+ });
+ } else if (name == "modular_set") {
+ return ffi::Function([analyzer](ffi::PackedArgs args, ffi::Any* ret) {
+ *ret = analyzer->modular_set(args[0].cast<PrimExpr>());
+ });
+ } else if (name == "const_int_bound_update") {
+ return ffi::Function([analyzer](ffi::PackedArgs args, ffi::Any* ret) {
+ analyzer->const_int_bound.Update(args[0].cast<Var>(),
args[1].cast<arith::ConstIntBound>(),
+ args[2].cast<bool>());
+ });
+ } else if (name == "const_int_bound_is_bound") {
+ return ffi::Function([analyzer](ffi::PackedArgs args, ffi::Any* ret) {
+ *ret = analyzer->const_int_bound.IsBound(args[0].cast<Var>());
+ });
+ } else if (name == "Simplify") {
+ return ffi::Function([analyzer](ffi::PackedArgs args, ffi::Any* ret) {
+ if (args.size() == 1) {
+ *ret = analyzer->Simplify(args[0].cast<PrimExpr>());
+ } else if (args.size() == 2) {
+ *ret = analyzer->Simplify(args[0].cast<PrimExpr>(),
args[1].cast<int>());
+ } else {
+ TVM_FFI_THROW(InternalError) << "Invalid size of argument (" <<
args.size() << ")";
+ }
+ });
+ } else if (name == "rewrite_simplify") {
+ return ffi::Function([analyzer](ffi::PackedArgs args, ffi::Any* ret) {
+ *ret = analyzer->rewrite_simplify(args[0].cast<PrimExpr>());
+ });
+ } else if (name == "get_rewrite_simplify_stats") {
+ return ffi::Function([analyzer](ffi::PackedArgs args, ffi::Any* ret) {
+ *ret = analyzer->rewrite_simplify.GetStatsCounters();
+ });
+ } else if (name == "reset_rewrite_simplify_stats") {
+ return ffi::Function([analyzer](ffi::PackedArgs args, ffi::Any* ret) {
+ analyzer->rewrite_simplify.ResetStatsCounters();
+ });
+ } else if (name == "canonical_simplify") {
+ return ffi::Function([analyzer](ffi::PackedArgs args, ffi::Any* ret) {
+ *ret = analyzer->canonical_simplify(args[0].cast<PrimExpr>());
+ });
+ } else if (name == "int_set") {
+ return ffi::Function([analyzer](ffi::PackedArgs args, ffi::Any* ret) {
+ *ret = analyzer->int_set(args[0].cast<PrimExpr>(),
+ args[1].cast<ffi::Map<Var, arith::IntSet>>());
+ });
+ } else if (name == "bind") {
+ return ffi::Function([analyzer](ffi::PackedArgs args, ffi::Any* ret) {
+ bool allow_override = args.size() >= 3 && args[2].cast<bool>();
+ if (auto opt_range = args[1].try_cast<Range>()) {
+ analyzer->Bind(args[0].cast<Var>(), opt_range.value(),
allow_override);
+ } else {
+ analyzer->Bind(args[0].cast<Var>(), args[1].cast<PrimExpr>(),
allow_override);
+ }
+ });
+ } else if (name == "can_prove") {
+ return ffi::Function([analyzer](ffi::PackedArgs args, ffi::Any* ret) {
+ int strength = args[1].cast<int>();
+ *ret = analyzer->CanProve(args[0].cast<PrimExpr>(),
+ static_cast<arith::ProofStrength>(strength));
+ });
+ } else if (name == "enter_constraint_context") {
+ return ffi::Function([analyzer](ffi::PackedArgs args, ffi::Any* ret) {
+ auto ctx = std::shared_ptr<With<arith::ConstraintContext>>(
Review Comment:
Using `new` with `std::shared_ptr` is generally discouraged in modern C++.
While `With<arith::ConstraintContext>` is a RAII guard typically used on the
stack, since you need to manage its lifetime across the FFI boundary via a
lambda capture, `std::make_shared` should be preferred for better memory
efficiency and exception safety.
```c
auto ctx =
std::make_shared<With<arith::ConstraintContext>>(analyzer.get(),
args[0].cast<PrimExpr>());
```
--
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]
