Tianqi Chen
a6c7377bae
AttrStmt.node is an ffi::Any field, so converting POD arguments to PrimExpr makes its representation depend on the caller rather than the declared container type. This change preserves values passed through AttrStmt and T.attr, uses raw zero sentinel nodes consistently, and updates the printer canonicalization.
227 lines
6.3 KiB
Python
227 lines
6.3 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.
|
|
"""Utility helpers for TIR IRBuilder."""
|
|
|
|
import contextlib
|
|
|
|
from tvm import tirx
|
|
from tvm.tirx import Buffer
|
|
|
|
from . import frame
|
|
from . import ir as T
|
|
|
|
|
|
class _FrameScope:
|
|
"""Context manager to enter multiple IRBuilder frames without deep nesting.
|
|
|
|
This class allows entering multiple frames in a single `with` statement,
|
|
avoiding the pyramid of nested context managers.
|
|
|
|
Parameters
|
|
----------
|
|
frames : List[IRBuilderFrame]
|
|
The list of frames to enter.
|
|
"""
|
|
|
|
def __init__(self, frames):
|
|
self.frames = frames if isinstance(frames, list | tuple) else [frames]
|
|
self._stack = None
|
|
|
|
def __enter__(self):
|
|
self._stack = contextlib.ExitStack()
|
|
self._stack.__enter__()
|
|
results = [self._stack.enter_context(f) for f in self.frames]
|
|
return tuple(results) if len(results) > 1 else results[0]
|
|
|
|
def __exit__(self, *args):
|
|
return self._stack.__exit__(*args)
|
|
|
|
|
|
def frame_scope(frames: list[frame.TIRFrame]) -> _FrameScope:
|
|
"""Enter multiple IRBuilder frames without deep nesting.
|
|
|
|
This function provides a way to enter multiple frames in a single `with`
|
|
statement, which is particularly useful when migrating from cases where
|
|
allocations don't require nested scopes.
|
|
|
|
Parameters
|
|
----------
|
|
frames : List[frame.TIRFrame]
|
|
The list of frames to enter. Each frame's `__enter__` return value
|
|
will be collected and returned as a tuple.
|
|
|
|
Returns
|
|
-------
|
|
_FrameScope
|
|
A context manager that enters all frames and returns their values.
|
|
"""
|
|
return _FrameScope(frames)
|
|
|
|
|
|
def seq_scope():
|
|
"""Create a scope that allows multiple consecutive statements.
|
|
|
|
The IRBuilder requires a parent frame when having multiple consecutive
|
|
top-level statements (e.g., multiple loops). This function creates a
|
|
dummy attr frame that serves as a parent scope.
|
|
|
|
Returns
|
|
-------
|
|
frame.AttrFrame
|
|
A dummy attribute frame that wraps multiple statements.
|
|
|
|
Examples
|
|
--------
|
|
Without seq_scope, multiple consecutive loops fail:
|
|
|
|
.. code-block:: python
|
|
|
|
with IRBuilder() as ib:
|
|
with T.serial(0, 10) as i:
|
|
T.evaluate(i)
|
|
with T.serial(0, 5) as j: # This would fail!
|
|
T.evaluate(j)
|
|
|
|
With seq_scope, multiple consecutive statements work:
|
|
|
|
.. code-block:: python
|
|
|
|
with IRBuilder() as ib:
|
|
with seq_scope():
|
|
with T.serial(0, 10) as i:
|
|
T.evaluate(i)
|
|
with T.serial(0, 5) as j:
|
|
T.evaluate(j)
|
|
result = ib.get()
|
|
"""
|
|
return T.attr(0, "pragma_scope", tirx.StringImm("seq"))
|
|
|
|
|
|
def _unravel_index(index, shape):
|
|
"""Convert a flat index to multi-dimensional indices.
|
|
|
|
Parameters
|
|
----------
|
|
index : Expr
|
|
The flat index.
|
|
shape : Tuple
|
|
The shape of the buffer.
|
|
|
|
Returns
|
|
-------
|
|
List[Expr]
|
|
The multi-dimensional indices.
|
|
"""
|
|
indices = []
|
|
for i, dim in enumerate(reversed(shape)):
|
|
if i == len(shape) - 1:
|
|
# Outermost dimension: use remaining quotient directly (no modulo)
|
|
indices.append(index)
|
|
else:
|
|
indices.append(index % dim)
|
|
index = index // dim
|
|
return list(reversed(indices))
|
|
|
|
|
|
class _BufferProxy:
|
|
"""Proxy for flat indexing on multi-dimensional buffers.
|
|
|
|
This class wraps a TIR Buffer and provides flat indexing that gets
|
|
automatically converted to multi-dimensional indices. It also supports
|
|
assignment syntax via __setitem__.
|
|
|
|
Parameters
|
|
----------
|
|
buf : Buffer
|
|
The TIR buffer to wrap.
|
|
|
|
Examples
|
|
--------
|
|
.. code-block:: python
|
|
|
|
buf = tvm.tirx.decl_buffer([2, 3], "float32")
|
|
ptr = buffer_proxy(buf)
|
|
|
|
# Read with flat index (converted to [0, 1])
|
|
val = ptr[1]
|
|
|
|
# Write with flat index
|
|
ptr[1] = 42.0
|
|
|
|
# Multi-dimensional access still works
|
|
val = ptr[0, 2]
|
|
"""
|
|
|
|
def __init__(self, buf):
|
|
self._buffer = buf
|
|
self.dtype = buf.dtype
|
|
self.shape = buf.shape
|
|
self.name = buf.name
|
|
self.data = buf.data
|
|
|
|
def _normalize_index(self, index):
|
|
"""Convert flat index to multi-dimensional indices if needed."""
|
|
try:
|
|
index = [*index]
|
|
except TypeError:
|
|
index = [index]
|
|
if len(index) == 1 and len(self._buffer.shape) != 1:
|
|
index = _unravel_index(index[0], self._buffer.shape)
|
|
return index
|
|
|
|
def __getitem__(self, index):
|
|
index = self._normalize_index(index)
|
|
return tirx.BufferLoad(self._buffer, index)
|
|
|
|
def __setitem__(self, index, value):
|
|
index = self._normalize_index(index)
|
|
T.buffer_store(self._buffer, value, index)
|
|
|
|
|
|
def buffer_proxy(buf: Buffer) -> _BufferProxy:
|
|
"""Create a buffer proxy for flat indexing on multi-dimensional buffers.
|
|
|
|
This provides flat indexing that gets converted to multi-dimensional indices.
|
|
It also supports assignment syntax via __setitem__.
|
|
|
|
Parameters
|
|
----------
|
|
buf : Buffer
|
|
The TIR buffer to wrap.
|
|
|
|
Returns
|
|
-------
|
|
_BufferProxy
|
|
A proxy object that supports flat indexing and assignment.
|
|
|
|
Examples
|
|
--------
|
|
.. code-block:: python
|
|
|
|
from tvm.tirx.script.builder.utils import buffer_proxy
|
|
|
|
buf = tvm.tirx.decl_buffer([2, 3], "float32")
|
|
ptr = buffer_proxy(buf)
|
|
|
|
# Flat indexing (index 1 -> indices [0, 1])
|
|
val = ptr[1]
|
|
|
|
# Assignment syntax
|
|
ptr[1] = 42.0
|
|
"""
|
|
return _BufferProxy(buf)
|