chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,191 @@
|
||||
# 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 Call, Op, is_prim_expr
|
||||
from tvm.tirx import (
|
||||
AllocBuffer,
|
||||
BufferLoad,
|
||||
BufferRegion,
|
||||
BufferStore,
|
||||
DeclBuffer,
|
||||
Evaluate,
|
||||
Expr,
|
||||
Stmt,
|
||||
TilePrimitiveCall,
|
||||
Var,
|
||||
decl_buffer,
|
||||
)
|
||||
from tvm.tirx.buffer import Buffer
|
||||
from tvm.tirx.layout import Iter, TileLayout
|
||||
from tvm.tirx.stmt_functor import StmtExprMutator, StmtMutator
|
||||
|
||||
|
||||
class BufferReplacer(StmtExprMutator):
|
||||
"""
|
||||
Replace buffer with another buffer.
|
||||
Also replace the data of the buffer with another var.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self, buffer_map: dict[Buffer, Buffer] | None = None, var_map: dict[Var, Var] | None = None
|
||||
):
|
||||
super().__init__()
|
||||
self.buffer_map = buffer_map if buffer_map is not None else {}
|
||||
self.var_map = var_map if var_map is not None else {}
|
||||
self.buffer_attr_var_mutated = False
|
||||
for old_buffer, new_buffer in self.buffer_map.items():
|
||||
self.var_map[old_buffer.data] = new_buffer.data
|
||||
|
||||
def mutate_buffer(self, buffer: Buffer):
|
||||
if buffer in self.buffer_map:
|
||||
return self.buffer_map[buffer]
|
||||
|
||||
# Track mutations for this specific buffer only. Without this reset,
|
||||
# unrelated buffers can be spuriously cloned and introduce alias buffers.
|
||||
prev_mutated = self.buffer_attr_var_mutated
|
||||
self.buffer_attr_var_mutated = False
|
||||
new_data = self.visit_expr(buffer.data)
|
||||
new_shape = [self.visit_expr(expr) for expr in buffer.shape]
|
||||
new_strides = [self.visit_expr(expr) for expr in buffer.strides]
|
||||
new_elem_offset = (
|
||||
self.visit_expr(buffer.elem_offset) if buffer.elem_offset is not None else None
|
||||
)
|
||||
if isinstance(buffer.layout, TileLayout):
|
||||
new_shard = []
|
||||
new_replicate = []
|
||||
for iter in buffer.layout.shard:
|
||||
new_iter = Iter(
|
||||
self.visit_expr(iter.extent), self.visit_expr(iter.stride), iter.axis
|
||||
)
|
||||
new_shard.append(new_iter)
|
||||
for iter in buffer.layout.replica:
|
||||
new_iter = Iter(
|
||||
self.visit_expr(iter.extent), self.visit_expr(iter.stride), iter.axis
|
||||
)
|
||||
new_replicate.append(new_iter)
|
||||
new_layout = TileLayout.from_iters(
|
||||
new_shard, new_replicate, offset=buffer.layout.offset
|
||||
)
|
||||
else:
|
||||
new_layout = buffer.layout
|
||||
buffer_attr_mutated = self.buffer_attr_var_mutated
|
||||
self.buffer_attr_var_mutated = prev_mutated or buffer_attr_mutated
|
||||
if not buffer_attr_mutated:
|
||||
return None
|
||||
new_buffer = decl_buffer(
|
||||
new_shape,
|
||||
buffer.dtype,
|
||||
buffer.name,
|
||||
new_data,
|
||||
new_strides,
|
||||
new_elem_offset,
|
||||
buffer.scope(),
|
||||
buffer.data_alignment,
|
||||
buffer.offset_factor,
|
||||
layout=new_layout,
|
||||
)
|
||||
self.buffer_map[buffer] = new_buffer
|
||||
return new_buffer
|
||||
|
||||
def visit_var_(self, op: Var):
|
||||
op = super().visit_var_(op)
|
||||
if op in self.var_map:
|
||||
self.buffer_attr_var_mutated = True
|
||||
return self.var_map[op]
|
||||
return op
|
||||
|
||||
def visit_buffer_load_(self, op: BufferLoad):
|
||||
new_buffer = self.mutate_buffer(op.buffer)
|
||||
op = super().visit_buffer_load_(op)
|
||||
if new_buffer is not None:
|
||||
return BufferLoad(new_buffer, op.indices)
|
||||
return op
|
||||
|
||||
def visit_buffer_store_(self, op: BufferStore):
|
||||
new_buffer = self.mutate_buffer(op.buffer)
|
||||
op = super().visit_buffer_store_(op)
|
||||
if new_buffer is not None:
|
||||
return BufferStore(new_buffer, op.value, op.indices)
|
||||
return op
|
||||
|
||||
def visit_buffer_region_(self, op: BufferRegion):
|
||||
new_buffer = self.mutate_buffer(op.buffer)
|
||||
op = super().visit_buffer_region_(op)
|
||||
if new_buffer is not None:
|
||||
return BufferRegion(new_buffer, op.region)
|
||||
return op
|
||||
|
||||
def visit_decl_buffer_(self, op: DeclBuffer):
|
||||
new_buffer = self.mutate_buffer(op.buffer)
|
||||
op = super().visit_decl_buffer_(op)
|
||||
if new_buffer is not None:
|
||||
return DeclBuffer(new_buffer, op.span)
|
||||
return op
|
||||
|
||||
def visit_array_prim_expr_(self, op: list[Expr]):
|
||||
return [self.visit_expr(expr) for expr in op]
|
||||
|
||||
def visit_alloc_buffer_(self, op: AllocBuffer):
|
||||
op = super().visit_alloc_buffer_(op)
|
||||
if op.buffer in self.buffer_map:
|
||||
return AllocBuffer(self.buffer_map[op.buffer], op.annotations, op.span)
|
||||
return op
|
||||
|
||||
def visit_op_call_(self, op):
|
||||
op = super().visit_op_call_(op)
|
||||
new_workspace = {}
|
||||
for key, value in op.workspace.items():
|
||||
new_buffer = self.mutate_buffer(value)
|
||||
if new_buffer is not None:
|
||||
new_workspace[key] = new_buffer
|
||||
else:
|
||||
new_workspace[key] = value
|
||||
new_config = {}
|
||||
for key, value in op.config.items():
|
||||
if is_prim_expr(value):
|
||||
new_config[key] = self.visit_expr(value)
|
||||
else:
|
||||
new_config[key] = value
|
||||
args = list()
|
||||
for arg in op.args:
|
||||
args.append(arg)
|
||||
return TilePrimitiveCall(
|
||||
*args,
|
||||
op=op.op,
|
||||
workspace=new_workspace,
|
||||
config=new_config,
|
||||
dispatch=op.dispatch,
|
||||
scope=op.scope,
|
||||
)
|
||||
|
||||
|
||||
class KernelReplacePointSearcher(StmtMutator):
|
||||
def __init__(self, body: Stmt):
|
||||
super().__init__()
|
||||
self.body = body
|
||||
|
||||
def visit_evaluate_(self, op: Evaluate):
|
||||
value = op.value
|
||||
if isinstance(value, Call) and value.op.same_as(Op.get("tirx.tvm_kernel_replace_point")):
|
||||
return self.body
|
||||
return super().visit_evaluate_(op)
|
||||
|
||||
|
||||
def seek_kernel_replace_point(stmt: Stmt, body: Stmt) -> Stmt:
|
||||
"""replace kernel replace point in stmt with body"""
|
||||
return KernelReplacePointSearcher(body)(stmt)
|
||||
Reference in New Issue
Block a user