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/topi/gpu/__init__.py
+25
View File
@@ -0,0 +1,25 @@
# 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=redefined-builtin, wildcard-import
"""GPU specific declaration."""
from .scan import cumsum, cumprod
from .scatter_elements import scatter_elements
from .scatter_nd import scatter_nd
from .sort import *
python/tvm/topi/gpu/scan.py
+774
View File
@@ -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,
)
python/tvm/topi/gpu/scatter_elements.py
+162
View File
@@ -0,0 +1,162 @@
# 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
"""scatter_elements related operators"""
import tvm
from tvm import te, tirx
from tvm.script.ir_builder import IRBuilder
from tvm.script.ir_builder import tirx as T
from .. import utils
from ..math import cast
from ..utils import ceil_div
def scatter_elements(data, indices, updates, axis=0, reduction="update"):
"""GPU implementation of scatter_elements with explicit thread bindings"""
if not isinstance(axis, int):
axis = utils.get_const_int(axis)
# Prepare ranges and strides
shape = data.shape
if axis < 0:
axis = len(shape) + axis
axis_range = cast(shape[axis], indices.dtype)
full_range = 1
after_axis_range = 1
for i, value in enumerate(shape, 0):
full_range *= value
if i > axis:
after_axis_range *= value
before_axis_stride = axis_range * after_axis_range
ind_shape = indices.shape
ind_axis_range = ind_shape[axis]
ind_before_axis_range = 1
ind_after_axis_range = 1
for i, value in enumerate(ind_shape, 0):
if i < axis:
ind_before_axis_range *= value
elif i > axis:
ind_after_axis_range *= value
ind_before_axis_stride = ind_axis_range * ind_after_axis_range
ind_full_range_excl_axis = ind_before_axis_range * ind_after_axis_range
def gen_ir(data_ptr, indices_ptr, updates_ptr, out_ptr, reduce_func):
# pylint: disable=invalid-name
data = T.buffer_proxy(data_ptr)
indices = T.buffer_proxy(indices_ptr)
updates = T.buffer_proxy(updates_ptr)
out = T.buffer_proxy(out_ptr)
max_threads = int(tvm.target.Target.current(allow_none=False).attrs["max_num_threads"])
with IRBuilder() as ib:
with T.seq_scope():
# Init
nthread_bx_init = cast(ceil_div(full_range, max_threads), "int32")
tx_init = te.thread_axis("threadIdx.x")
bx_init = te.thread_axis("blockIdx.x")
with T.frame_scope(
[
T.attr(bx_init, "thread_extent", nthread_bx_init),
T.attr(tx_init, "thread_extent", max_threads),
]
):
tid = bx_init * max_threads + tx_init
with T.If(tid < full_range):
with T.Then():
out[tid] = data[tid]
# Scatter
nthread_bx_scat = cast(ceil_div(ind_full_range_excl_axis, max_threads), "int32")
tx_scat = te.thread_axis("threadIdx.x")
bx_scat = te.thread_axis("blockIdx.x")
with T.frame_scope(
[
T.attr(bx_scat, "thread_extent", nthread_bx_scat),
T.attr(tx_scat, "thread_extent", max_threads),
]
):
fused = bx_scat * max_threads + tx_scat
with T.If(fused < ind_full_range_excl_axis):
with T.Then():
i = fused // ind_after_axis_range
j = fused % ind_after_axis_range
pre_index1 = i * ind_before_axis_stride + j
pre_index2 = i * before_axis_stride + j
with T.serial(0, ind_axis_range) as k:
# Offset along indices or updates
index1 = pre_index1 + k * ind_after_axis_range
# Get index and shift to positive side if need
k_new = indices[index1]
shifted_index = k_new + (k_new < 0) * axis_range
# Offset along data
index2 = pre_index2 + shifted_index * after_axis_range
reduce_func(out, index2, updates[index1])
return ib.get()
def update_func(dst_ptr, dst_index, update):
dst_ptr[dst_index] = update
def add_func(dst_ptr, dst_index, update):
dst_ptr[dst_index] += update
def mul_func(dst_ptr, dst_index, update):
dst_ptr[dst_index] *= update
def mean_func(dst_ptr, dst_index, update):
dst_ptr[dst_index] = (dst_ptr[dst_index] + update) / 2
def min_func(dst_ptr, dst_index, update):
dst_ptr[dst_index] = tirx.min(dst_ptr[dst_index], update)
def max_func(dst_ptr, dst_index, update):
dst_ptr[dst_index] = tirx.max(dst_ptr[dst_index], update)
reduce_func = None
if reduction == "update":
reduce_func = update_func
elif reduction == "add":
reduce_func = add_func
elif reduction == "mul":
reduce_func = mul_func
elif reduction == "mean":
reduce_func = mean_func
elif reduction == "min":
reduce_func = min_func
elif reduction == "max":
reduce_func = max_func
else:
raise NotImplementedError(
"scatter_elements reduction not in [update, add, mul, mean, min, max]:", reduction
)
out_buf = tirx.decl_buffer(data.shape, data.dtype, "out_buf", layout=None)
return te.extern(
[data.shape],
[data, indices, updates],
lambda ins, outs: gen_ir(ins[0], ins[1], ins[2], outs[0], reduce_func),
dtype=data.dtype,
out_buffers=[out_buf],
name="scatter_elements.gpu",
tag="scatter_elements.gpu",
)
python/tvm/topi/gpu/scatter_nd.py
+129
View File
@@ -0,0 +1,129 @@
# 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
# ruff: noqa: E741
"""scatter_nd related operators"""
import tvm
from tvm import te, tirx # hide redefinition of min and max
from tvm.script.ir_builder import IRBuilder
from tvm.script.ir_builder import tirx as T
from ..math import cast
from ..scatter import _verify_scatter_nd_inputs
from ..utils import ceil_div
def scatter_nd(data, indices, updates, mode):
"""GPU implementation of scatter_nd with explicit thread bindings."""
_verify_scatter_nd_inputs(data, indices, updates)
def gen_ir(data_ptr, indices_ptr, updates_ptr, out_ptr):
# pylint: disable=invalid-name
data = T.buffer_proxy(data_ptr)
indices = T.buffer_proxy(indices_ptr)
updates = T.buffer_proxy(updates_ptr)
out = T.buffer_proxy(out_ptr)
# We combine all the indices dimensions but the first one into a single
# dimension so we can iterate it in single loop instead of an arbitrary
# number of loops. We do the same thing for all the update dimensions.
fused_indices_dimension = 1
for i in indices_ptr.shape[1:]:
fused_indices_dimension *= i
fused_updates_dimension = 1
for i in updates_ptr.shape[len(indices_ptr.shape) - 1 :]:
fused_updates_dimension *= i
fused_shape = 1
for i in data_ptr.shape:
fused_shape *= i
max_threads = int(tvm.target.Target.current(allow_none=False).attrs["max_num_threads"])
with IRBuilder() as ib:
with T.seq_scope():
# Init
nthread_bx_init = cast(ceil_div(fused_shape, max_threads), "int32")
tx_init = te.thread_axis("threadIdx.x")
bx_init = te.thread_axis("blockIdx.x")
with T.frame_scope(
[
T.attr(bx_init, "thread_extent", nthread_bx_init),
T.attr(tx_init, "thread_extent", max_threads),
]
):
tid = bx_init * max_threads + tx_init
with T.If(tid < fused_shape):
with T.Then():
out[tid] = data[tid]
# Scatter
nthread_bx_scat = cast(ceil_div(fused_updates_dimension, max_threads), "int32")
tx_scat = te.thread_axis("threadIdx.x")
bx_scat = te.thread_axis("blockIdx.x")
with T.frame_scope(
[
T.attr(bx_scat, "thread_extent", nthread_bx_scat),
T.attr(tx_scat, "thread_extent", max_threads),
]
):
j = bx_scat * max_threads + tx_scat
with T.If(j < fused_updates_dimension):
with T.Then():
with T.serial(0, fused_indices_dimension) as i:
offset = fused_updates_dimension
index = j # x_M, .. x_{N-1} part of the index into out.
# Build up the indices[0, y_0, ..], ..,
# indices[M-1, y_0, ..] part of the index into out.
for l in reversed(range(indices_ptr.shape[0].value)):
# indices[l, y_0, ... y_{k-1}]
index += offset * indices[i + l * fused_indices_dimension]
offset *= data_ptr.shape[l]
if mode == "update":
out[index] = updates[i * fused_updates_dimension + j]
elif mode == "add":
out[index] += updates[i * fused_updates_dimension + j]
elif mode == "mul":
out[index] *= updates[i * fused_updates_dimension + j]
elif mode == "min":
out[index] = tirx.min(
out[index], updates[i * fused_updates_dimension + j]
)
elif mode == "max":
out[index] = tirx.max(
out[index], updates[i * fused_updates_dimension + j]
)
else:
raise NotImplementedError(
"scatter_nd mode not in [update, add, mul, min, max]:",
mode,
)
return ib.get()
out_buf = tirx.decl_buffer(data.shape, data.dtype, "out_buf", layout=None)
return te.extern(
[data.shape],
[data, indices, updates],
lambda ins, outs: gen_ir(ins[0], ins[1], ins[2], outs[0]),
dtype=data.dtype,
out_buffers=[out_buf],
name="scatter_nd.gpu",
tag="scatter_nd.gpu",
)
python/tvm/topi/gpu/sort.py
+1171
View File
File diff suppressed because it is too large Load Diff