chore: import upstream snapshot with attribution

python/tvm/backend/trn/transform/__init__.py

@@ -0,0 +1,53 @@
+# 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.
+"""Trainium-specific TIRX transformations."""
+# pylint: disable=invalid-name
+
+from tvm_ffi import get_global_func
+
+import tvm
+from tvm import tirx
+
+_LAZY_TRANSFORMS = {
+    "TrnNaiveAllocator": ".naive_allocator",
+    "TrnPrivateBufferAlloc": ".private_buffer_alloc",
+}
+
+
+def LowerTrainiumLayout():
+    """Lower Trainium layouts to backend physical buffer shapes and indices."""
+    return get_global_func("tirx.backend.trn.transform.LowerTrainiumLayout")()
+
+
+def LowerTIRx():
+    """Lower TIRx tile primitive calls for the Trainium backend."""
+    return tvm.ir.transform.Sequential(
+        [tirx.transform.TilePrimitiveDispatch(), LowerTrainiumLayout()],
+        name="tirx.backend.trn.LowerTIRx",
+    )
+
+
+def __getattr__(name):
+    target = _LAZY_TRANSFORMS.get(name)
+    if target is None:
+        raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
+    from importlib import import_module  # pylint: disable=import-outside-toplevel
+
+    return getattr(import_module(target, __name__), name)
+
+
+__all__ = ["LowerTIRx", "LowerTrainiumLayout", "TrnNaiveAllocator", "TrnPrivateBufferAlloc"]

python/tvm/backend/trn/transform/naive_allocator.py

@@ -0,0 +1,99 @@
+# 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.
+
+import functools
+
+from tvm.tirx import AllocBuffer, IntImm
+from tvm.tirx.buffer import Buffer
+from tvm.tirx.stmt_functor import StmtVisitor
+from tvm.tirx.transform.common import BufferReplacer
+from tvm.tirx.transform.function_pass import prim_func_pass
+
+
+def is_const_shape(buffer: Buffer) -> bool:
+    for i in buffer.shape:
+        if not isinstance(i, IntImm):
+            return False
+    return True
+
+
+def get_buffer_size(buffer: Buffer) -> int:
+    if buffer.scope() == "trn.sbuf":
+        if buffer.layout is None:
+            # the first dimension is partition size
+            num_elem = functools.reduce(lambda x, y: x * y, buffer.shape[1:])
+        else:
+            par_size = buffer.layout.size("P")
+            num_elem = functools.reduce(lambda x, y: x * y, buffer.shape) // par_size
+    elif buffer.scope().startswith("shared"):
+        num_elem = functools.reduce(lambda x, y: x * y, buffer.shape)
+    else:
+        return None
+    if not is_const_shape(buffer):
+        raise ValueError(
+            f"Buffer {buffer.name} has non-constant shape. Do not know how to allocate it."
+        )
+    return int(num_elem * buffer.dtype.dtype.itemsize)
+
+
+class AllocInfoCollector(StmtVisitor):
+    def __init__(self):
+        super().__init__()
+        self.alloc_pool_start = 0
+
+    def visit_alloc_buffer_(self, op: AllocBuffer):
+        super().visit_alloc_buffer_(op)
+        buffer = op.buffer
+        if len(buffer.allocated_addr) == 0:
+            return op
+        buffer_size = get_buffer_size(buffer)
+        if buffer_size is None:
+            return op
+        self.alloc_pool_start = max(self.alloc_pool_start, buffer.allocated_addr[-1] + buffer_size)
+
+
+class AllocMutator(BufferReplacer):
+    def __init__(self, alloc_pool_start: int):
+        super().__init__()
+        self.alloc_offset = alloc_pool_start
+
+    def visit_alloc_buffer_(self, op: AllocBuffer):
+        changed = False
+        buffer = op.buffer
+        buffer_size = get_buffer_size(buffer)
+        if len(buffer.allocated_addr) > 0 or buffer_size is None:
+            pass
+        else:
+            new_buffer = buffer.with_allocated_addr([self.alloc_offset])
+            self.buffer_map[buffer] = new_buffer
+            changed = True
+            self.alloc_offset += buffer_size
+
+        op = super().visit_alloc_buffer_(op)
+        if changed:
+            return AllocBuffer(new_buffer, op.annotations, op.span)
+        return op
+
+
+@prim_func_pass(opt_level=0, name="TrnNaiveAllocator")
+class TrnNaiveAllocator:
+    def transform_function(self, func, mod, ctx):
+        collector = AllocInfoCollector()
+        collector(func.body)
+        mutator = AllocMutator(collector.alloc_pool_start)
+        new_body = mutator(func.body)
+        return func.with_body(new_body)

