wehub/apache--tvm
1
0
Fork 0
Code Issues Pull Requests Actions Deployments Models Agent Notes Packages Projects Releases Wiki Activity

chore: import upstream snapshot with attribution
Lint / lint (push) Has been cancelled
Details
CI / MacOS (push) Has been cancelled
Details
CI / Windows (push) Has been cancelled
Details

Browse Source
This commit is contained in:
wehub-resource-sync
2026-07-13 13:36:25 +08:00
commit 26446540fa
3151 changed files with 974126 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/tvm/te/__init__.py
+37
View File
@@ -0,0 +1,37 @@
# 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.
# pylint: disable=unused-import, redefined-builtin, wildcard-import
"""Namespace for Tensor Expression Language"""
# expose all operators in tvm tirx.op
from tvm.tirx import any, all, min_value, max_value, trace
from tvm.tirx import exp, erf, tanh, sigmoid, log, tan, cos, sin, sqrt, rsqrt, floor, ceil
from tvm.tirx import sinh, cosh, log2, log10
from tvm.tirx import asin, asinh, acos, acosh, atan, atanh
from tvm.tirx import trunc, abs, round, nearbyint, power, popcount, fmod, if_then_else
from tvm.tirx import isnan, isfinite, isinf
from tvm.tirx import div, indexdiv, indexmod, truncdiv, truncmod, floordiv, floormod, logaddexp
from tvm.tirx import comm_reducer, min, max, sum
from .tensor import TensorSlice, Tensor
from .tag import tag_scope
from .operation import placeholder, compute, scan, extern, var, const
from .operation import thread_axis, reduce_axis, AXIS_SEPARATOR
from .operation import create_prim_func
from .operation import extern_primfunc
from .tensor import PlaceholderOp, ComputeOp, ScanOp, ExternOp
python/tvm/te/_ffi_api.py
+21
View File
@@ -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.te"""
import tvm_ffi
tvm_ffi.init_ffi_api("te", __name__)
python/tvm/te/_te_tensor_overload.py
+61
View File
@@ -0,0 +1,61 @@
# 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.
"""Tensor overload hooks for TE tensors and tensor slices."""
def __add__(_lhs, _rhs):
return NotImplemented
def __radd__(_lhs, _rhs):
return NotImplemented
def __sub__(_lhs, _rhs):
return NotImplemented
def __rsub__(_lhs, _rhs):
return NotImplemented
def __mul__(_lhs, _rhs):
return NotImplemented
def __rmul__(_lhs, _rhs):
return NotImplemented
def __div__(_lhs, _rhs):
return NotImplemented
def __rdiv__(_lhs, _rhs):
return NotImplemented
def __truediv__(_lhs, _rhs):
return NotImplemented
def __rtruediv__(_lhs, _rhs):
return NotImplemented
def astype(_value, _dtype, _span=None):
return NotImplemented
python/tvm/te/operation.py
+595
View File
@@ -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
python/tvm/te/tag.py
+96
View File
@@ -0,0 +1,96 @@
# 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.
"""Tag class for TVM operators."""
import functools
import warnings
class TagScope:
"""Tag scope object to set tag for operators, working as context
manager and decorator both. See also tag_scope.
"""
_current = None
@classmethod
def get_current(cls):
if cls._current:
cls._current.accessed = True
return cls._current
def __init__(self, tag):
self._old_scope = None
self.tag = tag
self.accessed = False
def __enter__(self):
if TagScope._current is not None:
raise ValueError("nested op_tag is not allowed for now")
self._old_scope = TagScope._current
TagScope._current = self
return self
def __exit__(self, ptype, value, trace):
assert self._old_scope is None
if not self.accessed:
warnings.warn(f"Tag '{self.tag}' declared via TagScope was not used.")
TagScope._current = self._old_scope
def __call__(self, fdecl):
@functools.wraps(fdecl)
def tagged_fdecl(*args, **kwargs):
with self:
return fdecl(*args, **kwargs)
return tagged_fdecl
def tag_scope(tag):
"""The operator tag scope.
Parameters
----------
tag: str
The tag name.
Returns
-------
tag_scope: TagScope
The tag scope object, which can be used as decorator or
context manger.
Example
-------
.. code-block:: python
n = te.var('n')
m = te.var('m')
l = te.var('l')
A = te.placeholder((n, l), name='A')
B = te.placeholder((m, l), name='B')
k = te.reduce_axis((0, l), name='k')
with tvm.te.tag_scope(tag='matmul'):
C = te.compute((n, m), lambda i, j: te.sum(A[i, k] * B[j, k], axis=k))
# or use tag_scope as decorator
@tvm.te.tag_scope(tag="conv")
def compute_relu(data):
return te.compute(data.shape, lambda *i: tvm.tirx.Select(data(*i) < 0, 0.0, data(*i)))
"""
return TagScope(tag)
python/tvm/te/tensor.py
+356
View File
@@ -0,0 +1,356 @@
# 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.
"""Tensor class for computation declaration."""
# pylint: disable=invalid-name
import tvm_ffi
from tvm.runtime import Object, ObjectConvertible, const
from tvm.tirx import DataProducer
from tvm.tirx import expr as _expr
from . import _ffi_api, _te_tensor_overload
def _as_scalar_operand(value):
return value.asobject() if isinstance(value, TensorSlice) else value
class TensorSlice(ObjectConvertible):
"""Auxiliary data structure for enable slicing syntax from tensor."""
def __init__(self, tensor, indices):
if not isinstance(indices, tuple):
indices = (indices,)
self.tensor = tensor
self.indices = indices
def __getitem__(self, indices):
if not isinstance(indices, tuple):
indices = (indices,)
return TensorSlice(self.tensor, self.indices + indices)
def asobject(self):
"""Convert slice to object."""
return self.tensor.__call__(*self.indices)
@property
def dtype(self):
"""Data content of the tensor."""
return self.tensor.dtype
def expr_ty(self):
"""Compile-time element type of the tensor."""
return self.tensor.expr_ty()
def __add__(self, other):
result = _te_tensor_overload.__add__(self, other)
if result is not NotImplemented:
return result
return _expr.ExprOp.__add__(self.asobject(), _as_scalar_operand(other))
def __radd__(self, other):
result = _te_tensor_overload.__radd__(self, other)
if result is not NotImplemented:
return result
return _expr.ExprOp.__radd__(self.asobject(), _as_scalar_operand(other))
def __sub__(self, other):
result = _te_tensor_overload.__sub__(self, other)
if result is not NotImplemented:
return result
return _expr.ExprOp.__sub__(self.asobject(), _as_scalar_operand(other))
def __rsub__(self, other):
result = _te_tensor_overload.__rsub__(self, other)
if result is not NotImplemented:
return result
return _expr.ExprOp.__rsub__(self.asobject(), _as_scalar_operand(other))
def __mul__(self, other):
result = _te_tensor_overload.__mul__(self, other)
if result is not NotImplemented:
return result
return _expr.ExprOp.__mul__(self.asobject(), _as_scalar_operand(other))
def __rmul__(self, other):
result = _te_tensor_overload.__rmul__(self, other)
if result is not NotImplemented:
return result
return _expr.ExprOp.__rmul__(self.asobject(), _as_scalar_operand(other))
def __div__(self, other):
result = _te_tensor_overload.__div__(self, other)
if result is not NotImplemented:
return result
return _expr.ExprOp.__div__(self.asobject(), _as_scalar_operand(other))
def __rdiv__(self, other):
result = _te_tensor_overload.__rdiv__(self, other)
if result is not NotImplemented:
return result
return _expr.ExprOp.__rdiv__(self.asobject(), _as_scalar_operand(other))
def __truediv__(self, other):
result = _te_tensor_overload.__truediv__(self, other)
if result is not NotImplemented:
return result
return _expr.ExprOp.__truediv__(self.asobject(), _as_scalar_operand(other))
def __rtruediv__(self, other):
result = _te_tensor_overload.__rtruediv__(self, other)
if result is not NotImplemented:
return result
return _expr.ExprOp.__rtruediv__(self.asobject(), _as_scalar_operand(other))
def __floordiv__(self, other):
return _expr.ExprOp.__floordiv__(self.asobject(), _as_scalar_operand(other))
def __rfloordiv__(self, other):
return _expr.ExprOp.__rfloordiv__(self.asobject(), _as_scalar_operand(other))
def __mod__(self, other):
return _expr.ExprOp.__mod__(self.asobject(), _as_scalar_operand(other))
def __rmod__(self, other):
return _expr.ExprOp.__rmod__(self.asobject(), _as_scalar_operand(other))
def __neg__(self):
return _expr.ExprOp.__neg__(self.asobject())
def __lshift__(self, other):
return _expr.ExprOp.__lshift__(self.asobject(), _as_scalar_operand(other))
def __rlshift__(self, other):
return _expr.ExprOp.__rlshift__(self.asobject(), _as_scalar_operand(other))
def __rshift__(self, other):
return _expr.ExprOp.__rshift__(self.asobject(), _as_scalar_operand(other))
def __rrshift__(self, other):
return _expr.ExprOp.__rrshift__(self.asobject(), _as_scalar_operand(other))
def __and__(self, other):
return _expr.ExprOp.__and__(self.asobject(), _as_scalar_operand(other))
def __rand__(self, other):
return _expr.ExprOp.__rand__(self.asobject(), _as_scalar_operand(other))
def __or__(self, other):
return _expr.ExprOp.__or__(self.asobject(), _as_scalar_operand(other))
def __ror__(self, other):
return _expr.ExprOp.__ror__(self.asobject(), _as_scalar_operand(other))
def __xor__(self, other):
return _expr.ExprOp.__xor__(self.asobject(), _as_scalar_operand(other))
def __rxor__(self, other):
return _expr.ExprOp.__rxor__(self.asobject(), _as_scalar_operand(other))
def __invert__(self):
return _expr.ExprOp.__invert__(self.asobject())
def __lt__(self, other):
return _expr.ExprOp.__lt__(self.asobject(), _as_scalar_operand(other))
def __le__(self, other):
return _expr.ExprOp.__le__(self.asobject(), _as_scalar_operand(other))
def __eq__(self, other):
return _expr.ExprOp.__eq__(self.asobject(), _as_scalar_operand(other))
def __ne__(self, other):
return _expr.ExprOp.__ne__(self.asobject(), _as_scalar_operand(other))
def __gt__(self, other):
return _expr.ExprOp.__gt__(self.asobject(), _as_scalar_operand(other))
def __ge__(self, other):
return _expr.ExprOp.__ge__(self.asobject(), _as_scalar_operand(other))
def __nonzero__(self):
return _expr.ExprOp.__nonzero__(self.asobject())
def __bool__(self):
return self.__nonzero__()
def equal(self, other, span=None):
return _expr.ExprOp.equal(self.asobject(), _as_scalar_operand(other), span)
def astype(self, dtype, span=None):
return _expr.ExprOp.astype(self.asobject(), dtype, span)
class TensorOpBase:
"""Operator overloads for whole TE Tensor values."""
def __add__(self, other):
return _te_tensor_overload.__add__(self, other)
def __radd__(self, other):
return _te_tensor_overload.__radd__(self, other)
def __sub__(self, other):
return _te_tensor_overload.__sub__(self, other)
def __rsub__(self, other):
return _te_tensor_overload.__rsub__(self, other)
def __mul__(self, other):
return _te_tensor_overload.__mul__(self, other)
def __rmul__(self, other):
return _te_tensor_overload.__rmul__(self, other)
def __div__(self, other):
return _te_tensor_overload.__div__(self, other)
def __rdiv__(self, other):
return _te_tensor_overload.__rdiv__(self, other)
def __truediv__(self, other):
return _te_tensor_overload.__truediv__(self, other)
def __rtruediv__(self, other):
return _te_tensor_overload.__rtruediv__(self, other)
def __neg__(self):
return self.__mul__(const(-1, self.expr_ty()))
def __nonzero__(self):
return _expr.ExprOp.__nonzero__(self)
def __bool__(self):
return self.__nonzero__()
def equal(self, other, span=None):
return _expr.ExprOp.equal(self, other, span)
def astype(self, dtype, span=None):
result = _te_tensor_overload.astype(self, dtype, span)
if result is NotImplemented:
raise TypeError("TE Tensor overload astype is not registered")
return result
@tvm_ffi.register_object("te.Tensor")
class Tensor(DataProducer, TensorOpBase):
"""Tensor object, to construct, see function.Tensor"""
def __call__(self, *indices):
ndim = self.ndim
if len(indices) != ndim:
raise ValueError(
f"Need to provide {ndim} index in tensor but {len(indices)} was provided"
)
return _expr.ProducerLoad(self, indices)
def __getitem__(self, indices):
return TensorSlice(self, indices)
def __hash__(self):
return _ffi_api.TensorHash(self)
def __eq__(self, other):
if not isinstance(other, Tensor):
if isinstance(other, _expr.ExprOp):
return _expr.EqualOp(self, other)
return False
if self.ndim == 0 and other.ndim == 0:
raise ValueError(
"Equal == comparison among rank-0 tensor is ambiguous, "
"use Tensor.equal for content expression equvalence, "
"use Tensor.same_as for exact reference comparison"
)
return _ffi_api.TensorEqual(self, other)
@property
def ndim(self):
"""Dimension of the tensor."""
return len(self.shape)
@property
def dtype(self):
"""Data content of the tensor."""
return _ffi_api.TensorDType(self)
def expr_ty(self):
"""Compile-time element type of the tensor."""
return self.dtype
@property
def name(self):
op = self.op
if op.num_outputs == 1:
return op.name
return f"{op.name}.v{self.value_index}"
@tvm_ffi.register_object("te.Operation")
class Operation(Object):
"""Represent an operation that generates a tensor"""
def output(self, index):
"""Get the index-th output of the operation
Parameters
----------
index : int
The index size.
Returns
-------
out : Tensor
The i-th output.
"""
return _ffi_api.OpGetOutput(self, index)
@property
def num_outputs(self):
"""Number of outputs from this op."""
return _ffi_api.OpNumOutputs(self)
@property
def input_tensors(self):
"""List of input tensors to this op."""
return _ffi_api.OpInputTensors(self)
@tvm_ffi.register_object("te.PlaceholderOp")
class PlaceholderOp(Operation):
"""Placeholder operation."""
@tvm_ffi.register_object("te.BaseComputeOp")
class BaseComputeOp(Operation):
"""Compute operation."""
@tvm_ffi.register_object("te.ComputeOp")
class ComputeOp(BaseComputeOp):
"""Scalar operation."""
@tvm_ffi.register_object("te.ScanOp")
class ScanOp(Operation):
"""Scan operation."""
@tvm_ffi.register_object("te.ExternOp")
class ExternOp(Operation):
"""External operation."""