chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,22 @@
|
||||
# isort: skip_file
|
||||
# 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.
|
||||
"""Namespace of all TIR transformations"""
|
||||
# pylint: disable=wildcard-import, invalid-name
|
||||
|
||||
from .function_pass import prim_func_pass, PrimFuncPass
|
||||
from .transform import *
|
||||
@@ -0,0 +1,21 @@
|
||||
# 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.
|
||||
"""FFI APIs for tvm.tirx.transform"""
|
||||
|
||||
import tvm_ffi
|
||||
|
||||
tvm_ffi.init_ffi_api("tirx.transform", __name__)
|
||||
@@ -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)
|
||||
@@ -0,0 +1,163 @@
|
||||
# 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.
|
||||
"""TIR specific function pass support."""
|
||||
|
||||
import functools
|
||||
import inspect
|
||||
from collections.abc import Callable
|
||||
|
||||
import tvm_ffi
|
||||
|
||||
from tvm.ir.transform import Pass, PassInfo
|
||||
|
||||
from . import _ffi_api
|
||||
|
||||
|
||||
@tvm_ffi.register_object("tirx.PrimFuncPass")
|
||||
class PrimFuncPass(Pass):
|
||||
"""A pass that works on each :py:func:`tvm.tirx.PrimFunc` in a module. A function
|
||||
pass class should be created through py:func:`tvm.tirx.transform.function_pass`.
|
||||
"""
|
||||
|
||||
|
||||
def _wrap_class_function_pass(pass_cls, pass_info):
|
||||
"""Wrap a python class as function pass"""
|
||||
|
||||
class PyFunctionPass(PrimFuncPass):
|
||||
"""Internal wrapper class to create a class instance."""
|
||||
|
||||
def __init__(self, *args, **kwargs):
|
||||
inst = pass_cls(*args, **kwargs)
|
||||
|
||||
# it is important not to capture self to
|
||||
# avoid a cyclic dependency
|
||||
def _pass_func(func, mod, ctx):
|
||||
return inst.transform_function(func, mod, ctx)
|
||||
|
||||
self.__init_handle_by_constructor__(
|
||||
_ffi_api.CreatePrimFuncPass,
|
||||
_pass_func,
|
||||
pass_info, # type: ignore
|
||||
)
|
||||
|
||||
self._inst = inst
|
||||
|
||||
def __getattr__(self, name):
|
||||
# fall back to instance attribute if there is not any
|
||||
return self._inst.__getattribute__(name)
|
||||
|
||||
functools.update_wrapper(PyFunctionPass.__init__, pass_cls.__init__)
|
||||
PyFunctionPass.__name__ = pass_cls.__name__
|
||||
PyFunctionPass.__doc__ = pass_cls.__doc__
|
||||
PyFunctionPass.__module__ = pass_cls.__module__
|
||||
return PyFunctionPass
|
||||
|
||||
|
||||
def prim_func_pass(
|
||||
pass_func=None,
|
||||
opt_level: int | None = None,
|
||||
name: str | None = None,
|
||||
required: list[str] | None = None,
|
||||
traceable=False,
|
||||
) -> Callable | PrimFuncPass:
|
||||
"""Decorate a function pass.
|
||||
|
||||
This function returns a callback when pass_func
|
||||
is provided. Otherwise, it returns the created function pass using the
|
||||
given optimization function.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
pass_func : Optional[Callable[(tvm.tirx.PrimFunc, IRModule, PassContext) -> tvm.tirx.PrimFunc]]
|
||||
The transformation function or class.
|
||||
|
||||
opt_level : int
|
||||
The optimization level of this module pass.
|
||||
|
||||
name : Optional[str]
|
||||
The name of the function pass. The name could be empty. In this case, the
|
||||
name of the optimization function will be used as the pass name.
|
||||
|
||||
required : Optional[List[str]]
|
||||
The list of passes that the function pass is dependent on.
|
||||
|
||||
Returns
|
||||
-------
|
||||
create_function_pass : Union[Callable, FunctionPass]
|
||||
|
||||
A decorator will be returned if pass_func is not provided,
|
||||
otherwise return the decorated result.
|
||||
The returned decorator has two behaviors depending on the input:
|
||||
A new FunctionPass will be returned when we decorate a pass function.
|
||||
A new FunctionPass class will be returned when we decorate a class type.
|
||||
|
||||
Examples
|
||||
--------
|
||||
The following code block decorates a function pass class.
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
@tvm.tirx.transform.prim_func_pass(opt_level=1)
|
||||
class TestReplaceFunc:
|
||||
def __init__(self, new_func):
|
||||
self.new_func = new_func
|
||||
|
||||
def transform_function(self, func, mod, ctx):
|
||||
# just for demo purposes
|
||||
# transform func to new_func
|
||||
return self.new_func
|
||||
|
||||
The following code creates a function pass by decorating
|
||||
a user defined transform function.
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
@tvm.tirx.transform.prim_func_pass(opt_level=2)
|
||||
def transform(func, mod, ctx):
|
||||
# my transformations here.
|
||||
return func
|
||||
|
||||
function_pass = transform
|
||||
assert isinstance(function_pass, transform.FunctionPass)
|
||||
assert function_pass.info.opt_level == 2
|
||||
|
||||
# Given a module m, the optimization could be invoked as the following:
|
||||
updated_mod = function_pass(m)
|
||||
# Now constant folding should have been applied to every function in
|
||||
# the provided module m. And the updated module will be returned.
|
||||
"""
|
||||
|
||||
if opt_level is None:
|
||||
raise ValueError("Please provide opt_level for the function pass.")
|
||||
|
||||
required = required if required else []
|
||||
if not isinstance(required, list | tuple):
|
||||
raise TypeError("Required is expected to be the type of " + "list/tuple.")
|
||||
|
||||
def create_function_pass(pass_arg):
|
||||
"""Internal function that creates a function pass"""
|
||||
fname = name if name else pass_arg.__name__
|
||||
info = PassInfo(opt_level, fname, required, traceable)
|
||||
if inspect.isclass(pass_arg):
|
||||
return _wrap_class_function_pass(pass_arg, info)
|
||||
if not callable(pass_arg):
|
||||
raise TypeError("pass_func must be a callable for Module pass")
|
||||
return _ffi_api.CreatePrimFuncPass(pass_arg, info) # type: ignore
|
||||
|
||||
if pass_func:
|
||||
return create_function_pass(pass_func)
|
||||
return create_function_pass
|
||||
@@ -0,0 +1,528 @@
|
||||
# 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.
|
||||
"""Wrapping existing transformations."""
|
||||
# pylint: disable=invalid-name, unsupported-binary-operation
|
||||
|
||||
import enum
|
||||
from collections.abc import Callable
|
||||
|
||||
import tvm_ffi as _ffi
|
||||
|
||||
from . import _ffi_api
|
||||
from . import function_pass as _fpass
|
||||
|
||||
|
||||
def Apply(ftransform):
|
||||
"""Apply ftransform to each function in the Module.
|
||||
|
||||
This function is a thin wrapper around tvm.tirx.transform.prim_func_pass
|
||||
|
||||
Parameters
|
||||
----------
|
||||
ftransform: tvm.tirx.PrimFunc -> tvm.tirx.PrimFunc
|
||||
The transformation pass.
|
||||
|
||||
Returns
|
||||
-------
|
||||
fpass : tvm.transform.Pass
|
||||
The result pass
|
||||
"""
|
||||
|
||||
# pylint: disable=unused-argument
|
||||
def _transform(func, mod, ctx):
|
||||
return ftransform(func)
|
||||
|
||||
return _fpass.prim_func_pass(_transform, opt_level=0, name="Apply") # type: ignore
|
||||
|
||||
|
||||
def VectorizeLoop(enable_vectorize: bool = True):
|
||||
"""Lower vectorization loops.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
enable_vectorize : bool
|
||||
Whether vectorization is enabled.
|
||||
Will lower to scalar loop when it is turned off.
|
||||
|
||||
Returns
|
||||
-------
|
||||
fpass : tvm.transform.Pass
|
||||
The result pass
|
||||
"""
|
||||
return _ffi_api.VectorizeLoop(enable_vectorize) # type: ignore
|
||||
|
||||
|
||||
def StorageRewrite():
|
||||
"""Rewrite storage allocation pattern.
|
||||
|
||||
Moves the allocation to outer most possible scope.
|
||||
Trying to share space between allocations to make
|
||||
a static allocation plan when possible.
|
||||
|
||||
Returns
|
||||
-------
|
||||
fpass : tvm.transform.Pass
|
||||
The result pass
|
||||
"""
|
||||
return _ffi_api.StorageRewrite() # type: ignore
|
||||
|
||||
|
||||
def InlinePrivateFunctions():
|
||||
"""Inline calls to private functions
|
||||
|
||||
Returns
|
||||
-------
|
||||
fpass : tvm.transform.Pass
|
||||
The result pass
|
||||
"""
|
||||
return _ffi_api.InlinePrivateFunctions() # type: ignore
|
||||
|
||||
|
||||
def PointerValueTypeRewrite():
|
||||
"""
|
||||
Rewrite the pointer content type of arguments, as well as Alloc internal to the function to use
|
||||
the most frequently accessed type for load/store to avoid pointer casting in backend when
|
||||
possible.
|
||||
|
||||
Returns
|
||||
-------
|
||||
fpass : tvm.transform.Pass
|
||||
The result pass
|
||||
"""
|
||||
return _ffi_api.PointerValueTypeRewrite() # type: ignore
|
||||
|
||||
|
||||
@_ffi.register_object("tirx.transform.UnrollLoopConfig")
|
||||
class UnrollLoopConfig(_ffi.Object):
|
||||
"""Config for unroll loop pass"""
|
||||
|
||||
|
||||
def UnrollLoop():
|
||||
"""Unroll the constant loop marked by unroll.
|
||||
|
||||
This pass also automatically attach pragma unroll tag to loops which meets the standard.
|
||||
|
||||
Returns
|
||||
-------
|
||||
fpass : tvm.transform.Pass
|
||||
The result pass
|
||||
"""
|
||||
return _ffi_api.UnrollLoop() # type: ignore
|
||||
|
||||
|
||||
@_ffi.register_object("tirx.transform.RemoveNoOpConfig")
|
||||
class RemoveNoOpConfig(_ffi.Object):
|
||||
"""Config for remove no op pass"""
|
||||
|
||||
|
||||
def RemoveNoOp():
|
||||
"""Remove No Op from the Stmt.
|
||||
|
||||
Returns
|
||||
-------
|
||||
fpass : tvm.transform.Pass
|
||||
The result pass
|
||||
"""
|
||||
return _ffi_api.RemoveNoOp() # type: ignore
|
||||
|
||||
|
||||
def RemoveAssume():
|
||||
"""Remove all instances of builtin::assume
|
||||
|
||||
Returns
|
||||
-------
|
||||
fpass : tvm.transform.Pass
|
||||
The result pass
|
||||
"""
|
||||
return _ffi_api.RemoveAssume() # type: ignore
|
||||
|
||||
|
||||
def BF16ComputeLegalize():
|
||||
"""Legalize bf16 compute Ops.
|
||||
|
||||
Returns
|
||||
-------
|
||||
fpass : tvm.transform.Pass
|
||||
The result pass
|
||||
"""
|
||||
return _ffi_api.BF16ComputeLegalize() # type: ignore
|
||||
|
||||
|
||||
def FP8ComputeLegalize(promote_dtype: str = "float32"):
|
||||
"""Legalize fp8 compute Ops.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
promote_dtype : str
|
||||
The data type we promote fp8 to, options: float16/float32.
|
||||
|
||||
Returns
|
||||
-------
|
||||
fpass : tvm.transform.Pass
|
||||
The result pass
|
||||
"""
|
||||
return _ffi_api.FP8ComputeLegalize(promote_dtype) # type: ignore
|
||||
|
||||
|
||||
def BF16StorageLegalize():
|
||||
"""Legalize bf16 storage types to u16.
|
||||
|
||||
Returns
|
||||
-------
|
||||
fpass : tvm.transform.Pass
|
||||
The result pass
|
||||
"""
|
||||
return _ffi_api.BF16StorageLegalize() # type: ignore
|
||||
|
||||
|
||||
def FP8StorageLegalize():
|
||||
"""Legalize fp8 storage types to u8.
|
||||
|
||||
Returns
|
||||
-------
|
||||
fpass : tvm.transform.Pass
|
||||
The result pass
|
||||
"""
|
||||
return _ffi_api.FP8StorageLegalize() # type: ignore
|
||||
|
||||
|
||||
def CommonSubexprElim():
|
||||
"""Replace redundant computations by new variables.
|
||||
|
||||
Returns
|
||||
-------
|
||||
fpass : tvm.transform.Pass
|
||||
The result pass
|
||||
"""
|
||||
return _ffi_api.CommonSubexprElim() # type: ignore
|
||||
|
||||
|
||||
@_ffi.register_object("tirx.transform.StmtSimplifyConfig")
|
||||
class StmtSimplifyConfig(_ffi.Object):
|
||||
"""Config for stmt simplify pass"""
|
||||
|
||||
|
||||
def StmtSimplify():
|
||||
"""Run statement-level arithmetic simplifications on the TIR PrimFunc.
|
||||
|
||||
Returns
|
||||
-------
|
||||
fpass : tvm.transform.Pass
|
||||
The result pass
|
||||
"""
|
||||
return _ffi_api.StmtSimplify() # type: ignore
|
||||
|
||||
|
||||
def ConvertSSA():
|
||||
"""Convert an IRModule to be SSA form.
|
||||
|
||||
This pass handles cases where the same `tirx.Var` appears in
|
||||
multiple functions within the same module. For example, after
|
||||
extracting a fragment from one function into another, where the
|
||||
same `tirx.Var` may be defined both as within the body of the
|
||||
original function, and as a parameter within the hoisted function.
|
||||
|
||||
Returns
|
||||
-------
|
||||
fpass : tvm.transform.Pass
|
||||
The result pass
|
||||
|
||||
"""
|
||||
return _ffi_api.ConvertSSA() # type: ignore
|
||||
|
||||
|
||||
def MakePackedAPI():
|
||||
"""Transform the PrimFuncs in the module to a packed func API.
|
||||
|
||||
Prior to this pass, the PrimFunc may have Buffer arguments defined
|
||||
in the `PrimFuncNode::buffer_map`. This pass consumes the
|
||||
`buffer_map`, using it to generate arguments that implement
|
||||
the packed based TVM FFI API.
|
||||
|
||||
For static shapes, the `BufferNode::shape`, `BufferNode::strides`,
|
||||
and `BufferNode::elem_offset` member variables are used to
|
||||
generate runtime checks on the corresponding member variables in
|
||||
the user-provided `DLTensor*` or `tvm.runtime.tensor` argument. (e.g. A
|
||||
PrimFunc that accepts a buffer of shape `[16,32]` validates that
|
||||
the `DLTensor::shape` array is `[16,32]`.)
|
||||
|
||||
For dynamic Buffers, in which one or more of these `BufferNode` member
|
||||
variables use `tirx.Var` that are not defined by other PrimFunc
|
||||
parameters, these are instead used to define the variables based on
|
||||
the corresponding `DLTensor` members. (e.g. A PrimFunc that accepts a
|
||||
buffer of shape `[tirx.Var("n"), tirx.Var("m")]`, when passed a
|
||||
`DLTensor` of shape `[16,32]`, will define `n = 16` and `n=32`, based
|
||||
on the argument's shape.
|
||||
|
||||
Returns
|
||||
-------
|
||||
fpass : tvm.transform.Pass
|
||||
The result pass
|
||||
"""
|
||||
return _ffi_api.MakePackedAPI() # type: ignore
|
||||
|
||||
|
||||
def SplitHostDevice():
|
||||
"""Annotate, split, and lower host/device functions.
|
||||
|
||||
This pass first annotates device regions within host functions,
|
||||
then splits them into host and device-side PrimFuncs, and finally
|
||||
lowers host-to-device calls into the device kernel launch ABI.
|
||||
|
||||
Returns
|
||||
-------
|
||||
fpass : tvm.transform.Pass
|
||||
The result pass
|
||||
"""
|
||||
return _ffi_api.SplitHostDevice() # type: ignore
|
||||
|
||||
|
||||
def SkipAssert():
|
||||
"""Skip assert stmt.
|
||||
|
||||
Returns
|
||||
-------
|
||||
fpass : tvm.transform.Pass
|
||||
The result pass
|
||||
"""
|
||||
return _ffi_api.SkipAssert() # type: ignore
|
||||
|
||||
|
||||
def LowerWarpMemory():
|
||||
"""Lower warp memory access to low-level device related function calls.
|
||||
|
||||
Returns
|
||||
-------
|
||||
fpass : tvm.transform.Pass
|
||||
The result pass
|
||||
"""
|
||||
return _ffi_api.LowerWarpMemory() # type: ignore
|
||||
|
||||
|
||||
def LowerTVMBuiltin():
|
||||
"""Lower tvm builtin intrinsics.
|
||||
|
||||
Returns
|
||||
-------
|
||||
fpass : tvm.transform.Pass
|
||||
The result pass
|
||||
"""
|
||||
return _ffi_api.LowerTVMBuiltin() # type: ignore
|
||||
|
||||
|
||||
def LowerIntrin():
|
||||
"""Lower target specific intrinsic calls.
|
||||
|
||||
Returns
|
||||
-------
|
||||
fpass : tvm.transform.Pass
|
||||
The result pass
|
||||
"""
|
||||
return _ffi_api.LowerIntrin() # type: ignore
|
||||
|
||||
|
||||
def NarrowDataType(target_bits: int):
|
||||
"""Narrow down Expr datatype in stmt to target_bits.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
target_bits : int
|
||||
The target bit configuration.
|
||||
|
||||
Returns
|
||||
-------
|
||||
fpass : tvm.transform.Pass
|
||||
The result pass
|
||||
|
||||
Note
|
||||
----
|
||||
Run this pass after FlattenBuffer.
|
||||
"""
|
||||
return _ffi_api.NarrowDataType(target_bits) # type: ignore
|
||||
|
||||
|
||||
def ForceNarrowIndexToInt32():
|
||||
"""Force narrow down indexing expressions and integer buffers to int32 dtype.
|
||||
|
||||
Returns
|
||||
-------
|
||||
fpass : tvm.transform.Pass
|
||||
The result pass
|
||||
|
||||
Note
|
||||
----
|
||||
This pass should not be used in default cases.
|
||||
"""
|
||||
return _ffi_api.ForceNarrowIndexToInt32() # type: ignore
|
||||
|
||||
|
||||
def VerifyMemory():
|
||||
"""Verify if func contains illegal host side direct memory access.
|
||||
|
||||
Returns
|
||||
-------
|
||||
fpass : tvm.transform.Pass
|
||||
The result pass
|
||||
"""
|
||||
return _ffi_api.VerifyMemory() # type: ignore
|
||||
|
||||
|
||||
@_ffi.register_object("s_tir.transform.HoistIfThenElseConfig")
|
||||
class HoistIfThenElseConfig(_ffi.Object):
|
||||
"""Config for hoist if then else pass"""
|
||||
|
||||
|
||||
class HoistedConditionals(enum.Flag):
|
||||
"""Flags for use in HoistExpressionConfig.conditional_types
|
||||
|
||||
Each bitflag represents a type of expression that should be
|
||||
hoisted to the outermost loop possible.
|
||||
"""
|
||||
|
||||
Never = 0
|
||||
""" No hoisting of conditionals """
|
||||
|
||||
IfElseStmt = 1
|
||||
""" If set, look for hoist candidates in IfElseStmt """
|
||||
|
||||
IfElseExpr = 2
|
||||
""" If set, look for hoist candidates in tirx.if_then_else """
|
||||
|
||||
BooleanExpression = 4
|
||||
""" If set, look for hoist candidates in all boolean expressions """
|
||||
|
||||
UsingBlockVar = 8
|
||||
""" If set, allow hoisting of conditionals that use a block variable (e.g. threadIdx.x) """
|
||||
|
||||
All = IfElseStmt | IfElseExpr | BooleanExpression | UsingBlockVar
|
||||
""" Enable all hoisting of conditionals"""
|
||||
|
||||
|
||||
class HoistedLetBindings(enum.Flag):
|
||||
"""Flags for use in HoistExpressionConfig.let_binding_types
|
||||
|
||||
Each bitflag represents a type of let binding expression that should be
|
||||
hoisted to the outermost loop possible.
|
||||
"""
|
||||
|
||||
Never = 0
|
||||
""" No hoisting of let bindings """
|
||||
|
||||
RequiredByConditional = 1
|
||||
""" Bindings that are used by a hoisted conditional """
|
||||
|
||||
Bind = 2
|
||||
""" Bindings occurring in Bind nodes """
|
||||
|
||||
LetExpr = 4
|
||||
""" Bindings occurring in Let expressions """
|
||||
|
||||
All = RequiredByConditional | Bind | LetExpr
|
||||
""" Enable all hoisting of let bindings """
|
||||
|
||||
|
||||
@_ffi.register_object("s_tir.transform.HoistExpressionConfig")
|
||||
class HoistExpressionConfig(_ffi.Object):
|
||||
"""Config for hoist expression pass"""
|
||||
|
||||
|
||||
def FlattenBuffer():
|
||||
"""Flatten the multi-dimensional BufferLoad and BufferStore to single dimensional
|
||||
BufferLoad/BufferStore for the TIR not contains opaque block.
|
||||
|
||||
Returns
|
||||
-------
|
||||
fpass : tvm.transform.Pass
|
||||
The result pass
|
||||
"""
|
||||
return _ffi_api.FlattenBuffer() # type: ignore
|
||||
|
||||
|
||||
def BindTarget(target):
|
||||
"""Annotate a PrimFunc with a given target.
|
||||
Parameters
|
||||
-------
|
||||
target : tvm.target.Target
|
||||
target
|
||||
|
||||
Returns
|
||||
-------
|
||||
fpass : tvm.transform.Pass
|
||||
The result pass
|
||||
"""
|
||||
return _ffi_api.BindTarget(target) # type: ignore
|
||||
|
||||
|
||||
def AnnotateEntryFunc():
|
||||
"""Set a PrimFunc as the entry point if it is only function in IRModule.
|
||||
|
||||
Returns
|
||||
-------
|
||||
fpass : tvm.transform.Pass
|
||||
The result pass
|
||||
"""
|
||||
return _ffi_api.AnnotateEntryFunc() # type: ignore
|
||||
|
||||
|
||||
def Filter(fcond: Callable):
|
||||
"""Filter out PrimFuncs that does not satisfy the given condition.
|
||||
`fcond` should be a function that takes a primfunc and returns boolean.
|
||||
|
||||
Returns
|
||||
-------
|
||||
fpass : tvm.transform.Pass
|
||||
The result pass
|
||||
"""
|
||||
return _ffi_api.Filter(fcond) # type: ignore
|
||||
|
||||
|
||||
def TilePrimitiveDispatch():
|
||||
"""Lower TIRx tile primitive calls through the active backend dispatch table.
|
||||
|
||||
Returns
|
||||
-------
|
||||
fpass : tvm.transform.Pass
|
||||
The result pass
|
||||
"""
|
||||
return _ffi_api.TilePrimitiveDispatch() # type: ignore
|
||||
|
||||
|
||||
def LowerTIRx():
|
||||
"""Lower TIR to a lower-level IR.
|
||||
|
||||
Returns
|
||||
-------
|
||||
fpass : tvm.transform.Pass
|
||||
The result pass
|
||||
"""
|
||||
return _ffi_api.LowerTIRx() # type: ignore
|
||||
|
||||
|
||||
def LowerTIRxOpaque():
|
||||
"""Lower opaque constructs in TIRX programs.
|
||||
|
||||
Handles AllocBuffer lowering, For(thread_binding) to AttrStmt(thread_extent)
|
||||
conversion, unit loop elimination, and pragma annotation handling.
|
||||
This is the tirx-specific counterpart of s_tir.LowerOpaqueBlock,
|
||||
without any SBlock/SBlockRealize handling.
|
||||
|
||||
Returns
|
||||
-------
|
||||
fpass : tvm.transform.Pass
|
||||
The result pass
|
||||
"""
|
||||
return _ffi_api.LowerTIRxOpaque() # type: ignore
|
||||
Reference in New Issue
Block a user