wehub-resource-sync
775 lines
28 KiB
Python
775 lines
28 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.
|
|
# pylint: disable=invalid-name, too-many-locals, too-many-statements
|
|
"Scan related operators"
|
|
|
|
import operator
|
|
from collections.abc import Callable
|
|
|
|
import tvm
|
|
from tvm import te
|
|
from tvm.contrib.thrust import can_use_rocthrust, can_use_thrust
|
|
from tvm.script.ir_builder import IRBuilder
|
|
from tvm.script.ir_builder import tirx as T
|
|
|
|
from ..math import cast, ceil_log2
|
|
from ..transform import expand_dims, reshape, squeeze, transpose
|
|
from ..utils import ceil_div, get_const_int, prod, swap
|
|
|
|
_THRUST_SUM_SCAN = "tvm.contrib.thrust.sum_scan"
|
|
|
|
|
|
def _get_thrust_func_name(tvmop):
|
|
if tvmop is not operator.add:
|
|
raise ValueError(f"{tvmop} not supported by thrust")
|
|
return _THRUST_SUM_SCAN
|
|
|
|
|
|
def _can_use_scan_thrust(binop):
|
|
"""
|
|
Check if scan_thrust can be utilized based on the current target and binary op.
|
|
"""
|
|
target = tvm.target.Target.current()
|
|
if target is None:
|
|
return False
|
|
return binop is operator.add and any(
|
|
[
|
|
can_use_thrust(target, _THRUST_SUM_SCAN),
|
|
can_use_rocthrust(target, _THRUST_SUM_SCAN),
|
|
]
|
|
)
|
|
|
|
|
|
def exclusive_scan_ir(data, output, reduction=None, binop=operator.add, identity_value=0):
|
|
"""Low level IR to do exclusive sum scan along rows of 2D input.
|
|
|
|
Parameters
|
|
----------
|
|
data : Buffer
|
|
Input N-D Buffer. Scan is done over the innermost axis.
|
|
|
|
output: Buffer
|
|
A buffer to store the output scan, of the same shape as data
|
|
|
|
reduction: Buffer, optional
|
|
(N-1)-D Buffer, to store the sum of each scan axis.
|
|
|
|
binop: function, optional
|
|
A binary associative op to use for scan. The function takes two TIR expressions
|
|
and produce a new TIR expression. By default it uses ``operator.add`` to compute prefix
|
|
sum.
|
|
|
|
identity_value: int or float
|
|
A value for the binary operation which provides the identity property. E.g. if * is
|
|
your operator and i is the identity_value then a * i = a for all a in the domain of
|
|
your operation.
|
|
"""
|
|
|
|
batch_size = cast(prod(data.shape[:-1]), "int32")
|
|
scan_axis_size = cast(data.shape[-1], "int32")
|
|
|
|
with IRBuilder() as ib:
|
|
data = T.buffer_proxy(data)
|
|
output = T.buffer_proxy(output)
|
|
out_dtype = output.dtype
|
|
|
|
if reduction is not None:
|
|
reduction = T.buffer_proxy(reduction)
|
|
|
|
max_threads = int(tvm.target.Target.current(allow_none=False).attrs["max_num_threads"])
|
|
|
|
with T.If(scan_axis_size == 0):
|
|
with T.Then():
|
|
bx = te.thread_axis("blockIdx.x")
|
|
with T.attr(bx, "thread_extent", batch_size):
|
|
with T.If(bx < batch_size):
|
|
with T.Then():
|
|
if reduction is not None:
|
|
reduction[bx] = cast(identity_value, out_dtype)
|
|
with T.Else():
|
|
nthread_tx = max_threads
|
|
nthread_bx = ceil_div(scan_axis_size, max_threads)
|
|
nthread_by = batch_size
|
|
|
|
# Copy data to output
|
|
tx = te.thread_axis("threadIdx.x")
|
|
bx = te.thread_axis("blockIdx.x")
|
|
by = te.thread_axis("blockIdx.y")
|
|
with T.frame_scope(
|
|
[
|
|
T.attr(tx, "thread_extent", nthread_tx),
|
|
T.attr(bx, "thread_extent", nthread_bx),
|
|
T.attr(by, "thread_extent", nthread_by),
|
|
]
|
|
):
|
|
tid = bx * nthread_tx + tx
|
|
with T.If(tid < scan_axis_size):
|
|
with T.Then():
|
|
output[by * scan_axis_size + tid] = cast(
|
|
data[by * scan_axis_size + tid], out_dtype
|
|
)
|
|
|
|
# The following algorithm performs parallel exclusive scan
|
|
# Up Sweep of exclusive scan
|
|
lim = ceil_log2(scan_axis_size)
|
|
|
|
with T.serial(0, cast(lim, "int32")) as l2_width:
|
|
width = 2 << l2_width
|
|
|
|
tx = te.thread_axis("threadIdx.x")
|
|
bx = te.thread_axis("blockIdx.x")
|
|
by = te.thread_axis("blockIdx.y")
|
|
start_buf = T.decl_buffer([1], "int32", scope="local")
|
|
middle_buf = T.decl_buffer([1], "int32", scope="local")
|
|
end_buf = T.decl_buffer([1], "int32", scope="local")
|
|
with T.frame_scope(
|
|
[
|
|
T.attr(tx, "thread_extent", nthread_tx),
|
|
T.attr(
|
|
bx,
|
|
"thread_extent",
|
|
cast(ceil_div(scan_axis_size, max_threads * width), "int32"),
|
|
),
|
|
T.attr(by, "thread_extent", nthread_by),
|
|
]
|
|
):
|
|
tid = bx * nthread_tx + tx
|
|
start = T.buffer_proxy(start_buf)
|
|
middle = T.buffer_proxy(middle_buf)
|
|
end = T.buffer_proxy(end_buf)
|
|
start[0] = width * tid
|
|
with T.If(start[0] < scan_axis_size):
|
|
with T.Then():
|
|
middle[0] = start[0] + tvm.tirx.indexdiv(width, 2)
|
|
end[0] = tvm.te.min(start[0] + width, scan_axis_size)
|
|
with T.If(middle[0] < scan_axis_size):
|
|
with T.Then():
|
|
output[by * scan_axis_size + end[0] - 1] = binop(
|
|
output[by * scan_axis_size + end[0] - 1],
|
|
output[by * scan_axis_size + middle[0] - 1],
|
|
)
|
|
|
|
# Down Sweep of exclusive scan
|
|
bx = te.thread_axis("blockIdx.x")
|
|
with T.attr(bx, "thread_extent", batch_size):
|
|
with T.If(bx < batch_size):
|
|
with T.Then():
|
|
if reduction is not None:
|
|
reduction[bx] = output[(bx + 1) * scan_axis_size - 1]
|
|
output[(bx + 1) * scan_axis_size - 1] = cast(identity_value, out_dtype)
|
|
|
|
with T.serial(0, cast(lim, "int32")) as l2_width:
|
|
width = 2 << (lim - l2_width - 1)
|
|
|
|
tx = te.thread_axis("threadIdx.x")
|
|
bx = te.thread_axis("blockIdx.x")
|
|
by = te.thread_axis("blockIdx.y")
|
|
start_buf = T.decl_buffer([1], "int32", scope="local")
|
|
middle_buf = T.decl_buffer([1], "int32", scope="local")
|
|
end_buf = T.decl_buffer([1], "int32", scope="local")
|
|
tmp_buf = T.decl_buffer([1], out_dtype, scope="local")
|
|
with T.frame_scope(
|
|
[
|
|
T.attr(tx, "thread_extent", nthread_tx),
|
|
T.attr(
|
|
bx,
|
|
"thread_extent",
|
|
cast(ceil_div(scan_axis_size, max_threads * width), "int32"),
|
|
),
|
|
T.attr(by, "thread_extent", nthread_by),
|
|
]
|
|
):
|
|
tid = bx * nthread_tx + tx
|
|
start = T.buffer_proxy(start_buf)
|
|
middle = T.buffer_proxy(middle_buf)
|
|
end = T.buffer_proxy(end_buf)
|
|
tmp = T.buffer_proxy(tmp_buf)
|
|
start[0] = width * tid
|
|
with T.If(tvm.tirx.all(start[0] < scan_axis_size)):
|
|
with T.Then():
|
|
middle[0] = start[0] + tvm.tirx.indexdiv(width, 2)
|
|
end[0] = tvm.tirx.min(start[0] + width, scan_axis_size)
|
|
with T.If(middle[0] < scan_axis_size):
|
|
with T.Then():
|
|
tmp[0] = output[by * scan_axis_size + middle[0] - 1]
|
|
output[by * scan_axis_size + middle[0] - 1] = output[
|
|
by * scan_axis_size + end[0] - 1
|
|
]
|
|
output[by * scan_axis_size + end[0] - 1] = binop(
|
|
output[by * scan_axis_size + end[0] - 1], tmp[0]
|
|
)
|
|
|
|
return ib.get()
|
|
|
|
|
|
def get_reduction_from_exclusive_scan(data, ex_scan_output, binop=operator.add):
|
|
"""Return the sum of the last element of data and the exclusive scan output.
|
|
The is the reduction of data along each row (for 2-D case).
|
|
|
|
Parameters
|
|
----------
|
|
data : tvm.te.Tensor
|
|
Input data of any shape
|
|
|
|
ex_scan_output : tvm.te.Tensor
|
|
The output of exclusive scan on data
|
|
|
|
binop: function, optional
|
|
A binary associative op to use for scan. The function takes two TIR expressions
|
|
and produce a new TIR expression. By default it uses ``operator.add`` to compute prefix
|
|
sum.
|
|
|
|
Returns
|
|
-------
|
|
reduction : tvm.te.Tensor
|
|
(N-1)-D tensor storing the reduction of each scan axis.
|
|
"""
|
|
ndim = len(data.shape)
|
|
if ndim == 1:
|
|
data = expand_dims(data, axis=0)
|
|
ex_scan_output = expand_dims(ex_scan_output, axis=0)
|
|
|
|
def ir(data_buf, data_ex_scan_buf, reduction_buf):
|
|
batch_size = cast(prod(data_buf.shape[:-1]), "int32")
|
|
scan_axis_size = cast(data_buf.shape[-1], "int32")
|
|
|
|
max_threads = int(tvm.target.Target.current(allow_none=False).attrs["max_num_threads"])
|
|
|
|
with IRBuilder() as ib:
|
|
data = T.buffer_proxy(data_buf)
|
|
data_ex_scan = T.buffer_proxy(data_ex_scan_buf)
|
|
reduction = T.buffer_proxy(reduction_buf)
|
|
|
|
nthread_tx = max_threads
|
|
nthread_bx = ceil_div(batch_size, max_threads)
|
|
tx = te.thread_axis("threadIdx.x")
|
|
bx = te.thread_axis("blockIdx.x")
|
|
with T.frame_scope(
|
|
[
|
|
T.attr(tx, "thread_extent", nthread_tx),
|
|
T.attr(bx, "thread_extent", nthread_bx),
|
|
]
|
|
):
|
|
tid = bx * max_threads + tx
|
|
with T.If(tid < batch_size):
|
|
with T.Then():
|
|
with T.If(scan_axis_size > 0):
|
|
with T.Then():
|
|
reduction[tid] = binop(
|
|
data_ex_scan[tid * scan_axis_size + scan_axis_size - 1],
|
|
data[tid * scan_axis_size + scan_axis_size - 1],
|
|
)
|
|
with T.Else():
|
|
reduction[tid] = cast(0, reduction_buf.dtype)
|
|
|
|
return ib.get()
|
|
|
|
data_buf = tvm.tirx.decl_buffer(
|
|
data.shape, data.dtype, "valid_indices_buf", data_alignment=8, layout=None
|
|
)
|
|
ex_scan_output_buf = tvm.tirx.decl_buffer(
|
|
ex_scan_output.shape,
|
|
ex_scan_output.dtype,
|
|
"ex_scan_output_buf",
|
|
data_alignment=8,
|
|
layout=None,
|
|
)
|
|
|
|
reduction = te.extern(
|
|
[data.shape[:-1]],
|
|
[data, ex_scan_output],
|
|
lambda ins, outs: ir(ins[0], ins[1], outs[0]),
|
|
dtype=[ex_scan_output.dtype],
|
|
in_buffers=[data_buf, ex_scan_output_buf],
|
|
name="ex_scan_reduction",
|
|
tag="ex_scan_reduction_gpu",
|
|
)
|
|
|
|
if ndim == 1:
|
|
return squeeze(reduction, 0)
|
|
|
|
return reduction
|
|
|
|
|
|
def scan_thrust(
|
|
data,
|
|
output_dtype,
|
|
exclusive=True,
|
|
return_reduction=False,
|
|
binop=operator.add,
|
|
workspace=None,
|
|
):
|
|
"""Do exclusive or inclusive scan on 1D or multidimensional input, using thrust.
|
|
|
|
Parameters
|
|
----------
|
|
data : tvm.te.Tensor
|
|
Input data of any shape. The scan is done over the innermost axis.
|
|
|
|
output_dtype: string
|
|
The dtype of the output scan tensor.
|
|
|
|
exclusive: bool, optional
|
|
Whether or not do exclusive or inclusive scan.
|
|
|
|
return_reduction: bool, optional
|
|
Whether or not return a (N-1)-D tensor storing the reduction of each scan axis.
|
|
Reductions are computed as part of the upsweep pass, so there is no extra cost.
|
|
If False, reductions are ignored. It must be False when exclusive is False.
|
|
|
|
binop: function, optional
|
|
A binary associative op to use for scan. Since we need to lookup the corresponding
|
|
thrust function, arbitrariy callables are not supported. Currently only
|
|
``operator.add`` can be passed in.
|
|
|
|
workspace: Optional[tvm.te.Tensor]
|
|
A buffer to store intermediate results. The size of the workspace should be sufficiently
|
|
large, this can be obtained by overestimation or memory usage profiling. If None, it will
|
|
fallback to use thrust internal memory allocation.
|
|
|
|
Returns
|
|
-------
|
|
output : tvm.te.Tensor
|
|
A N-D tensor of the same rank N and shape as the input data.
|
|
|
|
reduction : tvm.te.Tensor, optional
|
|
(N-1)-D tensor storing the reduction of each scan axis.
|
|
Returned if return_reduction is True.
|
|
"""
|
|
data_buf = tvm.tirx.decl_buffer(
|
|
data.shape, data.dtype, "data_buf", data_alignment=8, layout=None
|
|
)
|
|
output_buf = tvm.tirx.decl_buffer(
|
|
data.shape, output_dtype, "output_buf", data_alignment=8, layout=None
|
|
)
|
|
|
|
workspace_buf = (
|
|
tvm.tirx.decl_buffer(
|
|
workspace.shape, workspace.dtype, "workspace_buf", data_alignment=8, layout=None
|
|
)
|
|
if workspace is not None
|
|
else None
|
|
)
|
|
|
|
def f_compute(ins, outs):
|
|
args = [_get_thrust_func_name(binop), ins[0], outs[0], exclusive]
|
|
if workspace is not None:
|
|
args.append(ins[1])
|
|
return tvm.tirx.call_packed(*args)
|
|
|
|
output = te.extern(
|
|
[data.shape],
|
|
[data] if workspace is None else [data, workspace],
|
|
f_compute,
|
|
dtype=[output_dtype],
|
|
in_buffers=[data_buf] if workspace is None else [data_buf, workspace_buf],
|
|
out_buffers=[output_buf],
|
|
name="exclusive_scan_thrust",
|
|
tag="exclusive_scan_thrust_gpu",
|
|
)
|
|
|
|
if return_reduction:
|
|
assert exclusive, "return_reduction should be False for inclusive scan"
|
|
reduction = get_reduction_from_exclusive_scan(data, output, binop)
|
|
return output, reduction
|
|
|
|
return output
|
|
|
|
|
|
def exclusive_scan(
|
|
data,
|
|
axis=-1,
|
|
return_reduction=False,
|
|
output_dtype=None,
|
|
binop=operator.add,
|
|
identity_value=0,
|
|
workspace=None,
|
|
):
|
|
"""Do exclusive scan on 1D or multidimensional input.
|
|
|
|
Parameters
|
|
----------
|
|
data : tvm.te.Tensor
|
|
Input data of any shape.
|
|
|
|
axis: int, optional
|
|
The axis to do scan on. By default, scan is done on the innermost axis.
|
|
|
|
return_reduction: bool, optional
|
|
Whether or not return a tensor storing the reduction over each scan axis.
|
|
If the input rank is N, this tensor is of rank N - 1.
|
|
Reductions are computed as part of the upsweep pass, so there is no extra cost.
|
|
If False, reductions are ignored.
|
|
|
|
output_dtype: string, optional
|
|
The dtype of the output scan tensor. If not provided, the dtype of the input is used.
|
|
|
|
binop: function, optional
|
|
A binary associative op to use for scan. The function takes two TIR expressions
|
|
and produce a new TIR expression. By default it uses ``operator.add`` to compute prefix
|
|
sum.
|
|
|
|
identity_value: int or float
|
|
A value for the binary operation which provides the identity property. E.g. if * is
|
|
your operator and i is the identity_value then a * i = a for all a in the domain of
|
|
your operation.
|
|
|
|
workspace: Optional[tvm.te.Tensor]
|
|
A buffer to store intermediate results if thrust is enabled. The size of the workspace
|
|
should be sufficiently large, this can be obtained by overestimation or memory usage
|
|
profiling. If None, it will fallback to use thrust internal memory allocation.
|
|
|
|
Returns
|
|
-------
|
|
output : tvm.te.Tensor
|
|
A N-D tensor of the same rank N and shape as the input data.
|
|
|
|
reduction : tvm.te.Tensor, optional
|
|
(N-1)-D tensor storing the reduction of each scan axis.
|
|
Returned if return_reduction is True.
|
|
"""
|
|
|
|
def do_scan(data, output_dtype):
|
|
# TODO: add support for a prod_scan
|
|
if _can_use_scan_thrust(binop):
|
|
return scan_thrust(
|
|
data,
|
|
output_dtype,
|
|
exclusive=True,
|
|
return_reduction=return_reduction,
|
|
binop=binop,
|
|
workspace=workspace,
|
|
)
|
|
|
|
if ndim == 1:
|
|
# TIR exclusive scan accepts only 2D or higher-rank inputs.
|
|
data = expand_dims(data, axis=0)
|
|
|
|
data_buf = tvm.tirx.decl_buffer(
|
|
data.shape, data.dtype, "data_buf", data_alignment=8, layout=None
|
|
)
|
|
output_buf = tvm.tirx.decl_buffer(
|
|
data.shape, output_dtype, "output_buf", data_alignment=8, layout=None
|
|
)
|
|
|
|
if return_reduction:
|
|
output, reduction = te.extern(
|
|
[data.shape, data.shape[:-1]],
|
|
[data],
|
|
lambda ins, outs: exclusive_scan_ir(
|
|
ins[0], outs[0], outs[1], binop=binop, identity_value=identity_value
|
|
),
|
|
dtype=[output_dtype, output_dtype],
|
|
in_buffers=[data_buf],
|
|
name="exclusive_scan",
|
|
tag="exclusive_scan_gpu",
|
|
)
|
|
else:
|
|
output = te.extern(
|
|
[data.shape],
|
|
[data],
|
|
lambda ins, outs: exclusive_scan_ir(
|
|
ins[0], outs[0], binop=binop, identity_value=identity_value
|
|
),
|
|
dtype=[output_dtype],
|
|
in_buffers=[data_buf],
|
|
out_buffers=[output_buf],
|
|
name="exclusive_scan",
|
|
tag="exclusive_scan_gpu",
|
|
)
|
|
reduction = None
|
|
|
|
if ndim == 1:
|
|
output = squeeze(output, 0)
|
|
if return_reduction:
|
|
reduction = squeeze(reduction, 0)
|
|
|
|
if return_reduction:
|
|
return output, reduction
|
|
|
|
return output
|
|
|
|
if output_dtype is None or output_dtype == "":
|
|
output_dtype = data.dtype
|
|
|
|
ndim = len(data.shape)
|
|
if axis < 0:
|
|
axis += ndim
|
|
|
|
# If scan axis is not the innermost one, swap the scan and the innermost axes
|
|
# Scan is always done on the innermost axis, for performance reason.
|
|
if axis != ndim - 1:
|
|
axes = swap(list(range(ndim)), axis)
|
|
data = transpose(data, axes)
|
|
|
|
if return_reduction:
|
|
output, reduction = do_scan(data, output_dtype)
|
|
else:
|
|
output = do_scan(data, output_dtype)
|
|
|
|
if axis != ndim - 1:
|
|
axes = swap(list(range(ndim)), axis)
|
|
output = transpose(output, axes)
|
|
|
|
if return_reduction:
|
|
return output, reduction
|
|
|
|
return output
|
|
|
|
|
|
def inclusive_scan(
|
|
data, axis=-1, output_dtype=None, binop=operator.add, identity_value=0, workspace=None
|
|
):
|
|
"""Do inclusive scan on 1D or multidimensional input.
|
|
|
|
Parameters
|
|
----------
|
|
data : tvm.te.Tensor
|
|
Input data of any shape.
|
|
|
|
axis: int, optional
|
|
The axis to do scan on. By default, scan is done on the innermost axis.
|
|
|
|
output_dtype: string, optional
|
|
The dtype of the output scan tensor. If not provided, the dtype of the input is used.
|
|
|
|
binop: function, optional
|
|
A binary associative op to use for scan. The function takes two TIR expressions
|
|
and produce a new TIR expression. By default it uses ``operator.add`` to compute prefix
|
|
sum.
|
|
|
|
identity_value: int or float
|
|
A value for the binary operation which provides the identity property. E.g. if * is
|
|
your operator and i is the identity_value then a * i = a for all a in the domain of
|
|
your operation.
|
|
|
|
workspace: Optional[tvm.te.Tensor]
|
|
A buffer to store intermediate results if thrust is enabled. The size of the workspace
|
|
should be sufficiently large, this can be obtained by overestimation or memory usage
|
|
profiling. If None, it will fallback to use thrust internal memory allocation.
|
|
|
|
Returns
|
|
-------
|
|
output : tvm.te.Tensor
|
|
A N-D tensor of the same rank N as the input data.
|
|
"""
|
|
|
|
if _can_use_scan_thrust(binop):
|
|
if output_dtype is None or output_dtype == "":
|
|
output_dtype = data.dtype
|
|
ndim = len(data.shape)
|
|
if axis < 0:
|
|
axis += ndim
|
|
|
|
if axis != ndim - 1:
|
|
axes = swap(list(range(ndim)), axis)
|
|
data = transpose(data, axes)
|
|
output = scan_thrust(data, output_dtype, exclusive=False, binop=binop, workspace=workspace)
|
|
if axis != ndim - 1:
|
|
axes = swap(list(range(ndim)), axis)
|
|
output = transpose(output, axes)
|
|
return output
|
|
|
|
ex_scan = exclusive_scan(
|
|
data,
|
|
axis,
|
|
output_dtype=output_dtype,
|
|
binop=binop,
|
|
identity_value=identity_value,
|
|
workspace=workspace,
|
|
)
|
|
|
|
if output_dtype is not None and data.dtype != output_dtype and output_dtype != "":
|
|
data = cast(data, output_dtype)
|
|
|
|
return binop(data, ex_scan)
|
|
|
|
|
|
def scanop(
|
|
data: tvm.te.Tensor,
|
|
binop: Callable[["tvm.Expr", "tvm.Expr"], "tvm.Expr"],
|
|
identity_value: float | int,
|
|
axis: int | None = None,
|
|
dtype: str | None = None,
|
|
exclusive: bool | None = None,
|
|
workspace: tvm.te.Tensor | None = None,
|
|
) -> tvm.te.Tensor:
|
|
"""Cumulative binary operator (scan) with similar axis behavior as np.cumsum and np.cumprod.
|
|
|
|
See cumprod and cumsum for an example of use.
|
|
|
|
E.g. if * is your binary operator and the input tensor is [1, 2, 3, 4] the output may be
|
|
[1, 1 * 2, 1 * 2 * 3, 1 * 2 * 3 * 4]
|
|
|
|
Parameters
|
|
----------
|
|
data : tvm.te.Tensor
|
|
The input data to the operator.
|
|
|
|
binop: Callable (tvm.Expr, tvm.Expr) -> tvm.Expr
|
|
A binary operator which should be associative and commutative. E.g. if * is your
|
|
operator then a * (b * c) = (a * b) * c and a * b = b * a
|
|
|
|
identity_value: int or float
|
|
A value for the binary operation which provides the identity property. E.g. if * is
|
|
your operator and i is the identity_value then a * i = a for all a in the domain of
|
|
your operation.
|
|
|
|
axis : int, optional
|
|
Axis along which the operation is computed. The default (None) is to compute
|
|
the cumulative operation over the flattened array.
|
|
|
|
dtype : string, optional
|
|
Type of the returned array and of the accumulator in which the elements are computed.
|
|
If dtype is not specified, it defaults to the dtype of data.
|
|
|
|
exclusive : bool, optional
|
|
If true will return exclusive cumulative operation in which the first element is not
|
|
included. In other terms, if true, the j-th output element would be
|
|
the cumulative operation of the first (j-1) elements. Otherwise, it would be the
|
|
cumulative operation of the first j elements.
|
|
|
|
workspace: Optional[tvm.te.Tensor]
|
|
|
|
Returns
|
|
-------
|
|
result : tvm.te.Tensor
|
|
The result has the same size as data, and the same shape as data if axis is not None.
|
|
If axis is None, the result is a 1-d array.
|
|
"""
|
|
if axis is None:
|
|
axis = 0
|
|
data = reshape(data, (prod(data.shape),))
|
|
axis = get_const_int(axis)
|
|
if exclusive is not None and exclusive:
|
|
return exclusive_scan(
|
|
data,
|
|
axis,
|
|
output_dtype=dtype,
|
|
binop=binop,
|
|
identity_value=identity_value,
|
|
workspace=workspace,
|
|
)
|
|
return inclusive_scan(
|
|
data,
|
|
axis,
|
|
output_dtype=dtype,
|
|
binop=binop,
|
|
identity_value=identity_value,
|
|
workspace=workspace,
|
|
)
|
|
|
|
|
|
def cumsum(
|
|
data: tvm.te.Tensor,
|
|
axis: int | None = None,
|
|
dtype: int | None = None,
|
|
exclusive: bool | None = None,
|
|
workspace: tvm.te.Tensor | None = None,
|
|
) -> tvm.te.Tensor:
|
|
"""Numpy style cumsum op. Return the cumulative sum of the elements along a given axis.
|
|
|
|
Parameters
|
|
----------
|
|
data : tvm.te.Tensor
|
|
The input data to the operator.
|
|
|
|
axis : int, optional
|
|
Axis along which the cumulative sum is computed. The default (None) is to compute
|
|
the cumsum over the flattened array.
|
|
|
|
dtype : string, optional
|
|
Type of the returned array and of the accumulator in which the elements are summed.
|
|
If dtype is not specified, it defaults to the dtype of data.
|
|
|
|
exclusive : bool, optional
|
|
If true will return exclusive sum in which the first element is not
|
|
included. In other terms, if true, the j-th output element would be
|
|
the sum of the first (j-1) elements. Otherwise, it would be the sum of
|
|
the first j elements.
|
|
|
|
workspace: Optional[tvm.te.Tensor]
|
|
A buffer to store intermediate results if thrust is enabled. The size of the workspace
|
|
should be sufficiently large, this can be obtained by overestimation or memory usage
|
|
profiling. If None, it will fallback to use thrust internal memory allocation.
|
|
|
|
Returns
|
|
-------
|
|
result : tvm.te.Tensor
|
|
The result has the same size as data, and the same shape as data if axis is not None.
|
|
If axis is None, the result is a 1-d array.
|
|
"""
|
|
return scanop(
|
|
data=data,
|
|
binop=operator.add,
|
|
identity_value=0,
|
|
axis=axis,
|
|
dtype=dtype,
|
|
exclusive=exclusive,
|
|
workspace=workspace,
|
|
)
|
|
|
|
|
|
def cumprod(
|
|
data: tvm.te.Tensor,
|
|
axis: int | None = None,
|
|
dtype: int | None = None,
|
|
exclusive: bool | None = None,
|
|
workspace: tvm.te.Tensor | None = None,
|
|
):
|
|
"""Numpy style cumprod op. Return the cumulative product of the elements along a given axis.
|
|
|
|
Parameters
|
|
----------
|
|
data : tvm.te.Tensor
|
|
The input data to the operator.
|
|
|
|
axis : int, optional
|
|
Axis along which the cumulative product is computed. The default (None) is to compute
|
|
the cumproduct over the flattened array.
|
|
|
|
dtype : string, optional
|
|
Type of the returned array and of the accumulator in which the elements are multiplied.
|
|
If dtype is not specified, it defaults to the dtype of data.
|
|
|
|
exclusive : bool, optional
|
|
If True, will return exclusive product in which the first element is not
|
|
included. In other terms, if True, the j-th output element would be
|
|
the product of the first (j-1) elements. Otherwise, it would be the product of
|
|
the first j elements.
|
|
|
|
workspace: Optional[tvm.te.Tensor]
|
|
A buffer to store intermediate results if thrust is enabled. The size of the workspace
|
|
should be sufficiently large, this can be obtained by overestimation or memory usage
|
|
profiling. If None, it will fallback to use thrust internal memory allocation.
|
|
|
|
Returns
|
|
-------
|
|
result : tvm.te.Tensor
|
|
The result has the same size as data, and the same shape as data if axis is not None.
|
|
If axis is None, the result is a 1-d array.
|
|
"""
|
|
return scanop(
|
|
data=data,
|
|
binop=operator.mul,
|
|
identity_value=1,
|
|
axis=axis,
|
|
dtype=dtype,
|
|
exclusive=exclusive,
|
|
workspace=workspace,
|
|
)
|