wehub-resource-sync
139 lines
4.9 KiB
Python
139 lines
4.9 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.
|
|
|
|
|
|
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
|