chore: import upstream snapshot with attribution

2026-07-13 13:36:25 +08:00
commit 26446540fa
3151 changed files with 974126 additions and 0 deletions
@@ -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