From 2df09c17604e46aac7de00b0ec37cb20675866a8 Mon Sep 17 00:00:00 2001
From: Eric Lunderberg <elunderberg@octoml.ai>
Date: Thu, 12 Oct 2023 15:05:51 -0500
Subject: [PATCH] [Unity][TVMScript] Avoid dangling reference when printing
 Call attrs

Prior to this commit, the `tvm::script::printer::AttrPrinter` class
took the attribute path as a `const ObjectPath&`.  In both places
where an `AttrPrinter` is called, the temporary object
`n_p->Attr("attrs")` is passed for this argument.  While binding a
temporary object to a const reference can extend the lifetime of the
temporary, this requires the const reference to be in the same scope
as the temporary, and does not apply in this case (see [this
stackoverflow post](https://stackoverflow.com/a/2784304)).  Therefore,
this reference is only valid through the construction of `AttrPrinter
printer`, and is invalid during its usage on the following line.

This dangling reference has caused segfaults in CI for unrelated
changes ([example](https://ci.tlcpack.ai/blue/organizations/jenkins/tvm-unity/detail/PR-15904/3/pipeline)),
and can be reproduced with the following test case.

```python
import pytest

from tvm.script import relax as R

@pytest.mark.parametrize("iter", range(10000))
def test_argmax_without_specified_axis(iter):
    @R.function
    def func(x: R.Tensor((1, 2, 3, 4), "float32")):
        return R.argmax(x)

    func.script(show_meta=True)
```

This test case is not included in this commit, as the reproduction is
not consistent, with failure requiring on the order of 10k iterations
to trigger.  In addition, reproduction was sensitive to the following
conditions.

* The function being printed must contain at least one `relax::Call`
  node, with an operation that has attributes.

* TVM must be built with optimization enabled.  In gcc, the
  `-ftree-dse` optimization, which is part of `-O1`, is required to
  trigger the bug.

* Python's default allocation must be used.  If `PYTHONMALLOC=malloc`
  is set to instead use the system's `malloc`, the segfault was no
  longer triggered.

This commit updates `AttrPrinter` to accept the `ObjectPath` by value.
With the change applied, the above test ran 100k times without error.
---
 src/script/printer/relax/call.cc | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/src/script/printer/relax/call.cc b/src/script/printer/relax/call.cc
index dc57ca318406..79461df0833d 100644
--- a/src/script/printer/relax/call.cc
+++ b/src/script/printer/relax/call.cc
@@ -27,9 +27,9 @@ namespace printer {
 
 class AttrPrinter : public tvm::AttrVisitor {
  public:
-  explicit AttrPrinter(const ObjectPath& p, const IRDocsifier& d, Array<String>* keys,
+  explicit AttrPrinter(ObjectPath p, const IRDocsifier& d, Array<String>* keys,
                        Array<ExprDoc>* values)
-      : p(p), d(d), keys(keys), values(values) {}
+      : p(std::move(p)), d(d), keys(keys), values(values) {}
 
   void Visit(const char* key, double* value) final {
     keys->push_back(key);
@@ -79,7 +79,7 @@ class AttrPrinter : public tvm::AttrVisitor {
     LOG(FATAL) << "TypeError: NDArray is not allowed in Attrs";
   }
 
-  const ObjectPath& p;
+  ObjectPath p;
   const IRDocsifier& d;
   Array<String>* keys;
   Array<ExprDoc>* values;