python/tvm/backend/trn/transform/private_buffer_alloc.py

@@ -0,0 +1,138 @@
+# 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.
+
+
+from tvm.ir import Range
+from tvm.target import Target
+from tvm.tirx.buffer import Buffer
+from tvm.tirx.operator.tile_primitive.dispatch_context import DispatchContext
+from tvm.tirx.stmt import (
+    AllocBuffer,
+    AttrStmt,
+    For,
+    SeqStmt,
+    Stmt,
+    TilePrimitiveCall,
+)
+from tvm.tirx.stmt_functor import StmtMutator, StmtVisitor
+from tvm.tirx.transform.common import seek_kernel_replace_point
+from tvm.tirx.transform.function_pass import prim_func_pass
+
+
+class PrivateAllocCollector(StmtVisitor):
+    def __init__(self, target: Target):
+        super().__init__()
+        self.target = target
+        self.launch_params = {}
+        self.var_range_map = {}
+        self.buffer_dict = {}
+        self.private_buf_refs = {}
+
+    def visit_attr_(self, op: AttrStmt):
+        if op.attr_key == "thread_extent":
+            self.launch_params[op.node.thread_tag] = op.value
+        super().visit_attr_(op)
+
+    def visit_for_(self, op: For):
+        self.var_range_map[op.loop_var] = Range.from_min_extent(op.min, op.extent)
+        super().visit_for_(op)
+
+    def visit_op_call_(self, op: TilePrimitiveCall):
+        # Scope is a per-call field on the node; read it directly.
+        exec_scope = op.scope
+        scope_kind = op.scope.name
+        sctx = DispatchContext(
+            target=self.target,
+            exec_scope=exec_scope,
+            launch_params=self.launch_params,
+            var_range_map=self.var_range_map,
+            alloc_only=True,
+            scope_kind=scope_kind,
+        )
+        op = TilePrimitiveCall.downcast(op)
+        private_buf_refs = op.get_private_buffers(self.buffer_dict, sctx)
+        self.private_buf_refs[op] = private_buf_refs
+
+
+class PrivateAllocMutator(StmtMutator):
+    def __init__(
+        self,
+        alloc_buffers: list[Buffer],
+        init_stmts: list[Stmt],
+        added_workspace: dict[TilePrimitiveCall, dict[str, Buffer]],
+    ):
+        super().__init__()
+        self.alloc_buffers = alloc_buffers
+        self.init_stmts = init_stmts
+        self.added_workspace = added_workspace
+        self.is_outer_block = True
+
+    def visit_attr_(self, op: AttrStmt):
+        # AttrStmt(kDeviceEntry) marks the device-region root: inject the
+        # collected init stmts + alloc_buffers into its body.
+        if op.attr_key == "tirx.device_entry":
+            is_outer_block = self.is_outer_block
+            self.is_outer_block = False
+            op = super().visit_attr_(op)
+            if is_outer_block:
+                body = op.body
+                for stmt in self.init_stmts:
+                    body = seek_kernel_replace_point(stmt, body)
+                for buffer in reversed(self.alloc_buffers):
+                    body = SeqStmt([AllocBuffer(buffer), body])
+                return AttrStmt(op.node, op.attr_key, op.value, body)
+            return op
+        return super().visit_attr_(op)
+
+    def visit_op_call_(self, op):
+        if op not in self.added_workspace:
+            return op
+        new_workspace = dict(op.workspace)
+        new_workspace.update(self.added_workspace[op])
+        op = TilePrimitiveCall.downcast(op).with_workspace(new_workspace)
+        return op
+
+
+def private_alloc(stmt: Stmt, target: Target) -> Stmt:
+    collector = PrivateAllocCollector(target)
+    collector(stmt)
+
+    alloc_buffers = [buffer for buffer, _ in collector.buffer_dict.values()]
+    init_stmts = [stmt for _, stmt in collector.buffer_dict.values() if stmt is not None]
+    added_workspace = {
+        op: {
+            name: collector.buffer_dict[ref][0]
+            for name, ref in collector.private_buf_refs[op].items()
+        }
+        for op in collector.private_buf_refs
+    }
+
+    mutator = PrivateAllocMutator(alloc_buffers, init_stmts, added_workspace)
+    return mutator(stmt)
+
+
+@prim_func_pass(opt_level=0, name="TrnPrivateBufferAlloc")
+class TrnPrivateBufferAlloc:
+    """Generate private buffer allocations for each TilePrimitiveCall"""
+
+    def transform_function(self, func, mod, ctx):
+        target = func.attrs.get("target", None)
+        if target is None:
+            target = Target.current(allow_none=False)
+        new_body = private_alloc(func.body, target)
+        new_func = func.with_body(new_body)
+        return new_func