chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,774 @@
|
||||
# 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,
|
||||
)
|
||||
Reference in New Issue
Block a user