chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,595 @@
|
||||
# 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.
|
||||
"""Operation class for computation declaration."""
|
||||
|
||||
import inspect
|
||||
|
||||
# pylint: disable=invalid-name
|
||||
from numbers import Integral as _Integral
|
||||
|
||||
from tvm_ffi import Array
|
||||
|
||||
import tvm.arith._ffi_api
|
||||
import tvm.tirx
|
||||
import tvm.tirx._ffi_api
|
||||
from tvm.ir import is_prim_expr
|
||||
from tvm.runtime import convert
|
||||
|
||||
from . import _ffi_api
|
||||
from . import tag as _tag
|
||||
from . import tensor as _tensor
|
||||
|
||||
|
||||
def placeholder(shape, dtype=None, name="placeholder"):
|
||||
"""Construct an empty tensor object.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
shape: Tuple of Expr
|
||||
The shape of the tensor
|
||||
|
||||
dtype: str, optional
|
||||
The data type of the tensor
|
||||
|
||||
name: str, optional
|
||||
The name hint of the tensor
|
||||
|
||||
Returns
|
||||
-------
|
||||
tensor: Tensor
|
||||
The created tensor
|
||||
"""
|
||||
dtype = "float32" if dtype is None else dtype
|
||||
return _ffi_api.Placeholder(shape, dtype, name)
|
||||
|
||||
|
||||
def compute(shape, fcompute, name="compute", tag="", attrs=None, varargs_names=None):
|
||||
"""Construct a new tensor by computing over the shape domain.
|
||||
|
||||
The compute rule is result[axis] = fcompute(axis)
|
||||
|
||||
Parameters
|
||||
----------
|
||||
shape: Tuple of Expr
|
||||
The shape of the tensor
|
||||
|
||||
fcompute: lambda function of indices-> value
|
||||
Specifies the input source expression
|
||||
|
||||
name: str, optional
|
||||
The name hint of the tensor
|
||||
|
||||
tag: str, optional
|
||||
Additional tag information about the compute.
|
||||
|
||||
attrs: dict, optional
|
||||
The additional auxiliary attributes about the compute.
|
||||
|
||||
varargs_names: list, optional
|
||||
The names to use for each of the varargs. If not supplied, the varargs
|
||||
will be called i1, i2, ...
|
||||
|
||||
Returns
|
||||
-------
|
||||
tensor: Tensor
|
||||
The created tensor
|
||||
"""
|
||||
if _tag.TagScope.get_current() is not None:
|
||||
if tag != "":
|
||||
raise ValueError("nested tag is not allowed for now")
|
||||
tag = _tag.TagScope.get_current().tag
|
||||
shape = (shape,) if tvm.ir.is_prim_expr(shape) else shape
|
||||
# for python3
|
||||
shape = tuple([int(s) if isinstance(s, float) else s for s in shape])
|
||||
out_ndim = len(shape)
|
||||
|
||||
argspec = inspect.getfullargspec(fcompute)
|
||||
if len(argspec.args) == 0 and argspec.varargs is None:
|
||||
arg_names = [f"i{i}" for i in range(out_ndim)]
|
||||
elif argspec.varargs is not None:
|
||||
# if there is a varargs, it takes the remaining dimensions of out_ndim
|
||||
num_remaining_args = out_ndim - len(argspec.args)
|
||||
if varargs_names is not None:
|
||||
if len(varargs_names) != num_remaining_args:
|
||||
raise RuntimeError(
|
||||
f"Number of varargs ({num_remaining_args}) does not match number"
|
||||
f"of varargs_names ({len(varargs_names)})"
|
||||
)
|
||||
arg_names = argspec.args + varargs_names
|
||||
else:
|
||||
arg_names = argspec.args + [f"i{i}" for i in range(out_ndim - len(argspec.args))]
|
||||
else:
|
||||
arg_names = argspec.args
|
||||
# if there are fewer args than out dimensions, the remaining dimensions
|
||||
# are implicitly broadcast
|
||||
out_ndim = len(arg_names)
|
||||
assert argspec.varkw is None, "Variable keyword arguments not supported in fcompute"
|
||||
assert argspec.defaults is None, "Default arguments not supported in fcompute"
|
||||
assert len(argspec.kwonlyargs) == 0, "Keyword arguments are not supported in fcompute"
|
||||
|
||||
if out_ndim != len(arg_names):
|
||||
raise ValueError(
|
||||
"Number of args to fcompute does not match dimension, "
|
||||
f"args={len(arg_names)}, dimension={out_ndim}"
|
||||
)
|
||||
|
||||
dim_var = [tvm.tirx.IterVar((0, s), x, 0) for x, s in zip(arg_names, shape[:out_ndim])]
|
||||
body = fcompute(*[v.var for v in dim_var])
|
||||
|
||||
if not isinstance(body, list | tuple):
|
||||
body = [body]
|
||||
body = convert(body)
|
||||
op_node = _ffi_api.ComputeOp(name, tag, attrs, dim_var, body)
|
||||
|
||||
num = op_node.num_outputs
|
||||
outputs = tuple(op_node.output(i) for i in range(num))
|
||||
return outputs[0] if num == 1 else outputs
|
||||
|
||||
|
||||
def scan(init, update, state_placeholder, inputs=None, name="scan", tag="", attrs=None):
|
||||
"""Construct new tensors by scanning over axis.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
init: Tensor or list of Tensor
|
||||
The initial condition of first init.shape[0] timestamps
|
||||
|
||||
update: Tensor or list of Tensor
|
||||
The update rule of the scan given by symbolic tensor.
|
||||
|
||||
state_placeholder: Tensor or list of Tensor
|
||||
The placeholder variables used by update.
|
||||
|
||||
inputs: Tensor or list of Tensor, optional
|
||||
The list of inputs to the scan. This is not required, but can
|
||||
be useful for the compiler to detect scan body faster.
|
||||
|
||||
name: str, optional
|
||||
The name hint of the tensor
|
||||
|
||||
tag: str, optional
|
||||
Additonal tag information about the compute.
|
||||
|
||||
attrs: dict, optional
|
||||
The additional auxiliary attributes about the compute.
|
||||
|
||||
Returns
|
||||
-------
|
||||
tensor: Tensor or list of Tensors
|
||||
The created tensor or tuple of tensors contains multiple outputs.
|
||||
|
||||
Example
|
||||
-------
|
||||
.. code-block:: python
|
||||
|
||||
# The following code is equivalent to numpy.cumsum
|
||||
m = te.var("m")
|
||||
n = te.var("n")
|
||||
X = te.placeholder((m, n), name="X")
|
||||
s_state = te.placeholder((m, n))
|
||||
s_init = te.compute((1, n), lambda _, i: X[0, i])
|
||||
s_update = te.compute((m, n), lambda t, i: s_state[t-1, i] + X[t, i])
|
||||
res = tvm.te.scan(s_init, s_update, s_state, X)
|
||||
"""
|
||||
if _tag.TagScope.get_current() is not None:
|
||||
if tag != "":
|
||||
raise ValueError("nested tag is not allowed for now")
|
||||
tag = _tag.TagScope.get_current().tag
|
||||
if isinstance(init, _tensor.Tensor):
|
||||
init = [init]
|
||||
if isinstance(update, _tensor.Tensor):
|
||||
update = [update]
|
||||
if isinstance(state_placeholder, _tensor.Tensor):
|
||||
state_placeholder = [state_placeholder]
|
||||
if isinstance(inputs, _tensor.Tensor):
|
||||
inputs = [inputs]
|
||||
if inputs is None:
|
||||
inputs = []
|
||||
if len(init) != len(update) or len(init) != len(state_placeholder):
|
||||
raise ValueError("init, update, state_placeholder must have same length")
|
||||
axis = tvm.tirx.IterVar((init[0].shape[0], update[0].shape[0]), f"{name}.idx", 3)
|
||||
op = _ffi_api.ScanOp(name, tag, attrs, axis, init, update, state_placeholder, inputs)
|
||||
res = [op.output(i) for i in range(len(update))]
|
||||
return res[0] if len(res) == 1 else res
|
||||
|
||||
|
||||
def extern(
|
||||
shape,
|
||||
inputs,
|
||||
fcompute,
|
||||
name="extern",
|
||||
dtype=None,
|
||||
in_buffers=None,
|
||||
out_buffers=None,
|
||||
tag="",
|
||||
attrs=None,
|
||||
):
|
||||
"""Compute several tensors via an extern function.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
shape: tuple or list of tuples.
|
||||
The shape of the outputs.
|
||||
|
||||
inputs: list of Tensor
|
||||
The inputs
|
||||
|
||||
fcompute: lambda function of inputs, outputs-> stmt
|
||||
Specifies the IR statement to do the computation.
|
||||
See the following note for function signature of fcompute
|
||||
|
||||
.. note::
|
||||
**Parameters**
|
||||
|
||||
- **ins** (list of :any:`tvm.tirx.Buffer`) - Placeholder for each inputs
|
||||
- **outs** (list of :any:`tvm.tirx.Buffer`) - Placeholder for each outputs
|
||||
|
||||
**Returns**
|
||||
|
||||
- **stmt** (:any:`tvm.tirx.Stmt`) - The statement that carries out array computation.
|
||||
|
||||
name: str, optional
|
||||
The name hint of the tensor
|
||||
|
||||
dtype: str or list of str, optional
|
||||
The data types of outputs,
|
||||
by default dtype will be same as inputs.
|
||||
|
||||
in_buffers: tvm.tirx.Buffer or list of tvm.tirx.Buffer, optional
|
||||
Input buffers.
|
||||
|
||||
out_buffers: tvm.tirx.Buffer or list of tvm.tirx.Buffer, optional
|
||||
Output buffers.
|
||||
|
||||
|
||||
tag: str, optional
|
||||
Additonal tag information about the compute.
|
||||
|
||||
attrs: dict, optional
|
||||
The additional auxiliary attributes about the compute.
|
||||
|
||||
Returns
|
||||
-------
|
||||
tensor: Tensor or list of Tensors
|
||||
The created tensor or tuple of tensors contains multiple outputs.
|
||||
|
||||
Example
|
||||
-------
|
||||
In the code below, C is generated by calling external PackedFunc
|
||||
`tvm.contrib.cblas.matmul`
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
A = te.placeholder((n, l), name="A")
|
||||
B = te.placeholder((l, m), name="B")
|
||||
C = te.extern((n, m), [A, B],
|
||||
lambda ins, outs: tvm.tirx.call_packed(
|
||||
"tvm.contrib.cblas.matmul",
|
||||
ins[0], ins[1], outs[0], 0, 0), name="C")
|
||||
"""
|
||||
if _tag.TagScope.get_current() is not None:
|
||||
if tag != "":
|
||||
raise ValueError("nested tag is not allowed for now")
|
||||
tag = _tag.TagScope.get_current().tag
|
||||
shape = (shape,) if is_prim_expr(shape) or isinstance(shape, _Integral) else shape
|
||||
if shape == () or is_prim_expr(shape[0]) or isinstance(shape[0], _Integral):
|
||||
shape = [shape]
|
||||
if in_buffers is not None:
|
||||
in_buffers = [in_buffers] if not isinstance(in_buffers, list) else in_buffers
|
||||
if len(inputs) != len(in_buffers):
|
||||
raise RuntimeError(
|
||||
f"Number of inputs and in_buffers mismatch: {len(inputs)} vs {len(in_buffers)}."
|
||||
)
|
||||
if out_buffers is not None:
|
||||
out_buffers = [out_buffers] if not isinstance(out_buffers, list) else out_buffers
|
||||
if len(shape) != len(out_buffers):
|
||||
raise RuntimeError(
|
||||
f"Number of outputs and out_buffers mismatch: {len(shape)} vs {len(out_buffers)}."
|
||||
)
|
||||
input_placeholders = in_buffers or []
|
||||
output_placeholders = out_buffers or []
|
||||
types = set()
|
||||
for t in inputs:
|
||||
if not isinstance(t, _tensor.Tensor):
|
||||
raise ValueError("expect inputs to be tensor")
|
||||
if in_buffers is None:
|
||||
input_placeholders.append(
|
||||
tvm.tirx.decl_buffer(
|
||||
t.shape,
|
||||
t.dtype,
|
||||
t.op.name,
|
||||
elem_offset=tvm.tirx.Var("elem_offset", "int32"),
|
||||
layout=None,
|
||||
)
|
||||
)
|
||||
types.add(t.dtype)
|
||||
|
||||
if out_buffers is None:
|
||||
if dtype is None:
|
||||
if len(types) != 1:
|
||||
raise ValueError("Cannot infer output type, please provide dtype argument")
|
||||
infered_type = types.pop()
|
||||
dtype = [infered_type for _ in shape]
|
||||
if isinstance(dtype, str):
|
||||
dtype = [dtype]
|
||||
|
||||
for shp, dt in zip(shape, dtype):
|
||||
output_placeholders.append(
|
||||
tvm.tirx.decl_buffer(
|
||||
shp,
|
||||
dt,
|
||||
name,
|
||||
elem_offset=tvm.tirx.Var("elem_offset", "int32"),
|
||||
layout=None,
|
||||
)
|
||||
)
|
||||
body = fcompute(input_placeholders, output_placeholders)
|
||||
if tvm.ir.is_prim_expr(body):
|
||||
body = tvm.tirx.Evaluate(body)
|
||||
if not isinstance(body, tvm.tirx.Stmt):
|
||||
raise ValueError(
|
||||
f"Function '{fcompute.__name__}' should return Expr or Stmt, but it returned "
|
||||
f"'{type(body)}'"
|
||||
)
|
||||
|
||||
op = _ffi_api.ExternOp(name, tag, attrs, inputs, input_placeholders, output_placeholders, body)
|
||||
res = [op.output(i) for i in range(len(output_placeholders))]
|
||||
return res[0] if len(res) == 1 else res
|
||||
|
||||
|
||||
def extern_primfunc(input_tensors: list[_tensor.Tensor], primfunc: tvm.tirx.PrimFunc, **kwargs):
|
||||
"""Compute tensors via a schedulable TIR PrimFunc
|
||||
|
||||
Parameters
|
||||
----------
|
||||
input_tensors: list of Tensor
|
||||
Input tensors that map to the corresponding primfunc input params.
|
||||
|
||||
primfunc: PrimFunc
|
||||
The TIR PrimFunc
|
||||
|
||||
Returns
|
||||
-------
|
||||
tensor: Tensor or list of Tensors
|
||||
The created tensor or tuple of tensors if it contains multiple outputs.
|
||||
|
||||
Example
|
||||
-------
|
||||
In the code below, a TVMScript defined TIR PrimFunc is inlined into
|
||||
a TE ExternOp. Applying te.create_prim_func on this
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
A = te.placeholder((128, 128), name="A")
|
||||
B = te.placeholder((128, 128), name="B")
|
||||
|
||||
@T.prim_func(s_tir=True)
|
||||
def before_split(a: T.handle, b: T.handle) -> None:
|
||||
A = T.match_buffer(a, (128, 128))
|
||||
B = T.match_buffer(b, (128, 128))
|
||||
for i, j in T.grid(128, 128):
|
||||
with T.sblock("B"):
|
||||
vi, vj = T.axis.remap("SS", [i, j])
|
||||
B[vi, vj] = A[vi, vj] * 2.0
|
||||
|
||||
C = te.extern_primfunc([A, B], func)
|
||||
"""
|
||||
|
||||
# dt_access_map and primfunc.buffer_map are unordered, so use order from primfunc.params
|
||||
dt_access_map = tvm.arith._ffi_api.DomainTouchedAccessMap(primfunc)
|
||||
ordered_buffers = [primfunc.buffer_map[param] for param in primfunc.params]
|
||||
in_buffers = [buf for buf in ordered_buffers if len(dt_access_map[buf][0])]
|
||||
out_buffers = [buf for buf in ordered_buffers if len(dt_access_map[buf][1])]
|
||||
assert in_buffers, "PrimFunc has no input buffers"
|
||||
assert out_buffers, "PrimFunc has no output buffers"
|
||||
|
||||
outputs = []
|
||||
inplace = []
|
||||
input_buffers = in_buffers
|
||||
for obuf in out_buffers:
|
||||
if obuf in in_buffers:
|
||||
inplace.append(obuf)
|
||||
else:
|
||||
outputs.append(obuf)
|
||||
|
||||
if not outputs:
|
||||
iobuf = inplace.pop()
|
||||
input_buffers.remove(iobuf)
|
||||
outputs = [iobuf]
|
||||
|
||||
assert len(input_buffers) == len(input_tensors), (
|
||||
"The number of provided input input_tensors does not match the number of ",
|
||||
"input buffers in the primfunc",
|
||||
)
|
||||
for tensor, buffer in zip(input_tensors, input_buffers):
|
||||
# TODO(csullivan): Can a stronger comparison between Tensor<>Buffer be made?
|
||||
assert len(tensor.shape) == len(buffer.shape)
|
||||
for d1, d2 in zip(tensor.shape, buffer.shape):
|
||||
assert d1 == d2, (
|
||||
"The input input_tensors provided do not match the input buffers in the ",
|
||||
"primfunc. Please check that the order of input te.Input_Tensors and the ",
|
||||
"order of the primfunc variables in the params list agree.",
|
||||
)
|
||||
output = extern(
|
||||
[buf.shape for buf in outputs],
|
||||
input_tensors,
|
||||
lambda ins, outs: primfunc.body,
|
||||
in_buffers=input_buffers,
|
||||
out_buffers=outputs,
|
||||
**kwargs,
|
||||
)
|
||||
return output
|
||||
|
||||
|
||||
def var(name="tindex", dtype="int32", span=None):
|
||||
"""Create a new variable with specified name and dtype
|
||||
|
||||
Parameters
|
||||
----------
|
||||
name : str
|
||||
The name
|
||||
|
||||
dtype : str
|
||||
The data type
|
||||
|
||||
span : Optional[Span]
|
||||
The location of this variable in the source.
|
||||
|
||||
Returns
|
||||
-------
|
||||
var : tirx.Var
|
||||
The result symbolic variable.
|
||||
"""
|
||||
return tvm.tirx.Var(name, dtype, span)
|
||||
|
||||
|
||||
def const(value, dtype="int32", span=None):
|
||||
"""Create a new constant with specified value and dtype
|
||||
|
||||
Parameters
|
||||
----------
|
||||
value : Union[bool, int, float, numpy.ndarray, tvm.runtime.Tensor]
|
||||
The constant value.
|
||||
|
||||
dtype : str
|
||||
The data type
|
||||
|
||||
span : Optional[Span]
|
||||
The location of this variable in the source.
|
||||
|
||||
Returns
|
||||
-------
|
||||
const : Expr
|
||||
The result constant expr.
|
||||
"""
|
||||
return tvm.tirx.const(value, dtype, span)
|
||||
|
||||
|
||||
def thread_axis(dom=None, tag="", name="", span=None):
|
||||
"""Create a new IterVar to represent thread index.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
dom : Range or str
|
||||
The domain of iteration
|
||||
When str is passed, dom is set to None and str is used as tag
|
||||
|
||||
tag : str, optional
|
||||
The thread tag
|
||||
|
||||
name : str, optional
|
||||
The name of the var.
|
||||
|
||||
span : Optional[Span]
|
||||
The location of this variable in the source.
|
||||
|
||||
Returns
|
||||
-------
|
||||
axis : IterVar
|
||||
The thread itervar.
|
||||
"""
|
||||
if isinstance(dom, str):
|
||||
tag, dom = dom, None
|
||||
if not tag:
|
||||
raise ValueError("tag must be given as Positional or keyword argument")
|
||||
name = name if name else tag
|
||||
return tvm.tirx.IterVar(dom, name, 1, tag, span)
|
||||
|
||||
|
||||
def reduce_axis(dom, name="rv", thread_tag="", span=None):
|
||||
"""Create a new IterVar for reduction.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
dom : Range
|
||||
The domain of iteration.
|
||||
|
||||
name : str
|
||||
The name of the variable.
|
||||
|
||||
thread_tag : Optional[str]
|
||||
The name of the thread_tag.
|
||||
|
||||
span : Optional[Span]
|
||||
The location of this variable in the source.
|
||||
|
||||
Returns
|
||||
-------
|
||||
axis : IterVar
|
||||
An iteration variable representing the value.
|
||||
"""
|
||||
return tvm.tirx.IterVar(dom, name, 2, thread_tag, span)
|
||||
|
||||
|
||||
def create_prim_func(
|
||||
ops: list[_tensor.Tensor | tvm.tirx.Var], index_dtype_override: str | None = None
|
||||
) -> tvm.tirx.PrimFunc:
|
||||
"""Create a TensorIR PrimFunc from tensor expression
|
||||
|
||||
Parameters
|
||||
----------
|
||||
ops : List[Union[_tensor.Tensor, tvm.tirx.Var]]
|
||||
The source expression.
|
||||
|
||||
Example
|
||||
-------
|
||||
We define a matmul kernel using following code:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
import tvm
|
||||
from tvm import te
|
||||
from tvm.te import create_prim_func
|
||||
import tvm.script
|
||||
|
||||
A = te.placeholder((128, 128), name="A")
|
||||
B = te.placeholder((128, 128), name="B")
|
||||
k = te.reduce_axis((0, 128), "k")
|
||||
C = te.compute((128, 128), lambda x, y: te.sum(A[x, k] * B[y, k], axis=k), name="C")
|
||||
func = create_prim_func([A, B, C])
|
||||
print(func.script())
|
||||
|
||||
If we want to use TensorIR schedule to do transformations on such kernel,
|
||||
we need to use `create_prim_func([A, B, C])` to create a schedulable PrimFunc.
|
||||
The generated function looks like:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
@T.prim_func(s_tir=True)
|
||||
def tir_matmul(a: T.handle, b: T.handle, c: T.handle) -> None:
|
||||
A = T.match_buffer(a, (128, 128))
|
||||
B = T.match_buffer(b, (128, 128))
|
||||
C = T.match_buffer(c, (128, 128))
|
||||
|
||||
for i, j, k in T.grid(128, 128, 128):
|
||||
with T.sblock():
|
||||
vi, vj, vk = T.axis.remap("SSR", [i, j, k])
|
||||
with T.init():
|
||||
C[vi, vj] = 0.0
|
||||
C[vi, vj] += A[vi, vk] * B[vj, vk]
|
||||
|
||||
Returns
|
||||
-------
|
||||
func : tirx.PrimFunc
|
||||
The created function.
|
||||
"""
|
||||
if not isinstance(ops, list | tuple | Array):
|
||||
ops = [ops]
|
||||
return _ffi_api.CreatePrimFunc(ops, index_dtype_override)
|
||||
|
||||
|
||||
AXIS_SEPARATOR = tvm.tirx.IndexMap.AXIS_SEPARATOR
|
||||
Reference in New Issue
Block a user