This is an automated email from the ASF dual-hosted git repository.
tqchen pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/main by this push:
new 8b60e7f49b [REFACTOR][TIRx] Collect root statements without seq_scope
(#20564)
8b60e7f49b is described below
commit 8b60e7f49b80c9668352f27c466b9b95403cfcbc
Author: Tianqi Chen <[email protected]>
AuthorDate: Tue Oct 6 12:17:11 2026 -0400
[REFACTOR][TIRx] Collect root statements without seq_scope (#20564)
---
python/tvm/relax/transform/legalize_ops/grad.py | 88 +++---
python/tvm/tirx/script/ir_builder/__init__.py | 4 -
python/tvm/tirx/script/ir_builder/utils.py | 42 +--
python/tvm/topi/gpu/scatter_elements.py | 83 +++---
python/tvm/topi/gpu/scatter_nd.py | 201 +++++++-------
python/tvm/topi/index_put.py | 81 +++---
python/tvm/topi/scatter.py | 129 ++++-----
python/tvm/topi/scatter_elements.py | 35 ++-
python/tvm/topi/sparse_reshape.py | 303 ++++++++++-----------
python/tvm/topi/vision/nms.py | 67 +++--
src/tirx/script/ir_builder/utils.h | 27 +-
.../relax/test_transform_legalize_ops_grad.py | 2 -
.../test_transform_legalize_ops_manipulate.py | 3 -
13 files changed, 502 insertions(+), 563 deletions(-)
diff --git a/python/tvm/relax/transform/legalize_ops/grad.py
b/python/tvm/relax/transform/legalize_ops/grad.py
index 5f43d06a1c..a37829e3a9 100644
--- a/python/tvm/relax/transform/legalize_ops/grad.py
+++ b/python/tvm/relax/transform/legalize_ops/grad.py
@@ -178,56 +178,52 @@ def _grad_take_backward(bb: BlockBuilder, call: Call) ->
Expr:
indices_len = indices_ptr.shape[0]
with IRBuilder() as ib:
- with T.seq_scope():
- # Init loop (zero-fill output buffer)
- with T.serial(fused_shape) as i:
- T.buffer_store(
- out, tirx.const(0, dtype=x_ptr.dtype),
T.tensor_indices(out, i)
- )
-
- # Accumulation loop
- if axis is not None:
- fused_output_grad_shape_pre = 1
- fused_output_grad_shape_nxt = 1
- for i in range(len(output_grad_ptr.shape)):
- if i < axis:
- fused_output_grad_shape_pre *=
output_grad_ptr.shape[i]
- elif i > axis:
- fused_output_grad_shape_nxt *=
output_grad_ptr.shape[i]
-
- x_axis_len = x_ptr.shape[axis]
-
- with T.serial(
- fused_output_grad_shape_pre *
fused_output_grad_shape_nxt
- ) as fused:
- i = fused // fused_output_grad_shape_nxt
- j = fused % fused_output_grad_shape_nxt
- with T.serial(indices_len) as loop_l:
- out_idx = (
- i * fused_output_grad_shape_nxt *
x_axis_len
- + idx[T.tensor_indices(idx, loop_l)]
- * fused_output_grad_shape_nxt
- + j
- )
- grad_idx = (
- i * fused_output_grad_shape_nxt *
indices_len
- + loop_l * fused_output_grad_shape_nxt
- + j
- )
- T.buffer_store(
- out,
- out[T.tensor_indices(out, out_idx)]
- + grad[T.tensor_indices(grad, grad_idx)],
- T.tensor_indices(out, out_idx),
- )
- else:
+ # Init loop (zero-fill output buffer)
+ with T.serial(fused_shape) as i:
+ T.buffer_store(out, tirx.const(0, dtype=x_ptr.dtype),
T.tensor_indices(out, i))
+
+ # Accumulation loop
+ if axis is not None:
+ fused_output_grad_shape_pre = 1
+ fused_output_grad_shape_nxt = 1
+ for i in range(len(output_grad_ptr.shape)):
+ if i < axis:
+ fused_output_grad_shape_pre *=
output_grad_ptr.shape[i]
+ elif i > axis:
+ fused_output_grad_shape_nxt *=
output_grad_ptr.shape[i]
+
+ x_axis_len = x_ptr.shape[axis]
+
+ with T.serial(
+ fused_output_grad_shape_pre *
fused_output_grad_shape_nxt
+ ) as fused:
+ i = fused // fused_output_grad_shape_nxt
+ j = fused % fused_output_grad_shape_nxt
with T.serial(indices_len) as loop_l:
+ out_idx = (
+ i * fused_output_grad_shape_nxt * x_axis_len
+ + idx[T.tensor_indices(idx, loop_l)] *
fused_output_grad_shape_nxt
+ + j
+ )
+ grad_idx = (
+ i * fused_output_grad_shape_nxt * indices_len
+ + loop_l * fused_output_grad_shape_nxt
+ + j
+ )
T.buffer_store(
out,
- out[T.tensor_indices(out,
(idx[T.tensor_indices(idx, loop_l)]))]
- + grad[T.tensor_indices(grad, loop_l)],
- T.tensor_indices(out,
(idx[T.tensor_indices(idx, loop_l)])),
+ out[T.tensor_indices(out, out_idx)]
+ + grad[T.tensor_indices(grad, grad_idx)],
+ T.tensor_indices(out, out_idx),
)
+ else:
+ with T.serial(indices_len) as loop_l:
+ T.buffer_store(
+ out,
+ out[T.tensor_indices(out,
(idx[T.tensor_indices(idx, loop_l)]))]
+ + grad[T.tensor_indices(grad, loop_l)],
+ T.tensor_indices(out, (idx[T.tensor_indices(idx,
loop_l)])),
+ )
return ib.get()
diff --git a/python/tvm/tirx/script/ir_builder/__init__.py
b/python/tvm/tirx/script/ir_builder/__init__.py
index 5a3f5648c5..39673bcba4 100644
--- a/python/tvm/tirx/script/ir_builder/__init__.py
+++ b/python/tvm/tirx/script/ir_builder/__init__.py
@@ -53,9 +53,6 @@ from .tirx import (
from .utils import (
frame_scope as frame_scope,
)
-from .utils import (
- seq_scope as seq_scope,
-)
from .utils import (
tensor_indices as tensor_indices,
)
@@ -81,6 +78,5 @@ __all__ = [
"wg",
"tensor_indices",
"frame_scope",
- "seq_scope",
"supports_mutable_declarations",
]
diff --git a/python/tvm/tirx/script/ir_builder/utils.py
b/python/tvm/tirx/script/ir_builder/utils.py
index 00bc41fc2d..1eb9b64621 100644
--- a/python/tvm/tirx/script/ir_builder/utils.py
+++ b/python/tvm/tirx/script/ir_builder/utils.py
@@ -18,10 +18,9 @@
import contextlib
-from tvm.ir import StringImm
from tvm.tirx import Var
-from . import _ffi_api, frame
+from . import frame
class _FrameScope:
@@ -71,45 +70,6 @@ def frame_scope(frames: list[frame.TIRFrame]) -> _FrameScope:
return _FrameScope(frames)
-def seq_scope():
- """Create a scope that allows multiple consecutive statements.
-
- The IRBuilder requires a parent frame when having multiple consecutive
- top-level statements (e.g., multiple loops). This function creates a
- dummy attr frame that serves as a parent scope.
-
- Returns
- -------
- frame.AttrFrame
- A dummy attribute frame that wraps multiple statements.
-
- Examples
- --------
- Without seq_scope, multiple consecutive loops fail:
-
- .. code-block:: python
-
- with IRBuilder() as ib:
- with T.serial(0, 10) as i:
- T.evaluate(i)
- with T.serial(0, 5) as j: # This would fail!
- T.evaluate(j)
-
- With seq_scope, multiple consecutive statements work:
-
- .. code-block:: python
-
- with IRBuilder() as ib:
- with seq_scope():
- with T.serial(0, 10) as i:
- T.evaluate(i)
- with T.serial(0, 5) as j:
- T.evaluate(j)
- result = ib.get()
- """
- return _ffi_api.Attr(0, "pragma_scope", StringImm("seq"))
-
-
def tensor_indices(buffer: Var, index):
"""Translate logical flat or multidimensional indices for a concrete
buffer.
diff --git a/python/tvm/topi/gpu/scatter_elements.py
b/python/tvm/topi/gpu/scatter_elements.py
index 5086a36888..514e84f732 100644
--- a/python/tvm/topi/gpu/scatter_elements.py
+++ b/python/tvm/topi/gpu/scatter_elements.py
@@ -69,48 +69,47 @@ def scatter_elements(data, indices, updates, axis=0,
reduction="update"):
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")
- with T.frame_scope(
- [
- T.launch_thread("blockIdx.x", nthread_bx_init),
- T.launch_thread("threadIdx.x", max_threads),
- ]
- ) as (bx_init, tx_init):
- tid = bx_init * max_threads + tx_init
- with T.if_(tid < full_range):
- with T.then_():
- T.buffer_store(
- out,
- data[T.tensor_indices(data, tid)],
- T.tensor_indices(out, tid),
- )
-
- # Scatter
- nthread_bx_scat = cast(ceil_div(ind_full_range_excl_axis,
max_threads), "int32")
- with T.frame_scope(
- [
- T.launch_thread("blockIdx.x", nthread_bx_scat),
- T.launch_thread("threadIdx.x", max_threads),
- ]
- ) as (bx_scat, tx_scat):
- 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[T.tensor_indices(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[T.tensor_indices(updates, index1)])
+ # Init
+ nthread_bx_init = cast(ceil_div(full_range, max_threads), "int32")
+ with T.frame_scope(
+ [
+ T.launch_thread("blockIdx.x", nthread_bx_init),
+ T.launch_thread("threadIdx.x", max_threads),
+ ]
+ ) as (bx_init, tx_init):
+ tid = bx_init * max_threads + tx_init
+ with T.if_(tid < full_range):
+ with T.then_():
+ T.buffer_store(
+ out,
+ data[T.tensor_indices(data, tid)],
+ T.tensor_indices(out, tid),
+ )
+
+ # Scatter
+ nthread_bx_scat = cast(ceil_div(ind_full_range_excl_axis,
max_threads), "int32")
+ with T.frame_scope(
+ [
+ T.launch_thread("blockIdx.x", nthread_bx_scat),
+ T.launch_thread("threadIdx.x", max_threads),
+ ]
+ ) as (bx_scat, tx_scat):
+ 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[T.tensor_indices(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[T.tensor_indices(updates, index1)])
return ib.get()
diff --git a/python/tvm/topi/gpu/scatter_nd.py
b/python/tvm/topi/gpu/scatter_nd.py
index 99197b4c80..f9db776650 100644
--- a/python/tvm/topi/gpu/scatter_nd.py
+++ b/python/tvm/topi/gpu/scatter_nd.py
@@ -57,118 +57,113 @@ def scatter_nd(data, indices, updates, mode):
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")
- with T.frame_scope(
- [
- T.launch_thread("blockIdx.x", nthread_bx_init),
- T.launch_thread("threadIdx.x", max_threads),
- ]
- ) as (bx_init, tx_init):
- tid = bx_init * max_threads + tx_init
- with T.if_(tid < fused_shape):
- with T.then_():
- T.buffer_store(
- out,
- data[T.tensor_indices(data, tid)],
- T.tensor_indices(out, tid),
- )
+ # Init
+ nthread_bx_init = cast(ceil_div(fused_shape, max_threads), "int32")
+ with T.frame_scope(
+ [
+ T.launch_thread("blockIdx.x", nthread_bx_init),
+ T.launch_thread("threadIdx.x", max_threads),
+ ]
+ ) as (bx_init, tx_init):
+ tid = bx_init * max_threads + tx_init
+ with T.if_(tid < fused_shape):
+ with T.then_():
+ T.buffer_store(
+ out,
+ data[T.tensor_indices(data, tid)],
+ T.tensor_indices(out, tid),
+ )
- # Scatter
- nthread_bx_scat = cast(ceil_div(fused_updates_dimension,
max_threads), "int32")
- with T.frame_scope(
- [
- T.launch_thread("blockIdx.x", nthread_bx_scat),
- T.launch_thread("threadIdx.x", max_threads),
- ]
- ) as (bx_scat, tx_scat):
- 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[
+ # Scatter
+ nthread_bx_scat = cast(ceil_div(fused_updates_dimension,
max_threads), "int32")
+ with T.frame_scope(
+ [
+ T.launch_thread("blockIdx.x", nthread_bx_scat),
+ T.launch_thread("threadIdx.x", max_threads),
+ ]
+ ) as (bx_scat, tx_scat):
+ 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[
+ T.tensor_indices(indices, (i + l *
fused_indices_dimension))
+ ]
+ )
+ offset *= data_ptr.shape[l]
+ if mode == "update":
+ T.buffer_store(
+ out,
+ updates[
+ T.tensor_indices(updates, (i *
fused_updates_dimension + j))
+ ],
+ T.tensor_indices(out, index),
+ )
+ elif mode == "add":
+ T.buffer_store(
+ out,
+ out[T.tensor_indices(out, index)]
+ + (
+ updates[
+ T.tensor_indices(
+ updates, (i *
fused_updates_dimension + j)
+ )
+ ]
+ ),
+ T.tensor_indices(out, index),
+ )
+ elif mode == "mul":
+ T.buffer_store(
+ out,
+ out[T.tensor_indices(out, index)]
+ * (
+ updates[
T.tensor_indices(
- indices, (i + l *
fused_indices_dimension)
+ updates, (i *
fused_updates_dimension + j)
)
]
- )
- offset *= data_ptr.shape[l]
- if mode == "update":
- T.buffer_store(
- out,
+ ),
+ T.tensor_indices(out, index),
+ )
+ elif mode == "min":
+ T.buffer_store(
+ out,
+ tirx.min(
+ out[T.tensor_indices(out, index)],
+ updates[
+ T.tensor_indices(
+ updates, (i *
fused_updates_dimension + j)
+ )
+ ],
+ ),
+ T.tensor_indices(out, index),
+ )
+ elif mode == "max":
+ T.buffer_store(
+ out,
+ tirx.max(
+ out[T.tensor_indices(out, index)],
updates[
T.tensor_indices(
updates, (i *
fused_updates_dimension + j)
)
],
- T.tensor_indices(out, index),
- )
- elif mode == "add":
- T.buffer_store(
- out,
- out[T.tensor_indices(out, index)]
- + (
- updates[
- T.tensor_indices(
- updates, (i *
fused_updates_dimension + j)
- )
- ]
- ),
- T.tensor_indices(out, index),
- )
- elif mode == "mul":
- T.buffer_store(
- out,
- out[T.tensor_indices(out, index)]
- * (
- updates[
- T.tensor_indices(
- updates, (i *
fused_updates_dimension + j)
- )
- ]
- ),
- T.tensor_indices(out, index),
- )
- elif mode == "min":
- T.buffer_store(
- out,
- tirx.min(
- out[T.tensor_indices(out, index)],
- updates[
- T.tensor_indices(
- updates, (i *
fused_updates_dimension + j)
- )
- ],
- ),
- T.tensor_indices(out, index),
- )
- elif mode == "max":
- T.buffer_store(
- out,
- tirx.max(
- out[T.tensor_indices(out, index)],
- updates[
- T.tensor_indices(
- updates, (i *
fused_updates_dimension + j)
- )
- ],
- ),
- T.tensor_indices(out, index),
- )
- else:
- raise NotImplementedError(
- "scatter_nd mode not in [update, add,
mul, min, max]:",
- mode,
- )
+ ),
+ T.tensor_indices(out, index),
+ )
+ else:
+ raise NotImplementedError(
+ "scatter_nd mode not in [update, add, mul,
min, max]:",
+ mode,
+ )
return ib.get()
diff --git a/python/tvm/topi/index_put.py b/python/tvm/topi/index_put.py
index 9701aa1c65..1c4f6c4f28 100644
--- a/python/tvm/topi/index_put.py
+++ b/python/tvm/topi/index_put.py
@@ -95,48 +95,47 @@ def index_put(data, indices, values, accumulate=False):
out = out_ptr
with IRBuilder() as ib:
- with T.seq_scope():
- with T.parallel(0, full_range) as i:
- T.buffer_store(out, data[T.tensor_indices(data, i)],
T.tensor_indices(out, i))
-
- with T.parallel(0, index_len) as k:
- # Decompose k into multi-dimensional broadcast index
- k_temp = k
- broadcast_indices = []
+ with T.parallel(0, full_range) as i:
+ T.buffer_store(out, data[T.tensor_indices(data, i)],
T.tensor_indices(out, i))
+
+ with T.parallel(0, index_len) as k:
+ # Decompose k into multi-dimensional broadcast index
+ k_temp = k
+ broadcast_indices = []
+ for i in range(broadcast_ndim - 1, -1, -1):
+ broadcast_indices.insert(0, k_temp % broadcast_shape[i])
+ k_temp = k_temp // broadcast_shape[i]
+
+ flat_index = 0
+ stride = 1
+ for dim in range(len(shape) - 1, -1, -1):
+ # Get the index for this dimension using broadcasting
+ idx_shape = index_shapes[dim]
+ idx_ndim = len(idx_shape)
+
+ # Compute the linear index into this index tensor
+ idx_offset = 0
+ idx_stride = 1
for i in range(broadcast_ndim - 1, -1, -1):
- broadcast_indices.insert(0, k_temp %
broadcast_shape[i])
- k_temp = k_temp // broadcast_shape[i]
-
- flat_index = 0
- stride = 1
- for dim in range(len(shape) - 1, -1, -1):
- # Get the index for this dimension using broadcasting
- idx_shape = index_shapes[dim]
- idx_ndim = len(idx_shape)
-
- # Compute the linear index into this index tensor
- idx_offset = 0
- idx_stride = 1
- for i in range(broadcast_ndim - 1, -1, -1):
- # Right-align the index shape with broadcast shape
- dim_idx = idx_ndim - broadcast_ndim + i
- if dim_idx >= 0:
- dim_size = idx_shape[dim_idx]
- # Use broadcasting: if size is 1, use index 0
- # otherwise use broadcast_indices[i]
- if utils.equal_const_int(dim_size, 1):
- idx_in_dim = 0
- else:
- idx_in_dim = broadcast_indices[i]
- idx_offset += idx_in_dim * idx_stride
- idx_stride *= dim_size
-
- idx_val = indices[dim][T.tensor_indices(indices[dim],
(idx_offset))]
- shifted_idx = idx_val + (idx_val < 0) * shape[dim]
- flat_index += shifted_idx * stride
- stride *= shape[dim]
-
- reduce_func(out, flat_index,
values[T.tensor_indices(values, k)])
+ # Right-align the index shape with broadcast shape
+ dim_idx = idx_ndim - broadcast_ndim + i
+ if dim_idx >= 0:
+ dim_size = idx_shape[dim_idx]
+ # Use broadcasting: if size is 1, use index 0
+ # otherwise use broadcast_indices[i]
+ if utils.equal_const_int(dim_size, 1):
+ idx_in_dim = 0
+ else:
+ idx_in_dim = broadcast_indices[i]
+ idx_offset += idx_in_dim * idx_stride
+ idx_stride *= dim_size
+
+ idx_val = indices[dim][T.tensor_indices(indices[dim],
(idx_offset))]
+ shifted_idx = idx_val + (idx_val < 0) * shape[dim]
+ flat_index += shifted_idx * stride
+ stride *= shape[dim]
+
+ reduce_func(out, flat_index, values[T.tensor_indices(values,
k)])
return ib.get()
diff --git a/python/tvm/topi/scatter.py b/python/tvm/topi/scatter.py
index cf56d1474a..08cb1e5758 100644
--- a/python/tvm/topi/scatter.py
+++ b/python/tvm/topi/scatter.py
@@ -117,79 +117,64 @@ def scatter_nd(data, indices, updates, mode):
fused_shape *= i
with IRBuilder() as ib:
- with T.seq_scope():
- with T.serial(0, fused_shape) as i:
- T.buffer_store(out, data[T.tensor_indices(data, i)],
T.tensor_indices(out, i))
-
- with T.serial(0, fused_indices_dimension) as i:
- with T.parallel(0, fused_updates_dimension) as j:
- offset = fused_updates_dimension
- index = j # This is 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[
- T.tensor_indices(indices, i + l *
fused_indices_dimension)
- ]
- )
- offset *= data_ptr.shape[l]
- if mode == "update":
- T.buffer_store(
- out,
+ with T.serial(0, fused_shape) as i:
+ T.buffer_store(out, data[T.tensor_indices(data, i)],
T.tensor_indices(out, i))
+
+ with T.serial(0, fused_indices_dimension) as i:
+ with T.parallel(0, fused_updates_dimension) as j:
+ offset = fused_updates_dimension
+ index = j # This is 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[T.tensor_indices(indices, i + l *
fused_indices_dimension)]
+ )
+ offset *= data_ptr.shape[l]
+ if mode == "update":
+ T.buffer_store(
+ out,
+ updates[T.tensor_indices(updates, i *
fused_updates_dimension + j)],
+ T.tensor_indices(out, index),
+ )
+ elif mode == "add":
+ T.buffer_store(
+ out,
+ out[T.tensor_indices(out, index)]
+ + (updates[T.tensor_indices(updates, i *
fused_updates_dimension + j)]),
+ T.tensor_indices(out, index),
+ )
+ elif mode == "mul":
+ T.buffer_store(
+ out,
+ out[T.tensor_indices(out, index)]
+ * (updates[T.tensor_indices(updates, i *
fused_updates_dimension + j)]),
+ T.tensor_indices(out, index),
+ )
+ elif mode == "min":
+ T.buffer_store(
+ out,
+ tirx.min(
+ out[T.tensor_indices(out, index)],
updates[T.tensor_indices(updates, i *
fused_updates_dimension + j)],
- T.tensor_indices(out, index),
- )
- elif mode == "add":
- T.buffer_store(
- out,
- out[T.tensor_indices(out, index)]
- + (
- updates[
- T.tensor_indices(updates, i *
fused_updates_dimension + j)
- ]
- ),
- T.tensor_indices(out, index),
- )
- elif mode == "mul":
- T.buffer_store(
- out,
- out[T.tensor_indices(out, index)]
- * (
- updates[
- T.tensor_indices(updates, i *
fused_updates_dimension + j)
- ]
- ),
- T.tensor_indices(out, index),
- )
- elif mode == "min":
- T.buffer_store(
- out,
- tirx.min(
- out[T.tensor_indices(out, index)],
- updates[
- T.tensor_indices(updates, i *
fused_updates_dimension + j)
- ],
- ),
- T.tensor_indices(out, index),
- )
- elif mode == "max":
- T.buffer_store(
- out,
- tirx.max(
- out[T.tensor_indices(out, index)],
- updates[
- T.tensor_indices(updates, i *
fused_updates_dimension + j)
- ],
- ),
- T.tensor_indices(out, index),
- )
- else:
- raise NotImplementedError(
- "scatter_nd mode not in [update, add, mul,
min, max]:", mode
- )
+ ),
+ T.tensor_indices(out, index),
+ )
+ elif mode == "max":
+ T.buffer_store(
+ out,
+ tirx.max(
+ out[T.tensor_indices(out, index)],
+ updates[T.tensor_indices(updates, i *
fused_updates_dimension + j)],
+ ),
+ T.tensor_indices(out, index),
+ )
+ else:
+ raise NotImplementedError(
+ "scatter_nd mode not in [update, add, mul, min,
max]:", mode
+ )
return ib.get()
diff --git a/python/tvm/topi/scatter_elements.py
b/python/tvm/topi/scatter_elements.py
index 432c87abd7..4f88dcc4aa 100644
--- a/python/tvm/topi/scatter_elements.py
+++ b/python/tvm/topi/scatter_elements.py
@@ -105,24 +105,23 @@ def scatter_elements(data, indices, updates, axis=0,
reduction="update"):
# Copy initial input data to output
with IRBuilder() as ib:
- with T.seq_scope():
- with T.parallel(0, full_range) as i:
- T.buffer_store(out, data[T.tensor_indices(data, i)],
T.tensor_indices(out, i))
-
- with T.parallel(0, ind_before_axis_range *
ind_after_axis_range) as fused:
- 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[T.tensor_indices(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[T.tensor_indices(updates, index1)])
+ with T.parallel(0, full_range) as i:
+ T.buffer_store(out, data[T.tensor_indices(data, i)],
T.tensor_indices(out, i))
+
+ with T.parallel(0, ind_before_axis_range * ind_after_axis_range)
as fused:
+ 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[T.tensor_indices(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[T.tensor_indices(updates,
index1)])
return ib.get()
diff --git a/python/tvm/topi/sparse_reshape.py
b/python/tvm/topi/sparse_reshape.py
index a65e54241a..1024ec5760 100644
--- a/python/tvm/topi/sparse_reshape.py
+++ b/python/tvm/topi/sparse_reshape.py
@@ -79,192 +79,185 @@ def sparse_reshape(
out_new_shape_ptr,
):
with IRBuilder() as ib:
- with T.seq_scope():
- sparse_indices = sparse_indices_ptr
- prev_shape = prev_shape_ptr
+ sparse_indices = sparse_indices_ptr
+ prev_shape = prev_shape_ptr
- new_shape = new_shape_ptr
- out_new_shape = out_new_shape_ptr
- new_sparse_indices = new_sparse_indices_ptr
+ new_shape = new_shape_ptr
+ out_new_shape = out_new_shape_ptr
+ new_sparse_indices = new_sparse_indices_ptr
- prev_shape_size = prev_shape_ptr.shape[0]
- new_shape_size = new_shape_ptr.shape[0]
+ prev_shape_size = prev_shape_ptr.shape[0]
+ new_shape_size = new_shape_ptr.shape[0]
- multipliers_buf = T.alloc_tensor(
- [prev_shape_size], new_shape_ptr.dtype, scope="local"
- )
- multipliers = multipliers_buf
- dividers_buf = T.alloc_tensor([new_shape_size],
new_shape_ptr.dtype, scope="local")
- dividers = dividers_buf
- flattened_indices_buf = T.alloc_tensor(
- [sparse_indices_ptr.shape[0]], new_shape_ptr.dtype,
scope="local"
- )
- flattened_indices = flattened_indices_buf
- total_ele_buf = T.alloc_tensor([1], new_shape_ptr.dtype,
scope="local")
- total_ele = total_ele_buf
- division_total_ele_buf = T.alloc_tensor([1],
new_shape_ptr.dtype, scope="local")
- division_total_ele = division_total_ele_buf
- equal_shape_buf = T.alloc_tensor([1], "bool", scope="local")
- equal_shape = equal_shape_buf
+ multipliers_buf = T.alloc_tensor([prev_shape_size],
new_shape_ptr.dtype, scope="local")
+ multipliers = multipliers_buf
+ dividers_buf = T.alloc_tensor([new_shape_size],
new_shape_ptr.dtype, scope="local")
+ dividers = dividers_buf
+ flattened_indices_buf = T.alloc_tensor(
+ [sparse_indices_ptr.shape[0]], new_shape_ptr.dtype,
scope="local"
+ )
+ flattened_indices = flattened_indices_buf
+ total_ele_buf = T.alloc_tensor([1], new_shape_ptr.dtype,
scope="local")
+ total_ele = total_ele_buf
+ division_total_ele_buf = T.alloc_tensor([1], new_shape_ptr.dtype,
scope="local")
+ division_total_ele = division_total_ele_buf
+ equal_shape_buf = T.alloc_tensor([1], "bool", scope="local")
+ equal_shape = equal_shape_buf
- T.buffer_store(
- total_ele,
- prev_shape[T.tensor_indices(prev_shape, 0)],
- T.tensor_indices(total_ele, 0),
- )
+ T.buffer_store(
+ total_ele,
+ prev_shape[T.tensor_indices(prev_shape, 0)],
+ T.tensor_indices(total_ele, 0),
+ )
- # Cumulative Reverse Exclusive Multiply
+ # Cumulative Reverse Exclusive Multiply
+ T.buffer_store(
+ multipliers,
+ Cast(new_shape_ptr.dtype, 1),
+ T.tensor_indices(multipliers, prev_shape_size - 1),
+ )
+ with T.serial(0, prev_shape_size - 1) as i_:
+ i = i_ + 1
T.buffer_store(
multipliers,
- Cast(new_shape_ptr.dtype, 1),
- T.tensor_indices(multipliers, prev_shape_size - 1),
+ prev_shape[T.tensor_indices(prev_shape, prev_shape_size -
i)]
+ * multipliers[T.tensor_indices(multipliers,
prev_shape_size - i)],
+ T.tensor_indices(multipliers, prev_shape_size - 1 - i),
)
- with T.serial(0, prev_shape_size - 1) as i_:
- i = i_ + 1
- T.buffer_store(
- multipliers,
- prev_shape[T.tensor_indices(prev_shape,
prev_shape_size - i)]
- * multipliers[T.tensor_indices(multipliers,
prev_shape_size - i)],
- T.tensor_indices(multipliers, prev_shape_size - 1 - i),
- )
- T.buffer_store(
- total_ele,
- total_ele[T.tensor_indices(total_ele, 0)]
- * (prev_shape[T.tensor_indices(prev_shape,
prev_shape_size - i)]),
- T.tensor_indices(total_ele, 0),
- )
-
T.buffer_store(
- division_total_ele,
- Cast(new_shape_ptr.dtype, 1),
- T.tensor_indices(division_total_ele, 0),
+ total_ele,
+ total_ele[T.tensor_indices(total_ele, 0)]
+ * (prev_shape[T.tensor_indices(prev_shape, prev_shape_size
- i)]),
+ T.tensor_indices(total_ele, 0),
)
- with T.serial(0, new_shape_size) as i:
- with T.if_(new_shape[T.tensor_indices(new_shape, i)] !=
-1):
- with T.then_():
- T.buffer_store(
- division_total_ele,
-
division_total_ele[T.tensor_indices(division_total_ele, 0)]
- * (new_shape[T.tensor_indices(new_shape, i)]),
- T.tensor_indices(division_total_ele, 0),
- )
- # Compute true output shape (replace negative ones)
- with T.serial(0, new_shape_size) as i:
- with T.if_(new_shape[T.tensor_indices(new_shape, i)] ==
-1):
- with T.then_():
- T.buffer_store(
- out_new_shape,
- Cast(
- new_shape_ptr.dtype,
- div(
- total_ele[T.tensor_indices(total_ele,
0)],
-
division_total_ele[T.tensor_indices(division_total_ele, 0)],
- ),
- ),
- T.tensor_indices(out_new_shape, i),
- )
- with T.else_():
- T.buffer_store(
- out_new_shape,
- new_shape[T.tensor_indices(new_shape, i)],
- T.tensor_indices(out_new_shape, i),
- )
-
- # Check if prev_shape and new_shape are equal
- T.buffer_store(equal_shape, True,
T.tensor_indices(equal_shape, 0))
- with T.if_(prev_shape_size == new_shape_size):
+ T.buffer_store(
+ division_total_ele,
+ Cast(new_shape_ptr.dtype, 1),
+ T.tensor_indices(division_total_ele, 0),
+ )
+ with T.serial(0, new_shape_size) as i:
+ with T.if_(new_shape[T.tensor_indices(new_shape, i)] != -1):
with T.then_():
- with T.serial(0, prev_shape_size) as i:
- with T.if_(
- prev_shape[T.tensor_indices(prev_shape, i)]
- !=
out_new_shape[T.tensor_indices(out_new_shape, i)]
- ):
- with T.then_():
- T.buffer_store(
- equal_shape, False,
T.tensor_indices(equal_shape, 0)
- )
- with T.else_():
- T.buffer_store(equal_shape, False,
T.tensor_indices(equal_shape, 0))
+ T.buffer_store(
+ division_total_ele,
+
division_total_ele[T.tensor_indices(division_total_ele, 0)]
+ * (new_shape[T.tensor_indices(new_shape, i)]),
+ T.tensor_indices(division_total_ele, 0),
+ )
- # Return same inputs if shapes are equal
- with T.if_(equal_shape[T.tensor_indices(equal_shape, 0)]):
+ # Compute true output shape (replace negative ones)
+ with T.serial(0, new_shape_size) as i:
+ with T.if_(new_shape[T.tensor_indices(new_shape, i)] == -1):
with T.then_():
- with T.parallel(0, sparse_indices_ptr.shape[0]) as i:
- with T.serial(0, sparse_indices_ptr.shape[1]) as j:
- T.buffer_store(
- new_sparse_indices,
- sparse_indices[(i, j)],
- (i, j),
- )
-
- # Else compute new_sparse_indices
+ T.buffer_store(
+ out_new_shape,
+ Cast(
+ new_shape_ptr.dtype,
+ div(
+ total_ele[T.tensor_indices(total_ele, 0)],
+
division_total_ele[T.tensor_indices(division_total_ele, 0)],
+ ),
+ ),
+ T.tensor_indices(out_new_shape, i),
+ )
with T.else_():
T.buffer_store(
- dividers,
- Cast(new_shape_ptr.dtype, 1),
- T.tensor_indices(dividers, new_shape_size - 1),
+ out_new_shape,
+ new_shape[T.tensor_indices(new_shape, i)],
+ T.tensor_indices(out_new_shape, i),
)
- with T.serial(0, new_shape_size - 1) as i_:
- i = i_ + 1
+
+ # Check if prev_shape and new_shape are equal
+ T.buffer_store(equal_shape, True, T.tensor_indices(equal_shape, 0))
+ with T.if_(prev_shape_size == new_shape_size):
+ with T.then_():
+ with T.serial(0, prev_shape_size) as i:
+ with T.if_(
+ prev_shape[T.tensor_indices(prev_shape, i)]
+ != out_new_shape[T.tensor_indices(out_new_shape,
i)]
+ ):
+ with T.then_():
+ T.buffer_store(equal_shape, False,
T.tensor_indices(equal_shape, 0))
+ with T.else_():
+ T.buffer_store(equal_shape, False,
T.tensor_indices(equal_shape, 0))
+
+ # Return same inputs if shapes are equal
+ with T.if_(equal_shape[T.tensor_indices(equal_shape, 0)]):
+ with T.then_():
+ with T.parallel(0, sparse_indices_ptr.shape[0]) as i:
+ with T.serial(0, sparse_indices_ptr.shape[1]) as j:
T.buffer_store(
- dividers,
- dividers[T.tensor_indices(dividers,
new_shape_size - i)]
- * out_new_shape[
- T.tensor_indices(out_new_shape,
new_shape_size - i)
- ],
- T.tensor_indices(dividers, new_shape_size - 1
- i),
+ new_sparse_indices,
+ sparse_indices[(i, j)],
+ (i, j),
)
- with T.parallel(0, sparse_indices_ptr.shape[0]) as i:
+ # Else compute new_sparse_indices
+ with T.else_():
+ T.buffer_store(
+ dividers,
+ Cast(new_shape_ptr.dtype, 1),
+ T.tensor_indices(dividers, new_shape_size - 1),
+ )
+ with T.serial(0, new_shape_size - 1) as i_:
+ i = i_ + 1
+ T.buffer_store(
+ dividers,
+ dividers[T.tensor_indices(dividers, new_shape_size
- i)]
+ * out_new_shape[T.tensor_indices(out_new_shape,
new_shape_size - i)],
+ T.tensor_indices(dividers, new_shape_size - 1 - i),
+ )
+
+ with T.parallel(0, sparse_indices_ptr.shape[0]) as i:
+ T.buffer_store(
+ flattened_indices,
+ Cast(new_shape_ptr.dtype, 0),
+ T.tensor_indices(flattened_indices, i),
+ )
+ with T.serial(0, sparse_indices_ptr.shape[1]) as j:
T.buffer_store(
flattened_indices,
- Cast(new_shape_ptr.dtype, 0),
+
flattened_indices[T.tensor_indices(flattened_indices, i)]
+ + (
+ sparse_indices[(i, j)]
+ *
multipliers[T.tensor_indices(multipliers, j)]
+ ),
T.tensor_indices(flattened_indices, i),
)
- with T.serial(0, sparse_indices_ptr.shape[1]) as j:
- T.buffer_store(
- flattened_indices,
-
flattened_indices[T.tensor_indices(flattened_indices, i)]
- + (
- sparse_indices[(i, j)]
- *
multipliers[T.tensor_indices(multipliers, j)]
- ),
- T.tensor_indices(flattened_indices, i),
- )
- with T.parallel(0, new_sparse_indices_ptr.shape[0]) as
i:
- current_element_buf = T.alloc_tensor(
- [1], new_shape_ptr.dtype, scope="local"
+ with T.parallel(0, new_sparse_indices_ptr.shape[0]) as i:
+ current_element_buf = T.alloc_tensor(
+ [1], new_shape_ptr.dtype, scope="local"
+ )
+ current_element = current_element_buf
+ T.buffer_store(
+ current_element,
+
flattened_indices[T.tensor_indices(flattened_indices, i)],
+ T.tensor_indices(current_element, 0),
+ )
+
+ with T.serial(0, new_sparse_indices_ptr.shape[1]) as j:
+ T.buffer_store(
+ new_sparse_indices,
+ Cast(
+ sparse_indices_ptr.dtype,
+ floordiv(
+
current_element[T.tensor_indices(current_element, 0)],
+ dividers[T.tensor_indices(dividers,
j)],
+ ),
+ ),
+ (i, j),
)
- current_element = current_element_buf
T.buffer_store(
current_element,
-
flattened_indices[T.tensor_indices(flattened_indices, i)],
+ floormod(
+
current_element[T.tensor_indices(current_element, 0)],
+ dividers[T.tensor_indices(dividers, j)],
+ ),
T.tensor_indices(current_element, 0),
)
- with T.serial(0, new_sparse_indices_ptr.shape[1])
as j:
- T.buffer_store(
- new_sparse_indices,
- Cast(
- sparse_indices_ptr.dtype,
- floordiv(
-
current_element[T.tensor_indices(current_element, 0)],
-
dividers[T.tensor_indices(dividers, j)],
- ),
- ),
- (i, j),
- )
- T.buffer_store(
- current_element,
- floormod(
-
current_element[T.tensor_indices(current_element, 0)],
- dividers[T.tensor_indices(dividers,
j)],
- ),
- T.tensor_indices(current_element, 0),
- )
-
return ib.get()
new_sparse_indices_buf = decl_tensor(
diff --git a/python/tvm/topi/vision/nms.py b/python/tvm/topi/vision/nms.py
index 663eee83f2..89b24f4031 100644
--- a/python/tvm/topi/vision/nms.py
+++ b/python/tvm/topi/vision/nms.py
@@ -1058,44 +1058,43 @@ def _collect_selected_indices_ir(
batch_classes, _ = selected_indices.shape
with IRBuilder() as ib:
- with T.seq_scope():
- # Initialize output buffer to zero
- # Calculate the actual output shape based on
max_output_boxes_per_class
- if isinstance(max_output_boxes_per_class, int):
- max_output_rows = batch_classes * max_output_boxes_per_class
- else:
- # Fallback to a reasonable default if
max_output_boxes_per_class is not an integer
- max_output_rows = batch_classes * 10
- with T.serial(0, max_output_rows) as init_i:
- with T.serial(0, 3) as init_j: # 3 columns
- T.buffer_store(out, cast(0, "int64"), (init_i, init_j))
+ # Initialize output buffer to zero
+ # Calculate the actual output shape based on max_output_boxes_per_class
+ if isinstance(max_output_boxes_per_class, int):
+ max_output_rows = batch_classes * max_output_boxes_per_class
+ else:
+ # Fallback to a reasonable default if max_output_boxes_per_class
is not an integer
+ max_output_rows = batch_classes * 10
+ with T.serial(0, max_output_rows) as init_i:
+ with T.serial(0, 3) as init_j: # 3 columns
+ T.buffer_store(out, cast(0, "int64"), (init_i, init_j))
- with T.parallel(0, batch_classes) as i:
- i_64 = cast(i, "int64")
- batch_id = i_64 // num_class
- class_id = i_64 % num_class
+ with T.parallel(0, batch_classes) as i:
+ i_64 = cast(i, "int64")
+ batch_id = i_64 // num_class
+ class_id = i_64 % num_class
- if isinstance(max_output_boxes_per_class, int):
- limit = tvm.tirx.min(
- num_detections[i], tvm.tirx.IntImm("int32",
max_output_boxes_per_class)
- )
- elif isinstance(max_output_boxes_per_class, te.Tensor):
- if len(max_output_boxes_per_class.shape) == 0:
- max_boxes_val = max_output_boxes_per_class[()]
- else:
- max_boxes_val = max_output_boxes_per_class[0]
- limit = tvm.tirx.min(num_detections[i], max_boxes_val)
+ if isinstance(max_output_boxes_per_class, int):
+ limit = tvm.tirx.min(
+ num_detections[i], tvm.tirx.IntImm("int32",
max_output_boxes_per_class)
+ )
+ elif isinstance(max_output_boxes_per_class, te.Tensor):
+ if len(max_output_boxes_per_class.shape) == 0:
+ max_boxes_val = max_output_boxes_per_class[()]
else:
- limit = num_detections[i]
+ max_boxes_val = max_output_boxes_per_class[0]
+ limit = tvm.tirx.min(num_detections[i], max_boxes_val)
+ else:
+ limit = num_detections[i]
- with T.serial(0, limit) as j:
- T.buffer_store(out, batch_id, (row_offsets[i] + j, 0))
- T.buffer_store(out, class_id, (row_offsets[i] + j, 1))
- T.buffer_store(
- out,
- cast(selected_indices[i, j], "int64"),
- (row_offsets[i] + j, 2),
- )
+ with T.serial(0, limit) as j:
+ T.buffer_store(out, batch_id, (row_offsets[i] + j, 0))
+ T.buffer_store(out, class_id, (row_offsets[i] + j, 1))
+ T.buffer_store(
+ out,
+ cast(selected_indices[i, j], "int64"),
+ (row_offsets[i] + j, 2),
+ )
return ib.get()
diff --git a/src/tirx/script/ir_builder/utils.h
b/src/tirx/script/ir_builder/utils.h
index 7d20e18d3a..2578a2fb1e 100644
--- a/src/tirx/script/ir_builder/utils.h
+++ b/src/tirx/script/ir_builder/utils.h
@@ -42,9 +42,32 @@ inline void AddToParent(tvm::tirx::Stmt stmt, Span span) {
// Preserve an existing body location when flattening returns the body
itself.
if (stmt.defined() && !stmt->span.defined()) stmt->span = std::move(span);
if (builder->frames.empty()) {
- TVM_FFI_CHECK(!builder->result.has_value(), ValueError)
+ if (!builder->result.has_value()) {
+ if (stmt.as<tvm::tirx::SeqStmtNode>()) {
+ auto normalized = tvm::tirx::SeqStmt::Flatten(stmt);
+ if (!normalized->span.defined()) normalized->span = stmt->span;
+ builder->result = std::move(normalized);
+ } else {
+ builder->result = std::move(stmt);
+ }
+ return;
+ }
+ TVM_FFI_CHECK(builder->result.as<tvm::tirx::StmtNode>(), ValueError)
<< "Builder.result has already been set";
- builder->result = stmt;
+ ffi::Array<tvm::tirx::Stmt> incoming;
+ tvm::tirx::SeqStmt::Flattener{&incoming}(0, stmt);
+ if (incoming.empty()) return;
+ if (builder->result.as<tvm::tirx::SeqStmtNode>()) {
+ // Move the builder's ownership so unobserved results can grow in place.
+ // Copy-on-write preserves sequences and arrays retained by callers.
+ auto sequence =
std::move(builder->result).value().as_or_throw<tvm::tirx::SeqStmt>();
+ auto* node = sequence.CopyOnWrite();
+ for (const auto& child : incoming) node->seq.push_back(child);
+ builder->result = std::move(sequence);
+ } else {
+ builder->result = tvm::tirx::SeqStmt::Flatten(
+ builder->result.value().as_or_throw<tvm::tirx::Stmt>(), incoming);
+ }
} else if (const auto* tir_frame =
builder->frames.back().as<TIRFrameNode>()) {
ffi::GetRef<TIRFrame>(tir_frame)->stmts.push_back(stmt);
} else {
diff --git a/tests/python/relax/test_transform_legalize_ops_grad.py
b/tests/python/relax/test_transform_legalize_ops_grad.py
index 8b6db70db8..7f6f1fc3f2 100644
--- a/tests/python/relax/test_transform_legalize_ops_grad.py
+++ b/tests/python/relax/test_transform_legalize_ops_grad.py
@@ -316,7 +316,6 @@ def test_take_backward():
T.func_attr({"tirx.noalias": True})
with Ts.sblock("take_backward"):
- T.attr(0, "pragma_scope", "seq")
for i_index in range(T.int64(60)):
out_buf[i_index // T.int64(5) // T.int64(4), i_index //
T.int64(5) % T.int64(4), i_index % T.int64(5)] = T.float32(0)
for parallel, serial in T.grid(T.int64(15), T.int64(2)):
@@ -360,7 +359,6 @@ def test_take_backward_symbolic():
T.func_attr({"tirx.noalias": True})
with Ts.sblock("take_backward"):
- T.attr(0, "pragma_scope", "seq")
for i_1 in range(m_take_backward * n_take_backward):
out_buf[i_1 // n_take_backward, i_1 % n_take_backward] =
T.float32(0)
for parallel, serial in T.grid(m_take_backward,
i_take_backward):
diff --git a/tests/python/relax/test_transform_legalize_ops_manipulate.py
b/tests/python/relax/test_transform_legalize_ops_manipulate.py
index fc04589c96..2b7fe2ad14 100644
--- a/tests/python/relax/test_transform_legalize_ops_manipulate.py
+++ b/tests/python/relax/test_transform_legalize_ops_manipulate.py
@@ -1472,7 +1472,6 @@ def test_scatter_elements():
T.func_attr({"tirx.noalias": True})
with Ts.sblock("scatter_elements_generic"):
- T.attr(0, "pragma_scope", "seq")
for i in T.parallel(T.int64(16)):
out_buf[i // T.int64(4), i % T.int64(4)] = rxplaceholder[
i // T.int64(4), i % T.int64(4)
@@ -1575,7 +1574,6 @@ def test_scatter_elements_symbolic():
T.func_attr({"tirx.noalias": True})
with Ts.sblock("scatter_elements_generic"):
- T.attr(0, "pragma_scope", "seq")
for i in T.parallel(a_scatter_elements * b_scatter_elements):
out_buf[i // b_scatter_elements, i % b_scatter_elements] =
rxplaceholder[i // b_scatter_elements, i % b_scatter_elements]
for fused in T.parallel(m_scatter_elements):
@@ -1899,7 +1897,6 @@ def test_scatter_nd():
with Ts.sblock("scatter_nd_generic"):
Ts.reads()
Ts.writes()
- T.attr(0, "pragma_scope", "seq")
for i in range(T.int64(8)):
out_buf[i] = data[i]
for j in range(T.int64(4)):
