[Unity][BlockBuilder] Restore bb.get()

include/tvm/relax/block_builder.h

@@ -87,7 +87,11 @@ class BlockBuilderNode : public Object {
    * GlobalVars in the IRModule to ensure name uniqueness and the invariant:
    * every public function has the same name as its "global_symbol" attribute.
    *
-   * \return The IRModule in this BlockBuilder.
+   * \note this method should be called only once at the end of the building process, since it may
+   * invalidate global vars previously returned by this builder. See also
+   * transform::NormalizeGlobalVar.
+   *
+   * \return The result IRModule.
    */
   virtual IRModule Finalize() = 0;
 

python/tvm/relax/block_builder.py

@@ -22,7 +22,6 @@
 import tvm
 from tvm import relax as rx
 from tvm import tir
-from tvm.ir.base import deprecated
 from tvm.ir.module import IRModule
 from tvm.runtime import Object
 
@@ -654,26 +653,26 @@ def normalize(self, expr: Expr) -> Expr:
         """
         return _ffi_api.BlockBuilderNormalize(self, expr)  # type: ignore
 
-    @deprecated("tvm.relax.BlockBuilder.get", "tvm.relax.BlockBuilder.finalize")
     def get(self) -> tvm.IRModule:
-        """Return the IRModule being built.
+        """Return intermediate IRModule. For the situation where the IRModule is needed in the
+        middle of a building process.
 
         Returns
         -------
         ret : tvm.IRModule
             An IRModule with Relax and TIR functions being built.
         """
-        return self.finalize()
+        return _ffi_api.BlockBuilderGetContextIRModule(self)  # type: ignore
 
     def finalize(self) -> tvm.IRModule:
         """Finalize the building process and return the result IRModule.
 
         Possibly rename GlobalVars in the IRModule to ensure name uniqueness and the invariant:
         every public function has the same name as its "global_symbol" attribute.
 
-        Note this call may invalidate global vars previously returned by this builder
-        (see tvm.relax.transform.NormalizeGlobalVar), so it can only be called once at the end of
-        the building process.
+        Note this method should be called only once at the end of the building process, since it may
+        invalidate global vars previously returned by this builder.
+        See also tvm.relax.transform.NormalizeGlobalVar.
 
         Returns
         -------

tests/python/relax/test_blockbuilder_core.py

@@ -707,8 +707,10 @@ def te_one():
         gv1 = bb.emit_te(te_one, primfunc_name_hint="func")
         bb.emit_func_output((gv0, gv1))
 
-    mod = bb.finalize()
-    assert rx.analysis.well_formed(mod)
+    mod = bb.get()
+    assert not rx.analysis.well_formed(mod)
+    mod_final = bb.finalize()
+    assert rx.analysis.well_formed(mod_final)
 
     # relax function call
     bb = rx.BlockBuilder()
@@ -724,8 +726,10 @@ def te_one():
         gv0 = bb.emit(rx.Call(gvar1, []))
         bb.emit_func_output(gv0)
 
-    mod = bb.finalize()
-    assert rx.analysis.well_formed(mod)
+    mod = bb.get()
+    assert not rx.analysis.well_formed(mod)
+    mod_final = bb.finalize()
+    assert rx.analysis.well_formed(mod_final)
 
 
 def test_emit_nested_seqexpr_in_binding_block():