Hongyi Jin
ae99c3fd92
## Motivation and context The TVMScript parser already tracks Python AST locations for diagnostics, but TIRx statements and expression results emitted through `IRBuilder` did not retain those locations. After parsing, a direct intrinsic call, an inlined helper body, or a `TilePrimitiveCall` therefore could not be traced back to the source range that produced it. Inline expansion also needs more than a single flat location. The generated IR should retain both the caller location and the helper-definition location, while ordinary nested AST evaluation within one source should not accumulate redundant enclosing spans. ## Changes - Add an active source-span stack to `IRBuilder`, with scoped push/pop support. - Make the parser activate the current AST source range while visiting statements and evaluating expressions. - Attach the active span to emitted TIRx statements and to expression results that do not already carry an explicit span. - Normalize nested spans from the same source to the innermost relevant range. - Preserve cross-source inline expansion history as a `SequentialSpan`, ordered from the call site to the expanded definition. - Reuse the same source-coordinate calculation for diagnostics and IR spans so their line and column conventions remain consistent. Source spans remain diagnostic metadata: functions parsed from different source locations keep the same structural hash and remain structurally equal. ## Testing - Verify exact parser source coordinates against diagnostic coordinates. - Verify spans on direct intrinsic calls and `TilePrimitiveCall` nodes. - Verify that inline expansion produces a `SequentialSpan` containing caller and callee ranges. - Verify direct `IRBuilder.with_source_span` behavior. - Verify that source spans do not affect structural identity. - Run changed-files pre-commit checks, including clang-format. Focused result: 9 tests passed.
232 lines
7.7 KiB
Python
232 lines
7.7 KiB
Python
# Licensed to the Apache Software Foundation (ASF) under one
|
|
# or more contributor license agreements. See the NOTICE file
|
|
# distributed with this work for additional information
|
|
# regarding copyright ownership. The ASF licenses this file
|
|
# to you under the Apache License, Version 2.0 (the
|
|
# "License"); you may not use this file except in compliance
|
|
# with the License. You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing,
|
|
# software distributed under the License is distributed on an
|
|
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
|
# KIND, either express or implied. See the License for the
|
|
# specific language governing permissions and limitations
|
|
# under the License.
|
|
"""A generic IRBuilder across the TVM stack"""
|
|
|
|
from collections.abc import Callable
|
|
from contextlib import contextmanager
|
|
from typing import Any
|
|
|
|
from tvm_ffi import register_object as _register_object
|
|
|
|
from tvm.runtime import Object as _Object
|
|
|
|
from . import _ffi_api
|
|
|
|
|
|
@_register_object("script.ir_builder.IRBuilderFrame")
|
|
class IRBuilderFrame(_Object):
|
|
"""A stack frame of the IRBuilder used to keep track of the current scope.
|
|
|
|
Furthermore, the information stored in each stack frame can be useful for context-dependent
|
|
IR construction.
|
|
|
|
Examples
|
|
--------
|
|
|
|
The `T.match_buffer` below instead an element in the buffer map of `PrimFuncFrame`:
|
|
|
|
.. code-block:: python
|
|
|
|
from tvm.script.ir_builder import tirx as T
|
|
from tvm.script.ir_builder import IRBuilder
|
|
|
|
with IRBuilder() as builder:
|
|
with T.prim_func(...): # pushes a PrimFuncFrame (subclass of IRBuilderFrame)
|
|
# to `builder`'s stack of frames
|
|
buffer = T.match_buffer(...)
|
|
|
|
|
|
The `T.match_buffer` below instead generates `MatchBufferRegion` in a TIR block:
|
|
|
|
.. code-block:: python
|
|
|
|
from tvm.script.ir_builder import tirx as T
|
|
from tvm.script.ir_builder import IRBuilder
|
|
|
|
with IRBuilder() as builder:
|
|
with T.prim_func(...): # pushes a PrimFuncFrame (subclass of IRBuilderFrame)
|
|
# to `builder`'s stack of frames
|
|
with T.sblock(...): # pushes a BlockFrame (subclass of IRBuilderFrame)
|
|
# to `builder`'s stack of frames
|
|
buffer = T.match_buffer(...)
|
|
|
|
"""
|
|
|
|
def __enter__(self) -> "IRBuilderFrame":
|
|
_ffi_api.IRBuilderFrameEnter(self) # type: ignore[attr-defined] # pylint: disable=no-member
|
|
return self
|
|
|
|
def __exit__(self, exc_type, exc_value, trace) -> None: # pylint: disable=unused-argument
|
|
if exc_type is None and exc_value is None:
|
|
# Do not execute `FrameExit` if the with scope exits because of exceptions
|
|
_ffi_api.IRBuilderFrameExit(self) # type: ignore[attr-defined] # pylint: disable=no-member
|
|
|
|
def add_callback(self, callback: Callable[[], None]) -> None:
|
|
"""Add a callback method invoked when exiting the with-scope.
|
|
|
|
Parameters
|
|
----------
|
|
callback : Callable[[], None]
|
|
The callback method to be invoked.
|
|
"""
|
|
_ffi_api.IRBuilderFrameAddCallback( # type: ignore[attr-defined] # pylint: disable=no-member
|
|
self, callback
|
|
)
|
|
|
|
|
|
@_register_object("script.ir_builder.IRBuilder")
|
|
class IRBuilder(_Object):
|
|
"""A dialect-agnostic IRBuilder that constructs any IR of TVM.
|
|
|
|
Examples
|
|
--------
|
|
An idiomatic use of this class is to put this inside the with-scope,
|
|
call dialect-specific methods accordingly. Upon exiting the scope.
|
|
|
|
.. code-block:: python
|
|
|
|
from tvm.script.ir_builder import tirx as T
|
|
from tvm.script.ir_builder import IRBuilder
|
|
|
|
with IRBuilder() as builder:
|
|
with T.prim_func(...): # pushes a PrimFuncFrame (subclass of IRBuilderFrame)
|
|
# to `builder`'s stack of frames
|
|
buffer = T.match_buffer(...)
|
|
|
|
return builder.get() # returns the constructed IR, i.e. tirx.PrimFunc
|
|
"""
|
|
|
|
def __init__(self) -> None:
|
|
"""Construct an IRBuilder."""
|
|
self.__init_handle_by_constructor__(
|
|
_ffi_api.IRBuilder # type: ignore[attr-defined] # pylint: disable=no-member
|
|
)
|
|
|
|
def __enter__(self) -> "IRBuilder":
|
|
"""Enter the with-scope for IRBuilder, which allows the IRBuilder to be discoverable
|
|
using `IRBuilder.current()`.
|
|
|
|
Examples
|
|
--------
|
|
.. code-block:: python
|
|
|
|
from tvm.script.ir_builder import IRBuilder
|
|
|
|
with IRBuilder() as builder:
|
|
assert IRBuilder.current() == builder
|
|
|
|
"""
|
|
_ffi_api.IRBuilderEnter(self) # type: ignore[attr-defined] # pylint: disable=no-member
|
|
return self
|
|
|
|
def __exit__(self, ptype, value, trace) -> None: # pylint: disable=unused-argument
|
|
_ffi_api.IRBuilderExit(self) # type: ignore[attr-defined] # pylint: disable=no-member
|
|
|
|
@staticmethod
|
|
def current() -> "IRBuilder":
|
|
"""Get the current IRBuilder put in the with-scope.
|
|
|
|
Returns
|
|
-------
|
|
builder : IRBuilder
|
|
The current IRBuilder.
|
|
"""
|
|
return _ffi_api.IRBuilderCurrent() # type: ignore[attr-defined] # pylint: disable=no-member
|
|
|
|
@staticmethod
|
|
def is_in_scope() -> bool:
|
|
"""See if the current thread-local scope has an IRBuilder.
|
|
|
|
Returns
|
|
-------
|
|
bool
|
|
Whether the current thread-local scope has an IRBuilder
|
|
"""
|
|
return _ffi_api.IRBuilderIsInScope() # type: ignore[attr-defined] # pylint: disable=no-member
|
|
|
|
def get(self) -> _Object:
|
|
"""Get the constructed IR."""
|
|
return _ffi_api.IRBuilderGet(self) # type: ignore[attr-defined] # pylint: disable=no-member
|
|
|
|
@contextmanager
|
|
def with_source_span(self, span):
|
|
"""Attach ``span`` to IR nodes constructed in the nested scope.
|
|
|
|
Nested scopes are retained as a ``SequentialSpan`` when they describe
|
|
distinct source ranges, such as a TVMScript inline expansion.
|
|
|
|
Parameters
|
|
----------
|
|
span : tvm.ir.Span
|
|
The frontend source range active in the nested scope.
|
|
"""
|
|
_ffi_api.IRBuilderPushSourceSpan( # type: ignore[attr-defined] # pylint: disable=no-member
|
|
self, span
|
|
)
|
|
try:
|
|
yield
|
|
finally:
|
|
_ffi_api.IRBuilderPopSourceSpan( # type: ignore[attr-defined] # pylint: disable=no-member
|
|
self
|
|
)
|
|
|
|
def _set_current_source_span(self, value):
|
|
"""Attach the active source span to an expression without one."""
|
|
return _ffi_api.IRBuilderSetCurrentSourceSpan( # type: ignore[attr-defined] # pylint: disable=no-member
|
|
self, value
|
|
)
|
|
|
|
@staticmethod
|
|
def name(s: str, v: Any) -> Any:
|
|
"""Set the name of an object.
|
|
|
|
Parameters
|
|
----------
|
|
s : str
|
|
The name of the object.
|
|
v : Any
|
|
The object to name.
|
|
|
|
Returns
|
|
-------
|
|
v : Any
|
|
The same object with the name set.
|
|
"""
|
|
return _ffi_api.IRBuilderName(s, v) # type: ignore[attr-defined] # pylint: disable=no-member
|
|
|
|
@staticmethod
|
|
def name_many( # pylint: disable=invalid-name
|
|
s: list[str],
|
|
vs: list[Any],
|
|
) -> list[Any]:
|
|
"""Set the name of a list of objects.
|
|
|
|
Parameters
|
|
----------
|
|
s : List[str]
|
|
The names of the objects.
|
|
vs : List[Any]
|
|
The objects to name.
|
|
|
|
Returns
|
|
-------
|
|
vs : List[Any]
|
|
The same objects with the names set.
|
|
"""
|
|
assert len(s) == len(vs)
|
|
return [IRBuilder.name(i, v) for i, v in zip(s, vs)]
|