(tvm) branch main updated: [TVMScript] Add block name suffix management for TIR macros (#18465)

[tlopex]

This is an automated email from the ASF dual-hosted git repository.

tlopex pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git


The following commit(s) were added to refs/heads/main by this push:
     new ced7181708 [TVMScript] Add block name suffix management for TIR macros 
(#18465)
ced7181708 is described below

commit ced7181708b6359d86d6e5f7196daa51c552e628
Author: Guan-Ming (Wesley) Chiu <105915352+guan404m...@users.noreply.github.com>
AuthorDate: Tue Nov 25 14:34:23 2025 +0800

    [TVMScript] Add block name suffix management for TIR macros (#18465)
    
    ## Related Issue
    
    closes https://github.com/apache/tvm/issues/18344
    
    ## Why
    
    When a `T.macro` containing a block was called multiple times in a TIR
    function, all expanded blocks had the same name, causing a "Duplicated
    block name" error in meta_schedule.
    
    ## How
    
    Implemented automatic block name suffixing during macro expansion
---
 python/tvm/script/ir_builder/tir/ir.py | 35 ++++++++++++++++++++++++++++++++++
 python/tvm/script/parser/tir/entry.py  | 20 ++++++++++++++++---
 2 files changed, 52 insertions(+), 3 deletions(-)

diff --git a/python/tvm/script/ir_builder/tir/ir.py 
b/python/tvm/script/ir_builder/tir/ir.py
index 31e48260f5..a08e66789f 100644
--- a/python/tvm/script/ir_builder/tir/ir.py
+++ b/python/tvm/script/ir_builder/tir/ir.py
@@ -16,9 +16,11 @@
 # under the License.
 """IRBuilder for TIR"""
 
+import contextlib
 import functools
 import inspect
 import sys
+import threading
 from numbers import Integral
 from typing import Any, Callable, Dict, List, Optional, Tuple, Union
 
@@ -87,6 +89,35 @@ from .external_kernel import call_kernel
 # pylint: enable=unused-import
 
 
+_block_name_suffix = threading.local()
+
+
+def _get_block_name_suffix() -> str:
+    """Get the current block name suffix for macro expansion."""
+    return getattr(_block_name_suffix, "value", "")
+
+
+@contextlib.contextmanager
+def block_name_suffix_context(block_suffix: str):
+    """Context manager to set block name suffix during macro expansion.
+
+    Parameters
+    ----------
+    block_suffix : str
+        The suffix to append to block names (e.g., "_1", "_2").
+
+    Yields
+    ------
+    None
+    """
+    old_suffix = getattr(_block_name_suffix, "value", "")
+    _block_name_suffix.value = block_suffix
+    try:
+        yield
+    finally:
+        _block_name_suffix.value = old_suffix
+
+
 def buffer(
     shape: Union[List[PrimExpr], Tuple[PrimExpr], PrimExpr, Integral],
     dtype: str = "float32",
@@ -352,6 +383,9 @@ def block(name: str = "", no_realize: bool = False) -> 
frame.BlockFrame:
     res : frame.BlockFrame
         The BlockFrame.
     """
+    block_suffix = _get_block_name_suffix()
+    if block_suffix and name:
+        name = name + block_suffix
     return _ffi_api.Block(name, no_realize)  # type: ignore[attr-defined] # 
pylint: disable=no-member
 
 
@@ -2135,6 +2169,7 @@ __all__ = float_types + [
     "func_ret",
     "match_buffer",
     "block",
+    "block_name_suffix_context",
     "init",
     "where",
     "reads",
diff --git a/python/tvm/script/parser/tir/entry.py 
b/python/tvm/script/parser/tir/entry.py
index c7d5dc756b..bcac49733d 100644
--- a/python/tvm/script/parser/tir/entry.py
+++ b/python/tvm/script/parser/tir/entry.py
@@ -21,7 +21,7 @@ from typing import Callable, Optional, Union
 from tvm.ir.base import deprecated
 from tvm.tir import Buffer, PrimFunc
 
-from ...ir_builder.tir import buffer, ptr
+from ...ir_builder.tir import block_name_suffix_context, buffer, ptr
 from .._core import parse, scan_macro, utils
 from ..core.parser import Parser, ScriptMacro
 
@@ -90,11 +90,25 @@ setattr(prim_func, "dispatch_token", "tir")
 
 
 class TIRMacro(ScriptMacro):
-    """Specialization of the ScriptMacro class for TIR."""
+    """Specialization of the ScriptMacro class for TIR.
+
+    Attributes
+    ----------
+    call_count : int
+        Counter for the number of times this macro has been invoked.
+        Used to generate unique block name suffixes.
+    """
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.call_count = 0
 
     def parse_macro(self, parser: Parser) -> None:
         macro_def = self.get_macro_def()
-        parser.visit_body(macro_def.body)
+        suffix = f"_{self.call_count}" if self.call_count > 0 else ""
+        self.call_count += 1
+        with block_name_suffix_context(suffix):
+            parser.visit_body(macro_def.body)
 
 
 def macro(*args, hygienic: bool = True) -> Callable: