chore: import upstream snapshot with attribution

python/tvm/topi/__init__.py

@@ -0,0 +1,64 @@
+# 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
+"""TVM Operator Inventory.
+
+TOPI is the operator collection library for TVM, to provide sugars
+for constructing compute declaration as well as optimized schedules.
+
+Some of the schedule function may have been specially optimized for a
+specific workload.
+"""
+
+from tvm.libinfo import __version__
+
+# Ensure C++ schedules get registered first, so python schedules can
+# override them.
+from . import cpp
+
+from .math import *
+from .tensor import *
+from .index_put import *
+from .reduction import *
+from .transform import *
+from .broadcast import *
+from . import _te_tensor_overload
+from .sort import *
+from .scatter import *
+from .scatter_elements import *
+from .slice_scatter import *
+from .sparse_reshape import *
+from .scan import *
+from .einsum import *
+from .unique import *
+from .searchsorted import *
+from .signal import *
+from . import nn
+from . import utils
+from . import image
+from . import vision
+from . import gpu
+
+# error reporting
+from .utils import InvalidShapeError
+
+# not import testing by default
+# because testing can have extra deps that are not necessary
+# we can import them from test cases explicitly
+# from . import testing

python/tvm/topi/_te_tensor_overload.py

@@ -0,0 +1,64 @@
+# 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.
+"""Register TOPI implementations for TE tensor overload hooks."""
+
+from tvm import te
+from tvm.te import _te_tensor_overload as _overload
+from tvm.tirx import expr as _expr
+
+from . import broadcast as _broadcast
+from . import math as _math
+
+
+def _is_integer(value):
+    if isinstance(value, te.Tensor | te.TensorSlice):
+        return value.dtype.matches_code(_expr.DataTypeCode.INT)
+    return _expr._dtype_is_int(value)
+
+
+def _binary(op, reflected=False, check_integer=False):
+    def implementation(lhs, rhs):
+        if not isinstance(lhs, te.Tensor) and not isinstance(rhs, te.Tensor):
+            return NotImplemented
+        if reflected:
+            lhs, rhs = rhs, lhs
+        if check_integer and _is_integer(lhs) and _is_integer(rhs):
+            raise _expr.div_ambiguity_error()
+        return op(lhs, rhs)
+
+    return implementation
+
+
+_overload.__add__ = _binary(_broadcast.add)
+_overload.__radd__ = _binary(_broadcast.add, reflected=True)
+_overload.__sub__ = _binary(_broadcast.subtract)
+_overload.__rsub__ = _binary(_broadcast.subtract, reflected=True)
+_overload.__mul__ = _binary(_broadcast.multiply)
+_overload.__rmul__ = _binary(_broadcast.multiply, reflected=True)
+_overload.__div__ = _binary(_broadcast.divide, check_integer=True)
+_overload.__rdiv__ = _binary(_broadcast.divide, reflected=True, check_integer=True)
+_overload.__truediv__ = _overload.__div__
+_overload.__rtruediv__ = _overload.__rdiv__
+
+
+def _astype(value, dtype, span=None):
+    if not isinstance(value, te.Tensor):
+        return NotImplemented
+    return _math.cast(value, dtype, span)
+
+
+_overload.astype = _astype

python/tvm/topi/broadcast.py

@@ -0,0 +1,566 @@
+# 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.
+"""Broadcast operators"""
+
+from . import cpp as _cpp
+
+
+def broadcast_to(data, shape):
+    """Broadcast the src to the target shape
+
+    We follows the numpy broadcasting rule.
+    See also https://docs.scipy.org/doc/numpy/user/basics.broadcasting.html
+
+    Parameters
+    ----------
+    data : tvm.te.Tensor
+        The input data
+
+    shape : list or tuple
+        The target shape to be broadcasted.
+
+    Returns
+    -------
+    ret : tvm.te.Tensor
+    """
+    return _cpp.broadcast_to(data, shape)
+
+
+def add(lhs, rhs):
+    """Addition with auto-broadcasting
+
+    Parameters
+    ----------
+    lhs : tvm.te.Tensor or Expr
+        The left operand
+    rhs : tvm.te.Tensor or Expr
+        The right operand
+
+    Returns
+    -------
+    ret : tvm.te.Tensor or Expr
+        Returns Expr if both operands are Expr.
+        Otherwise returns Tensor.
+    """
+    return _cpp.add(lhs, rhs)
+
+
+def subtract(lhs, rhs):
+    """Subtraction with auto-broadcasting
+
+    Parameters
+    ----------
+    lhs : tvm.te.Tensor or Expr
+        The left operand
+    rhs : tvm.te.Tensor or Expr
+        The right operand
+
+    Returns
+    -------
+    ret : tvm.te.Tensor or Expr
+        Returns Expr if both operands are Expr.
+        Otherwise returns Tensor.
+    """
+    return _cpp.subtract(lhs, rhs)
+
+
+def multiply(lhs, rhs):
+    """Multiplication with auto-broadcasting
+
+    Parameters
+    ----------
+    lhs : tvm.te.Tensor or Expr
+        The left operand
+    rhs : tvm.te.Tensor or Expr
+        The right operand
+
+    Returns
+    -------
+    ret : tvm.te.Tensor or Expr
+        Returns Expr if both operands are Expr.
+        Otherwise returns Tensor.
+    """
+    return _cpp.multiply(lhs, rhs)
+
+
+def divide(lhs, rhs):
+    """Division with auto-broadcasting
+
+    Parameters
+    ----------
+    lhs : tvm.te.Tensor or Expr
+        The left operand
+    rhs : tvm.te.Tensor or Expr
+        The right operand
+
+    Returns
+    -------
+    ret : tvm.te.Tensor or Expr
+        Returns Expr if both operands are Expr.
+        Otherwise returns Tensor.
+    """
+    return _cpp.divide(lhs, rhs)
+
+
+def floor_divide(lhs, rhs):
+    """Floor division with auto-broadcasting
+
+    Parameters
+    ----------
+    lhs : tvm.te.Tensor or Expr
+        The left operand
+    rhs : tvm.te.Tensor or Expr
+        The right operand
+
+    Returns
+    -------
+    ret : tvm.te.Tensor or Expr
+        Returns Expr if both operands are Expr.
+        Otherwise returns Tensor.
+    """
+    return _cpp.floor_divide(lhs, rhs)
+
+
+def log_add_exp(lhs, rhs):
+    """Log-sum-exp operation with auto-broadcasting.
+
+    Parameters
+    ----------
+    x1 : tvm.te.Tensor or Expr
+        The first input tensor or expression.
+    x2 : tvm.te.Tensor or Expr
+        The second input tensor or expression.
+
+    Returns
+    -------
+    ret : tvm.te.Tensor or Expr
+        Returns an Expr if both operands are Expr.
+        Otherwise, returns a Tensor.
+    """
+    return _cpp.log_add_exp(lhs, rhs)
+
+
+def mod(lhs, rhs):
+    """Modulus with auto-broadcasting
+
+    Parameters
+    ----------
+    lhs : tvm.te.Tensor or Expr
+        The left operand
+    rhs : tvm.te.Tensor or Expr
+        The right operand
+
+    Returns
+    -------
+    ret : tvm.te.Tensor or Expr
+        Returns Expr if both operands are Expr.
+        Otherwise returns Tensor.
+    """
+    return _cpp.mod(lhs, rhs)
+
+
+def floor_mod(lhs, rhs):
+    """Floor modulus with auto-broadcasting
+
+    Parameters
+    ----------
+    lhs : tvm.te.Tensor or Expr
+        The left operand
+    rhs : tvm.te.Tensor or Expr
+        The right operand
+
+    Returns
+    -------
+    ret : tvm.te.Tensor or Expr
+        Returns Expr if both operands are Expr.
+        Otherwise returns Tensor.
+    """
+    return _cpp.floor_mod(lhs, rhs)
+
+
+def maximum(lhs, rhs):
+    """Take element-wise maximum of two tensors with auto-broadcasting
+
+    Parameters
+    ----------
+    lhs : tvm.te.Tensor or Expr
+        The left operand
+    rhs : tvm.te.Tensor or Expr
+        The right operand
+
+    Returns
+    -------
+    ret : tvm.te.Tensor or Expr
+        Returns Expr if both operands are Expr.
+        Otherwise returns Tensor.
+    """
+    return _cpp.maximum(lhs, rhs)
+
+
+def minimum(lhs, rhs):
+    """Take element-wise maximum of two tensors with auto-broadcasting
+
+    Parameters
+    ----------
+    lhs : tvm.te.Tensor or Expr
+        The left operand
+    rhs : tvm.te.Tensor or Expr
+        The right operand
+
+    Returns
+    -------
+    ret : tvm.te.Tensor or Expr
+        Returns Expr if both operands are Expr.
+        Otherwise returns Tensor.
+    """
+    return _cpp.minimum(lhs, rhs)
+
+
+def power(lhs, rhs):
+    """Power with auto-broadcasting
+
+    Parameters
+    ----------
+    lhs : tvm.te.Tensor or Expr
+        The left operand
+    rhs : tvm.te.Tensor or Expr
+        The right operand
+
+    Returns
+    -------
+    ret : tvm.te.Tensor or Expr
+        Returns Expr if both operands are Expr.
+        Otherwise returns Tensor.
+    """
+    return _cpp.power(lhs, rhs)
+
+
+def atan2(lhs, rhs):
+    """Atan2 with auto-broadcasting.
+
+    Parameters
+    ----------
+    lhs : tvm.te.Tensor or Expr
+        The left operand (y-coordinates).
+    rhs : tvm.te.Tensor or Expr
+        The right operand (x-coordinates).
+
+    Returns
+    -------
+    ret : tvm.te.Tensor or Expr
+        Returns Expr if both operands are Expr.
+        Otherwise returns Tensor.
+    """
+    return _cpp.atan2(lhs, rhs)
+
+
+def left_shift(lhs, rhs):
+    """Left shift with auto-broadcasting
+
+    Parameters
+    ----------
+    lhs : tvm.te.Tensor or Expr
+        The left operand
+    rhs : tvm.te.Tensor or Expr
+        The right operand
+
+    Returns
+    -------
+    ret : tvm.te.Tensor or Expr
+        Returns Expr if both operands are Expr.
+        Otherwise returns Tensor.
+    """
+    return _cpp.left_shift(lhs, rhs)
+
+
+def right_shift(lhs, rhs):
+    """Right shift with auto-broadcasting
+
+    Parameters
+    ----------
+    lhs : tvm.te.Tensor or Expr
+        The left operand
+    rhs : tvm.te.Tensor or Expr
+        The right operand
+
+    Returns
+    -------
+    ret : tvm.te.Tensor or Expr
+        Returns Expr if both operands are Expr.
+        Otherwise returns Tensor.
+    """
+    return _cpp.right_shift(lhs, rhs)
+
+
+def greater(lhs, rhs):
+    """Compute (lhs>rhs) with auto-broadcasting
+
+    Parameters
+    ----------
+    lhs : tvm.te.Tensor or Expr
+        The left operand
+    rhs : tvm.te.Tensor or Expr
+        The right operand
+
+    Returns
+    -------
+    ret : tvm.te.Tensor or Expr
+        Returns Expr if both operands are Expr.
+        Otherwise returns Tensor.
+    """
+    return _cpp.greater(lhs, rhs)
+
+
+def less(lhs, rhs):
+    """Compute (lhs<rhs) with auto-broadcasting
+
+    Parameters
+    ----------
+    lhs : tvm.te.Tensor or Expr
+        The left operand
+    rhs : tvm.te.Tensor or Expr
+        The right operand
+
+    Returns
+    -------
+    ret : tvm.te.Tensor or Expr
+        Returns Expr if both operands are Expr.
+        Otherwise returns Tensor.
+    """
+    return _cpp.less(lhs, rhs)
+
+
+def equal(lhs, rhs):
+    """Compute (lhs==rhs) with auto-broadcasting
+
+    Parameters
+    ----------
+    lhs : tvm.te.Tensor or Expr
+        The left operand
+    rhs : tvm.te.Tensor or Expr
+        The right operand
+
+    Returns
+    -------
+    ret : tvm.te.Tensor or Expr
+        Returns Expr if both operands are Expr.
+        Otherwise returns Tensor.
+    """
+    return _cpp.equal(lhs, rhs)
+
+
+def not_equal(lhs, rhs):
+    """Compute (lhs!=rhs) with auto-broadcasting
+
+    Parameters
+    ----------
+    lhs : tvm.te.Tensor or Expr
+        The left operand
+    rhs : tvm.te.Tensor or Expr
+        The right operand
+
+    Returns
+    -------
+    ret : tvm.te.Tensor or Expr
+        Returns Expr if both operands are Expr.
+        Otherwise returns Tensor.
+    """
+    return _cpp.not_equal(lhs, rhs)
+
+
+def greater_equal(lhs, rhs):
+    """Compute (lhs>=rhs) with auto-broadcasting
+
+    Parameters
+    ----------
+    lhs : tvm.te.Tensor or Expr
+        The left operand
+    rhs : tvm.te.Tensor or Expr
+        The right operand
+
+    Returns
+    -------
+    ret : tvm.te.Tensor or Expr
+        Returns Expr if both operands are Expr.
+        Otherwise returns Tensor.
+    """
+    return _cpp.greater_equal(lhs, rhs)
+
+
+def less_equal(lhs, rhs):
+    """Compute (lhs<=rhs) with auto-broadcasting
+
+    Parameters
+    ----------
+    lhs : tvm.te.Tensor or Expr
+        The left operand
+    rhs : tvm.te.Tensor or Expr
+        The right operand
+
+    Returns
+    -------
+    ret : tvm.te.Tensor or Expr
+        Returns Expr if both operands are Expr.
+        Otherwise returns Tensor.
+    """
+    return _cpp.less_equal(lhs, rhs)
+
+
+def logical_and(lhs, rhs):
+    """Compute element-wise logical and of data.
+
+    Parameters
+    ----------
+    lhs : tvm.te.Tensor or Expr
+          The left operand
+    rhs : tvm.te.Tensor or Expr
+          The right operand
+
+    Returns
+    -------
+    ret : tvm.te.Tensor or Expr
+          Returns Expr if both operands are Expr.
+          Otherwise returns Tensor.
+    """
+    return _cpp.logical_and(lhs, rhs)
+
+
+def logical_or(lhs, rhs):
+    """Compute element-wise logical or of data.
+
+    Parameters
+    ----------
+    lhs : tvm.te.Tensor or Expr
+          The left operand
+    rhs : tvm.te.Tensor or Expr
+          The right operand
+
+    Returns
+    -------
+    ret : tvm.te.Tensor or Expr
+          Returns Expr if both operands are Expr.
+          Otherwise returns Tensor.
+    """
+    return _cpp.logical_or(lhs, rhs)
+
+
+def logical_xor(lhs, rhs):
+    """Compute element-wise logical xor of data.
+
+    Parameters
+    ----------
+    lhs : tvm.te.Tensor or Expr
+          The left operand
+    rhs : tvm.te.Tensor or Expr
+          The right operand
+
+    Returns
+    -------
+    ret : tvm.te.Tensor or Expr
+          Returns Expr if both operands are Expr.
+          Otherwise returns Tensor.
+    """
+    return _cpp.logical_xor(lhs, rhs)
+
+
+def bitwise_and(lhs, rhs):
+    """Compute element-wise bitwise and of data.
+
+    Parameters
+    ----------
+    lhs : tvm.te.Tensor or Expr
+          The left operand
+    rhs : tvm.te.Tensor or Expr
+          The right operand
+
+    Returns
+    -------
+    ret : tvm.te.Tensor or Expr
+          Returns Expr if both operands are Expr.
+          Otherwise returns Tensor.
+    """
+    return _cpp.bitwise_and(lhs, rhs)
+
+
+def bitwise_or(lhs, rhs):
+    """Compute element-wise bitwise or of data.
+
+    Parameters
+    ----------
+    lhs : tvm.te.Tensor or Expr
+          The left operand
+    rhs : tvm.te.Tensor or Expr
+          The right operand
+
+    Returns
+    -------
+    ret : tvm.te.Tensor or Expr
+          Returns Expr if both operands are Expr.
+          Otherwise returns Tensor.
+    """
+    return _cpp.bitwise_or(lhs, rhs)
+
+
+def bitwise_xor(lhs, rhs):
+    """Compute element-wise bitwise xor of data.
+
+    Parameters
+    ----------
+    lhs : tvm.te.Tensor or Expr
+          The left operand
+    rhs : tvm.te.Tensor or Expr
+          The right operand
+
+    Returns
+    -------
+    ret : tvm.te.Tensor or Expr
+          Returns Expr if both operands are Expr.
+          Otherwise returns Tensor.
+    """
+    return _cpp.bitwise_xor(lhs, rhs)
+
+
+def logical_not(data):
+    """Compute element-wise logical not of data.
+
+    Parameters
+    ----------
+    data : tvm.te.Tensor or Expr
+
+    Returns
+    -------
+    ret : tvm.te.Tensor or Expr
+          Returns Expr if the operand are Expr.
+          Otherwise returns Tensor.
+    """
+    return _cpp.logical_not(data)
+
+
+def bitwise_not(data):
+    """Compute element-wise bitwise not of data.
+
+    Parameters
+    ----------
+    data : tvm.te.Tensor or Expr
+
+    Returns
+    -------
+    ret : tvm.te.Tensor or Expr
+          Returns Expr if the operand are Expr.
+          Otherwise returns Tensor.
+    """
+    return _cpp.bitwise_not(data)

python/tvm/topi/cpp/__init__.py

@@ -0,0 +1,28 @@
+# 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.
+
+"""FFI for C++ TOPI ops and schedules"""
+
+from .impl import *  # pylint: disable=wildcard-import
+from . import cuda
+from . import nn
+from . import vision
+from . import x86
+from . import generic
+from . import rocm
+from . import utils

python/tvm/topi/cpp/cuda.py

@@ -0,0 +1,21 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""FFI for CUDA TOPI ops and schedules"""
+
+import tvm_ffi
+
+tvm_ffi.init_ffi_api("topi.cuda", "tvm.topi.cpp.cuda")

python/tvm/topi/cpp/generic.py

@@ -0,0 +1,21 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""FFI for generic TOPI ops and schedules"""
+
+import tvm_ffi
+
+tvm_ffi.init_ffi_api("topi.generic", "tvm.topi.cpp.generic")

python/tvm/topi/cpp/impl.py

@@ -0,0 +1,21 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Load Lib for C++ TOPI ops and schedules"""
+
+import tvm_ffi
+
+tvm_ffi.init_ffi_api("topi", "tvm.topi.cpp")

python/tvm/topi/cpp/nn.py

@@ -0,0 +1,21 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""FFI for NN TOPI ops and schedules"""
+
+import tvm_ffi
+
+tvm_ffi.init_ffi_api("topi.nn", "tvm.topi.cpp.nn")

python/tvm/topi/cpp/rocm.py

@@ -0,0 +1,21 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""FFI for Rocm TOPI ops and schedules"""
+
+import tvm_ffi
+
+tvm_ffi.init_ffi_api("topi.rocm", "tvm.topi.cpp.rocm")

python/tvm/topi/cpp/utils.py

@@ -0,0 +1,21 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""FFI for TOPI utility functions"""
+
+import tvm_ffi
+
+tvm_ffi.init_ffi_api("topi.utils", "tvm.topi.cpp.utils")

python/tvm/topi/cpp/vision/__init__.py

@@ -0,0 +1,26 @@
+# 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.
+
+"""FFI for vision TOPI ops and schedules"""
+
+import tvm_ffi
+
+from . import yolo
+from ...vision import nms
+
+tvm_ffi.init_ffi_api("topi.vision", "tvm.topi.cpp.vision")

python/tvm/topi/cpp/vision/yolo.py

@@ -0,0 +1,21 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""FFI for Yolo TOPI ops and schedules"""
+
+import tvm_ffi
+
+tvm_ffi.init_ffi_api("topi.vision.yolo", "tvm.topi.cpp.vision.yolo")

python/tvm/topi/cpp/x86.py

@@ -0,0 +1,21 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""FFI for x86 TOPI ops and schedules"""
+
+import tvm_ffi
+
+tvm_ffi.init_ffi_api("topi.x86", "tvm.topi.cpp.x86")

python/tvm/topi/einsum.py

@@ -0,0 +1,45 @@
+# 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,consider-using-enumerate,redefined-outer-name
+"""Einsum operator"""
+
+from . import cpp
+
+
+def einsum(subscripts, *operand):
+    """Evaluates the Einstein summation convention on the operands.
+
+    Parameters
+    ----------
+    subscripts : string
+        Specifies the subscripts for summation as comma separated list of subscript labels.
+        An implicit (classical Einstein summation) calculation is performed unless the
+        explicit indicator '->' is included as well as subscript labels of the precise
+        output form.
+
+    a_tuple : tuple of tvm.te.Tensor
+        These are the Tensors for the operation.
+        The only difference of einsum between in tvm and numpy is it needs an extra brackets
+        for the tensors. For example, topi.einsum("ij, jk -> ik", (A, B)).
+
+    Returns
+    -------
+    out : tvm.te.Tensor
+        The calculation based on the Einstein summation convention.
+    """
+
+    return cpp.einsum(subscripts, operand)

python/tvm/topi/gpu/__init__.py

@@ -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

@@ -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

@@ -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

@@ -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/image/__init__.py

@@ -0,0 +1,24 @@
+# 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=wildcard-import
+"""IMAGE network operators"""
+
+from .resize import *
+from .dilation2d import *
+from .grid_sample import *

python/tvm/topi/image/dilation2d.py

@@ -0,0 +1,178 @@
+# 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, unused-variable, too-many-locals
+# pylint: disable=unused-argument, redefined-builtin
+"""Dilation2D operators"""
+
+from tvm import te
+from tvm.topi.utils import simplify
+
+from ..nn.pad import pad
+from ..nn.utils import get_pad_tuple
+
+
+def dilation2d_nchw(input, filter, stride, padding, dilations, out_dtype=None):
+    """Morphological dilation operator in NCHW layout.
+
+    Parameters
+    ----------
+    input : tvm.te.Tensor
+        4-D with shape [batch, in_channel, in_height, in_width]
+
+    filter : tvm.te.Tensor
+        3-D with shape [ in_channel, filter_height, filter_width]
+
+    stride : int or a list/tuple of two ints
+        Stride size, or [stride_height, stride_width]
+
+    padding : int or str
+        Padding size
+
+    dilations: int or a list/tuple of two ints
+        dilation size, or [dilation_height, dilation_width]
+
+    out_dtype : Optional[str]
+        Specifies the output data type.
+
+    Returns
+    -------
+    Output : tvm.te.Tensor
+        4-D with shape [batch, in_channel, out_height, out_width]
+    """
+    if out_dtype is None:
+        out_dtype = input.dtype
+    assert isinstance(stride, int) or len(stride) == 2
+    assert isinstance(dilations, int) or len(dilations) == 2
+    if isinstance(stride, int):
+        stride_h = stride_w = stride
+    else:
+        stride_h, stride_w = stride
+
+    if isinstance(dilations, int):
+        dilation_h = dilation_w = dilations
+    else:
+        dilation_h, dilation_w = dilations
+
+    batch, in_channel, in_height, in_width = input.shape
+    channel, kernel_h, kernel_w = filter.shape
+    assert in_channel.value == channel.value, (
+        "For Dilation2D input and filter channels should be same."
+    )
+
+    # compute the output shape
+    dilated_kernel_h = (kernel_h - 1) * dilation_h + 1
+    dilated_kernel_w = (kernel_w - 1) * dilation_w + 1
+    pad_top, pad_left, pad_down, pad_right = get_pad_tuple(
+        padding, (dilated_kernel_h, dilated_kernel_w)
+    )
+
+    out_height = simplify((in_height - dilated_kernel_h + pad_top + pad_down) // stride_h + 1)
+    out_width = simplify((in_width - dilated_kernel_w + pad_left + pad_right) // stride_w + 1)
+    # compute graph
+    pad_before = [0, 0, pad_top, pad_left]
+    pad_after = [0, 0, pad_down, pad_right]
+    temp = pad(input, pad_before, pad_after, name="pad_temp")
+    ry = te.reduce_axis((0, kernel_h), name="ry")
+    rx = te.reduce_axis((0, kernel_w), name="rx")
+
+    return te.compute(
+        (batch, in_channel, out_height, out_width),
+        lambda nn, ff, yy, xx: te.max(
+            temp[nn, ff, yy * stride_h + ry * dilation_h, xx * stride_w + rx * dilation_w].astype(
+                out_dtype
+            )
+            + filter[ff, ry, rx].astype(out_dtype),
+            axis=[ry, rx],
+        ),
+        tag="dilation2d_nchw",
+    )
+
+
+def dilation2d_nhwc(input, filter, stride, padding, dilations, out_dtype=None):
+    """Morphological 2d dilation NHWC layout.
+
+    Parameters
+    ----------
+    input : tvm.te.Tensor
+        4-D with shape [batch, in_height, in_width, in_channel]
+
+    filter : tvm.te.Tensor
+        3-D with shape [filter_height, filter_width, in_channel]
+
+    stride : int or a list/tuple of two ints
+        Stride size, or [stride_height, stride_width]
+
+    padding : int
+        Padding size
+
+    dilations: int or a list/tuple of two ints
+        dilation size, or [dilation_height, dilation_width]
+
+    out_dtype : Optional[str]
+        Specifies the output data type.
+
+    Returns
+    -------
+    Output : tvm.te.Tensor
+        4-D with shape [batch, out_height, out_width, in_channel]
+    """
+    if out_dtype is None:
+        out_dtype = input.dtype
+    assert isinstance(stride, int) or len(stride) == 2
+    assert isinstance(dilations, int) or len(dilations) == 2
+    if isinstance(stride, int):
+        stride_h = stride_w = stride
+    else:
+        stride_h, stride_w = stride
+
+    if isinstance(dilations, int):
+        dilation_h = dilation_w = dilations
+    else:
+        dilation_h, dilation_w = dilations
+
+    batch, in_height, in_width, in_channel = input.shape
+    kernel_h, kernel_w, channel = filter.shape
+    assert in_channel.value == channel.value, (
+        "For Dilation2D input and filter channels should be same."
+    )
+
+    # compute the output shape
+    dilated_kernel_h = (kernel_h - 1) * dilation_h + 1
+    dilated_kernel_w = (kernel_w - 1) * dilation_w + 1
+    pad_top, pad_left, pad_down, pad_right = get_pad_tuple(
+        padding, (dilated_kernel_h, dilated_kernel_w)
+    )
+
+    out_height = simplify((in_height - dilated_kernel_h + pad_top + pad_down) // stride_h + 1)
+    out_width = simplify((in_width - dilated_kernel_w + pad_left + pad_right) // stride_w + 1)
+    pad_before = [0, pad_top, pad_left, 0]
+    pad_after = [0, pad_down, pad_right, 0]
+    padded_input = pad(input, pad_before, pad_after, name="padded_input")
+    ry = te.reduce_axis((0, kernel_h), name="ry")
+    rx = te.reduce_axis((0, kernel_w), name="rx")
+
+    return te.compute(
+        (batch, out_height, out_width, in_channel),
+        lambda nn, yy, xx, ff: te.max(
+            padded_input[
+                nn, yy * stride_h + ry * dilation_h, xx * stride_w + rx * dilation_w, ff
+            ].astype(out_dtype)
+            + filter[ry, rx, ff].astype(out_dtype),
+            axis=[ry, rx],
+        ),
+        tag="dilation2d_nhcw",
+    )

python/tvm/topi/image/grid_sample.py

@@ -0,0 +1,544 @@
+# 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
+"""affine_grid and grid_sample operator"""
+
+from tvm import te, tirx
+
+
+def affine_grid(data, target_shape, align_corners=True):
+    """affine_grid operator that generates a 2D or 3D sampling grid.
+
+    This operation is described in https://arxiv.org/pdf/1506.02025.pdf. It generates a uniform
+    sampling grid within the target shape and normalizes it to [-1, 1]. The provided affine
+    transformation is then applied on the sampling grid.
+
+    Parameters
+    ----------
+    data : tvm.Tensor
+        3-D with shape [batch, 2, 3] for 2D or [batch, 3, 4] for 3D. The affine matrix.
+
+    target_shape: list/tuple of int
+        Specifies the output spatial shape (H, W) for 2D or (D, H, W) for 3D.
+
+    align_corners : bool
+        If True, normalized coordinates map to corner pixels; if False, to pixel centers
+        (the PyTorch / ONNX default).
+
+    Returns
+    -------
+    Output : tvm.Tensor
+        [batch, 2, H, W] for 2D or [batch, 3, D, H, W] for 3D.
+    """
+    assert len(target_shape) in (2, 3)
+    if align_corners:
+        assert all(s > 1 for s in target_shape), (
+            "target spatial dims should be greater than 1 when align_corners is True"
+        )
+
+    dtype = data.dtype
+    if align_corners:
+        starts = [tirx.const(-1.0, dtype=dtype) for _ in target_shape]
+        steps = [tirx.const((2.0 - 1e-7) / (s - 1), dtype=dtype) for s in target_shape]
+    else:
+        # Pixel centers: coordinate i maps to (2 * i + 1) / size - 1.
+        starts = [tirx.const(-1.0 + 1.0 / s, dtype=dtype) for s in target_shape]
+        steps = [tirx.const(2.0 / s, dtype=dtype) for s in target_shape]
+
+    ndim = len(target_shape)
+
+    def _compute(n, dim, *coords):
+        # coords are ordered slowest-to-fastest (e.g. (k, i, j)); the affine matrix
+        # columns are fastest-to-slowest (x, y, z), so index it in reverse.
+        val = data[n, dim, ndim]  # translation column
+        for r in range(ndim):
+            coord = starts[r] + coords[r] * steps[r]
+            val += data[n, dim, ndim - 1 - r] * coord
+        return val
+
+    oshape = (data.shape[0], ndim, *target_shape)
+    return te.compute(oshape, _compute, tag="affine_grid")
+
+
+def _grid_sample_2d(
+    data, grid, method="bilinear", layout="NCHW", padding_mode="zeros", align_corners=True
+):
+    """Applies bilinear/nearest/bicubic sampling to input feature map.
+
+    Given :math:`data` and :math:`grid` assuming NCHW layout, then the output is computed by
+
+    .. math::
+
+        x_{src} = grid[batch, 0, y_{dst}, x_{dst}] \\
+        y_{src} = grid[batch, 1, y_{dst}, x_{dst}] \\
+        output[batch, channel, y_{dst}, x_{dst}] = G(data[batch, channel, y_{src}, x_{src})
+
+    :math:`x_{dst}`, :math:`y_{dst}` enumerate all spatial locations in :math:`output`, and
+    :math:`G()` denotes the interpolation method.
+
+    The out-boundary points will be padded with zeros if padding_mode is "zeros", or
+    border pixel value if padding_mode is "border", or
+    inner pixel value if padding_mode is "reflection".
+
+    The left-top corner (-1, -1) and right-bottom corner (1, 1) in grid will be map to
+    (0, 0) and (h - 1, w - 1) of data if align_corners is "True", or
+    (-0.5, -0.5) and (h - 0.5, w - 0.5) of data if align_corners is "False".
+
+    The shape of the output will be (data.shape[0], data.shape[1], grid.shape[2], grid.shape[3]).
+
+    The operator assumes that :math:`grid` has been normalized to [-1, 1].
+
+    grid_sample often cooperates with affine_grid which generates sampling grids for grid_sample.
+
+    Parameters
+    ----------
+    data : tvm.Tensor
+        4-D with shape [batch, in_channel, in_height, in_width]
+
+    grid : tvm.Tensor
+        4-D with shape [batch, 2, out_height, out_width]
+
+    method : str
+        The interpolation method "nearest", "bilinear", "bicubic" are supported.
+
+    layout : str
+        The layout of input data and the output.
+
+    padding_mode : str
+        The padding mode for outside grid values, "zeros", "border", "reflection" are supported.
+
+    align_corners: bool
+        Geometrically, we consider the pixels of the input as squares rather than points.
+        If set to "True", the extrema ("-1" and "1") are considered as referring
+        to the center points of the input corner pixels. If set to "False", they
+        are instead considered as referring to the corner points of the input corner
+        pixels, making the sampling more resolution agnostic.
+
+    Returns
+    -------
+    Output : tvm.Tensor
+        4-D with shape [batch, in_channel, out_height, out_width]
+    """
+
+    assert method in ("bilinear", "nearest", "bicubic"), f"{method} is not supported"
+    assert padding_mode in ("zeros", "border", "reflection"), f"{padding_mode} is not supported"
+    assert layout == "NCHW", f"{layout} is not supported"
+
+    batch, in_channel, in_height, in_width = data.shape
+    out_height, out_width = grid.shape[2:]
+
+    def _get_pixel_value(n, c, h, w):
+        return te.if_then_else(
+            te.all(h >= 0, w >= 0, h < in_height, w < in_width),
+            data[n, c, h, w],
+            tirx.const(0.0, dtype=data.dtype),
+        )
+
+    def _unnormalize(h, w):
+        if align_corners:
+            y = (h + 1) * (in_height - 1) / 2
+            x = (w + 1) * (in_width - 1) / 2
+        else:
+            y = -0.5 + (h + 1) * in_height / 2
+            x = -0.5 + (w + 1) * in_width / 2
+        return (y, x)
+
+    def _clip_coordinates(x, size):
+        return te.min(te.max(x, 0), size - 1)
+
+    def _compute_source_index(n, h, w):
+        y = grid[n, 1, h, w]
+        x = grid[n, 0, h, w]
+        y, x = _unnormalize(y, x)
+
+        if padding_mode == "reflection":
+            y = _reflect_coordinates(y, in_height)
+            x = _reflect_coordinates(x, in_width)
+            y = _clip_coordinates(y, in_height)
+            x = _clip_coordinates(x, in_width)
+        elif padding_mode == "border":
+            y = _clip_coordinates(y, in_height)
+            x = _clip_coordinates(x, in_width)
+
+        return (y, x)
+
+    def _reflect_coordinates(x, size):
+        def __refelection(x, size, corner_start):
+            def __reflect(index, size, corner_start):
+                index_align_corner = te.abs(corner_start - index)
+                size_times = te.truncdiv(index_align_corner.astype("int32"), size).astype("int32")
+                t = tirx.Mod(size_times, 2)
+                extra = index_align_corner - size_times * size
+                return tirx.if_then_else(
+                    tirx.EQ(t, 0), extra + corner_start, size - extra + corner_start
+                )
+
+            return tirx.if_then_else(
+                tirx.all(x >= corner_start, x <= size + corner_start),
+                x,
+                __reflect(x, size, corner_start),
+            )
+
+        if align_corners:
+            new_x = __refelection(x, size - 1, 0)
+        else:
+            new_x = __refelection(x, size, -0.5)
+        return new_x
+
+    def _bilinear_sample(n, c, h, w):
+        y, x = _compute_source_index(n, h, w)
+        y0 = te.floor(y).astype("int32")
+        x0 = te.floor(x).astype("int32")
+        y1 = y0 + tirx.const(1, "int32")
+        x1 = x0 + tirx.const(1, "int32")
+
+        return (
+            _get_pixel_value(n, c, y0, x0) * (1.0 - (y - y0)) * (1.0 - (x - x0))
+            + _get_pixel_value(n, c, y0, x1) * (1.0 - (y - y0)) * (x - x0)
+            + _get_pixel_value(n, c, y1, x0) * (y - y0) * (1.0 - (x - x0))
+            + _get_pixel_value(n, c, y1, x1) * (y - y0) * (x - x0)
+        )
+
+    def _nearest_sample(n, c, h, w):
+        y, x = _compute_source_index(n, h, w)
+        y_new = te.nearbyint(y).astype("int32")
+        x_new = te.nearbyint(x).astype("int32")
+
+        return _get_pixel_value(n, c, y_new, x_new)
+
+    def _bicubic_sample(n, c, h, w):
+        A = -0.75  # -0.75 is used in pytorch, it maybe different in other frameworks
+
+        def cubic_weight_1(fraction):
+            return ((A + 2) * fraction - (A + 3)) * fraction * fraction + 1
+
+        def cubic_weight_2(fraction):
+            return ((A * fraction - 5 * A) * fraction + 8 * A) * fraction - 4 * A
+
+        def cubic_interp_1d(pixel_0, pixel_1, pixel_2, pixel_3, fraction):
+            weights = [0] * 4
+            weights[0] = cubic_weight_2(fraction + 1)
+            weights[1] = cubic_weight_1(fraction)
+            weights[2] = cubic_weight_1(1 - fraction)
+            weights[3] = cubic_weight_2(2 - fraction)
+            return (
+                pixel_0 * weights[0]
+                + pixel_1 * weights[1]
+                + pixel_2 * weights[2]
+                + pixel_3 * weights[3]
+            )
+
+        y = grid[n, 1, h, w]
+        x = grid[n, 0, h, w]
+        y, x = _unnormalize(y, x)
+        y_floor = te.floor(y).astype("int32")
+        x_floor = te.floor(x).astype("int32")
+        y_fraction = y - y_floor
+        x_fraction = x - x_floor
+
+        coefficients = [0] * 4
+
+        for i in range(4):
+            y_ = y_floor - 1 + i
+            x_0 = x_floor - 1
+            x_1 = x_floor + 0
+            x_2 = x_floor + 1
+            x_3 = x_floor + 2
+
+            if padding_mode == "border":
+                y_ = _clip_coordinates(y_, in_height).astype("int32")
+                x_0 = _clip_coordinates(x_0, in_width).astype("int32")
+                x_1 = _clip_coordinates(x_1, in_width).astype("int32")
+                x_2 = _clip_coordinates(x_2, in_width).astype("int32")
+                x_3 = _clip_coordinates(x_3, in_width).astype("int32")
+
+            elif padding_mode == "reflection":
+                y_ = _reflect_coordinates(y_, in_height)
+                x_0 = _reflect_coordinates(x_0, in_width)
+                x_1 = _reflect_coordinates(x_1, in_width)
+                x_2 = _reflect_coordinates(x_2, in_width)
+                x_3 = _reflect_coordinates(x_3, in_width)
+
+                y_ = _clip_coordinates(y_, in_height).astype("int32")
+                x_0 = _clip_coordinates(x_0, in_width).astype("int32")
+                x_1 = _clip_coordinates(x_1, in_width).astype("int32")
+                x_2 = _clip_coordinates(x_2, in_width).astype("int32")
+                x_3 = _clip_coordinates(x_3, in_width).astype("int32")
+
+            coefficients[i] = cubic_interp_1d(
+                _get_pixel_value(n, c, y_, x_0),
+                _get_pixel_value(n, c, y_, x_1),
+                _get_pixel_value(n, c, y_, x_2),
+                _get_pixel_value(n, c, y_, x_3),
+                x_fraction,
+            )
+
+        return cubic_interp_1d(
+            coefficients[0], coefficients[1], coefficients[2], coefficients[3], y_fraction
+        )
+
+    if method == "bilinear":
+        interpolation = _bilinear_sample
+    elif method == "nearest":
+        interpolation = _nearest_sample
+    else:  # method == "bicubic"
+        interpolation = _bicubic_sample
+
+    return te.compute((batch, in_channel, out_height, out_width), interpolation, tag="grid_sample")
+
+
+def _grid_sample_3d(
+    data, grid, method="bilinear", layout="NCDHW", padding_mode="zeros", align_corners=True
+):
+    """Applies bilinear/nearest sampling to input feature map.
+
+    Given :math:`data` and :math:`grid` assuming NCDHW layout, then the output is computed by
+
+    .. math::
+
+        x_{src} = grid[batch, 0, z_{dst}, y_{dst}, x_{dst}] \\
+        y_{src} = grid[batch, 1, z_{dst}, y_{dst}, x_{dst}] \\
+        z_{src} = grid[batch, 2, z_{dst}, y_{dst}, x_{dst}] \\
+        output[batch, channel, z_{src}, y_{dst}, x_{dst}]
+        = G(data[batch, channel, z_{src}, y_{src}, x_{src})
+
+    :math:`x_{dst}`, :math:`y_{dst}`, :math:`z_{dst}` enumerate all spatial locations
+    in :math:`output`, and :math:`G()` denotes the interpolation method.
+
+    The out-boundary points will be padded with zeros if padding_mode is "zeros", or
+    border pixel value if padding_mode is "border", or
+    inner pixel value if padding_mode is "reflection".
+
+    The left-top corner (-1, -1, -1) and right-bottom corner (1, 1, 1) in grid will be map to
+    (0, 0, 0) and (d - 1, h - 1, w - 1) of data if align_corners is "True", or
+    (-0.5, -0.5, -0.5) and (d - 0.5, h - 0.5, w - 0.5) of data if align_corners is "False".
+
+    The shape of the output will be
+    (data.shape[0], data.shape[1], grid.shape[2], grid.shape[3], grid.shape[4]).
+
+    The operator assumes that :math:`grid` has been normalized to [-1, 1].
+
+    grid_sample often cooperates with affine_grid which generates sampling grids for grid_sample.
+
+    Parameters
+    ----------
+    data : tvm.Tensor
+        5-D with shape [batch, in_channel, in_depth, in_height, in_width]
+
+    grid : tvm.Tensor
+        5-D with shape [batch, 3, out_depth, out_height, out_width]
+
+    method : str
+        The interpolation method "nearest", "bilinear"("trilinear") are supported.
+
+    layout : str
+        The layout of input data and the output.
+
+    padding_mode : str
+        The padding mode for outside grid values, "zeros", "border", "reflection" are supported.
+
+    align_corners: bool
+        Geometrically, we consider the pixels of the input as squares rather than points.
+        If set to "True", the extrema ("-1" and "1") are considered as referring
+        to the center points of the input corner pixels. If set to "False", they
+        are instead considered as referring to the corner points of the input corner
+        pixels, making the sampling more resolution agnostic.
+
+    Returns
+    -------
+    Output : tvm.Tensor
+        5-D with shape [batch, in_channel, out_depth, out_height, out_width]
+    """
+
+    assert method in ("bilinear", "nearest"), f"{method} is not supported"
+    assert padding_mode in ("zeros", "border", "reflection"), f"{padding_mode} is not supported"
+    assert layout == "NCDHW", f"{layout} is not supported"
+
+    batch, in_channel, in_depth, in_height, in_width = data.shape
+    out_depth, out_height, out_width = grid.shape[2:]
+
+    def _get_pixel_value(n, c, d, h, w):
+        return te.if_then_else(
+            te.all(d >= 0, h >= 0, w >= 0, d < in_depth, h < in_height, w < in_width),
+            data[n, c, d, h, w],
+            tirx.const(0.0, dtype=data.dtype),
+        )
+
+    def _compute_source_index(n, d, h, w):
+        z = grid[n, 2, d, h, w]
+        y = grid[n, 1, d, h, w]
+        x = grid[n, 0, d, h, w]
+
+        if align_corners:
+            z = (z + 1) * (in_depth - 1) / 2
+            y = (y + 1) * (in_height - 1) / 2
+            x = (x + 1) * (in_width - 1) / 2
+        else:
+            z = -0.5 + (z + 1) * in_depth / 2
+            y = -0.5 + (y + 1) * in_height / 2
+            x = -0.5 + (x + 1) * in_width / 2
+
+        if padding_mode == "reflection":
+            z = _reflect_coordinates(z, in_depth)
+            y = _reflect_coordinates(y, in_height)
+            x = _reflect_coordinates(x, in_width)
+            z = _clip_coordinates(z, in_depth)
+            y = _clip_coordinates(y, in_height)
+            x = _clip_coordinates(x, in_width)
+        elif padding_mode == "border":
+            z = _clip_coordinates(z, in_depth)
+            y = _clip_coordinates(y, in_height)
+            x = _clip_coordinates(x, in_width)
+
+        return (z, y, x)
+
+    def _clip_coordinates(x, size):
+        return te.min(te.max(x, 0), size - 1)
+
+    def _reflect_coordinates(x, size):
+        def __refelection(x, size, corner_start):
+            def __reflect(index, size, corner_start):
+                index_align_corner = te.abs(corner_start - index)
+                size_times = te.truncdiv(index_align_corner.astype("int32"), size).astype("int32")
+                t = tirx.Mod(size_times, 2)
+                extra = index_align_corner - size_times * size
+                return tirx.if_then_else(
+                    tirx.EQ(t, 0), extra + corner_start, size - extra + corner_start
+                )
+
+            return tirx.if_then_else(
+                tirx.all(x >= corner_start, x <= size + corner_start),
+                x,
+                __reflect(x, size, corner_start),
+            )
+
+        if align_corners:
+            return __refelection(x, size - 1, 0)
+        return __refelection(x, size, -0.5)
+
+    def _trilinear_sample(n, c, d, h, w):
+        z, y, x = _compute_source_index(n, d, h, w)
+        z0 = te.floor(z).astype("int32")
+        y0 = te.floor(y).astype("int32")
+        x0 = te.floor(x).astype("int32")
+        z1 = z0 + tirx.const(1, "int32")
+        y1 = y0 + tirx.const(1, "int32")
+        x1 = x0 + tirx.const(1, "int32")
+
+        return (
+            _get_pixel_value(n, c, z0, y0, x0) * (1 - (x - x0)) * (1 - (y - y0)) * (1 - (z - z0))
+            + _get_pixel_value(n, c, z0, y0, x1) * (x - x0) * (1 - (y - y0)) * (1 - (z - z0))
+            + _get_pixel_value(n, c, z1, y1, x0) * (1 - (x - x0)) * (y - y0) * (z - z0)
+            + _get_pixel_value(n, c, z1, y1, x1) * (x - x0) * (y - y0) * (z - z0)
+            + _get_pixel_value(n, c, z0, y1, x0) * (1 - (x - x0)) * (y - y0) * (1 - (z - z0))
+            + _get_pixel_value(n, c, z1, y0, x1) * (x - x0) * (1 - (y - y0)) * (z - z0)
+            + _get_pixel_value(n, c, z1, y0, x0) * (1 - (x - x0)) * (1 - (y - y0)) * (z - z0)
+            + _get_pixel_value(n, c, z0, y1, x1) * (x - x0) * (y - y0) * (1 - (z - z0))
+        )
+
+    def _nearest_sample(n, c, d, h, w):
+        z, y, x = _compute_source_index(n, d, h, w)
+        z_new = te.nearbyint(z).astype("int32")
+        y_new = te.nearbyint(y).astype("int32")
+        x_new = te.nearbyint(x).astype("int32")
+        return _get_pixel_value(n, c, z_new, y_new, x_new)
+
+    if method == "bilinear":
+        interpolation = _trilinear_sample
+    else:  # method == "nearest"
+        interpolation = _nearest_sample
+
+    return te.compute(
+        (batch, in_channel, out_depth, out_height, out_width), interpolation, tag="grid_sample"
+    )
+
+
+def grid_sample(
+    data, grid, method="bilinear", layout="NCHW", padding_mode="zeros", align_corners=True
+):
+    """Applies grid sampling to input feature map.
+
+    Given :math:`data` and :math:`grid`, then for 4-D the output is computed by
+
+    .. math::
+
+        x_{src} = grid[batch, 0, y_{dst}, x_{dst}] \\
+        y_{src} = grid[batch, 1, y_{dst}, x_{dst}] \\
+        output[batch, channel, y_{dst}, x_{dst}] = G(data[batch, channel, y_{src}, x_{src}])
+
+    :math:`x_{dst}`, :math:`y_{dst}` enumerate all spatial locations in :math:`output`, and
+    :math:`G()` denotes the interpolation function.
+
+    The out-boundary points will be padded with zeros if padding_mode is "zeros", or
+    border pixel value if padding_mode is "border", or
+    inner pixel value if padding_mode is "reflection".
+
+    The left-top corner (-1, -1) and right-bottom corner (1, 1) in grid will be map to
+    (0, 0) and (h - 1, w - 1) of data if align_corners is "True", or
+    (-0.5, -0.5) and (h - 0.5, w - 0.5) of data if align_corners is "False".
+
+    The shape of the output will be
+    4-D (data.shape[0], data.shape[1], grid.shape[2], grid.shape[3]), or
+    5-D (data.shape[0], data.shape[1], grid.shape[2], grid.shape[3], grid.shape[4]).
+
+    The operator assumes that :math:`grid` has been normalized to [-1, 1].
+
+    grid_sample often cooperates with affine_grid which generates sampling grids for grid_sample.
+
+    Parameters
+    ----------
+    data : tvm.Tensor
+        4-D with shape [batch, in_channel, in_height, in_width], or
+        5-D with shape [batch, in_channel, in_depth, in_height, in_width]
+
+    grid : tvm.Tensor
+        4-D with shape [batch, 2, out_height, out_width], or
+        5-D with shape [batch, 3, out_depth, out_height, out_width]
+
+    method : str
+        The interpolation method, 4-D "nearest", "bilinear", "bicubic" and
+        5-D "nearest", "bilinear"("trilinear") are supported.
+
+    layout : str
+        The layout of input data and the output.
+
+    padding_mode : str
+        The padding mode for outside grid values, "zeros", "border", "reflection" are supported.
+
+    align_corners: bool
+        Geometrically, we consider the pixels of the input as squares rather than points.
+        If set to "True", the extrema ("-1" and "1") are considered as referring
+        to the center points of the input corner pixels. If set to "False", they
+        are instead considered as referring to the corner points of the input corner
+        pixels, making the sampling more resolution agnostic.
+
+    Returns
+    -------
+    Output : tvm.Tensor
+        4-D with shape [batch, in_channel, out_height, out_width], or
+        5-D with shape [batch, in_channel, out_depth, out_height, out_width]
+    """
+
+    if len(layout) == 4:
+        compute = _grid_sample_2d
+    elif len(layout) == 5:
+        compute = _grid_sample_3d
+    else:
+        msg = f"layout {layout} is not supported"
+        raise ValueError(msg)
+
+    return compute(data, grid, method, layout, padding_mode, align_corners)

python/tvm/topi/index_put.py

@@ -0,0 +1,165 @@
+# 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
+"""IndexPut operator"""
+
+from tvm import te, tirx
+from tvm.script.ir_builder import IRBuilder
+from tvm.script.ir_builder import tirx as T
+
+from . import utils
+
+
+def index_put(data, indices, values, accumulate=False):
+    """Put values into an array according to indices.
+
+    Parameters
+    ----------
+    data : tvm.te.Tensor
+        The source array to be modified.
+
+    indices : Tuple[tvm.te.Tensor]
+        Tuple of index tensors (can be multi-dimensional) specifying positions.
+        Index tensors are broadcast together following NumPy broadcasting rules.
+
+    values : tvm.te.Tensor
+        The values to place at the specified indices.
+
+    accumulate : bool, optional
+        Whether to accumulate (add) values rather than replace.
+        If True, performs tensor[indices] += values
+        If False, performs tensor[indices] = values
+        Default is False.
+
+    Returns
+    -------
+    ret : tvm.te.Tensor
+    """
+    if not isinstance(indices, list | tuple):
+        indices = [indices]
+
+    # Check indices match data dimensions
+    if len(indices) != len(data.shape):
+        raise ValueError(
+            f"Number of index tensors ({len(indices)}) must match "
+            f"data dimensions ({len(data.shape)})"
+        )
+
+    # Prepare ranges and strides
+    shape = data.shape
+    full_range = 1
+    for dim in shape:
+        full_range *= dim
+
+    index_shapes = [idx.shape for idx in indices]
+    broadcast_ndim = max(len(s) for s in index_shapes)
+    broadcast_shape = []
+
+    for i in range(broadcast_ndim):
+        max_dim = 1
+        for idx_shape in index_shapes:
+            # Right-align shapes
+            dim_idx = len(idx_shape) - broadcast_ndim + i
+            if dim_idx >= 0:
+                dim_size = idx_shape[dim_idx]
+                if not utils.equal_const_int(dim_size, 1):
+                    if utils.equal_const_int(max_dim, 1):
+                        max_dim = dim_size
+                    elif not utils.equal_const_int(dim_size, max_dim):
+                        raise ValueError(f"Cannot broadcast index shapes: {index_shapes}")
+        broadcast_shape.append(max_dim)
+
+    # Compute total number of elements after broadcasting
+    index_len = 1
+    for dim in broadcast_shape:
+        index_len *= dim
+
+    def gen_ir(data_ptr, index_ptrs, values_ptr, out_ptr, reduce_func):
+        data = T.buffer_proxy(data_ptr)
+        indices = [T.buffer_proxy(idx) for idx in index_ptrs]
+        values = T.buffer_proxy(values_ptr)
+        out = T.buffer_proxy(out_ptr)
+
+        with IRBuilder() as ib:
+            with T.seq_scope():
+                with T.parallel(0, full_range) as i:
+                    out[i] = data[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):
+                            # 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][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[k])
+
+            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
+
+    reduce_func = add_func if accumulate else update_func
+
+    # Prepare input buffers
+    in_buffers = [data]
+    in_buffers.extend(indices)
+    in_buffers.append(values)
+
+    out_buf = tirx.decl_buffer(data.shape, data.dtype, "out_buf", layout=None)
+    return te.extern(
+        [data.shape],
+        in_buffers,
+        lambda ins, outs: gen_ir(ins[0], ins[1:-1], ins[-1], outs[0], reduce_func),
+        dtype=data.dtype,
+        out_buffers=[out_buf],
+        name="index_put.generic",
+        tag="index_put.generic",
+    )

python/tvm/topi/math.py

@@ -0,0 +1,883 @@
+# 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.
+"""Elementwise operators"""
+
+# pylint: disable=redefined-builtin,unused-argument
+import tvm
+from tvm import DataTypeCode, te
+
+from . import cpp, tag
+from .utils import get_const_tuple
+
+
+def _require_float_tensor(op_name, x):
+    if not x.dtype.matches_code(DataTypeCode.FLOAT, DataTypeCode.BFLOAT):
+        raise TypeError(f"topi.{op_name} only supports floating-point inputs, but got {x.dtype}")
+    return x
+
+
+def _is_integer_tensor(x):
+    return x.dtype.matches_code(DataTypeCode.INT, DataTypeCode.UINT)
+
+
+@tvm.te.tag_scope(tag=tag.ELEMWISE)
+def identity(x):
+    """Take identity of input x.
+
+    Parameters
+    ----------
+    x : tvm.te.Tensor
+        Input argument.
+
+    Returns
+    -------
+    y : tvm.te.Tensor
+        The result.
+    """
+    # pylint: disable=unnecessary-lambda
+    return te.compute(x.shape, lambda *i: x(*i))
+
+
+@tvm.te.tag_scope(tag=tag.ELEMWISE)
+def negative(x):
+    """Take negation of input x.
+
+    Parameters
+    ----------
+    x : tvm.te.Tensor
+        Input argument.
+
+    Returns
+    -------
+    y : tvm.te.Tensor
+        The result.
+    """
+    # pylint: disable=unnecessary-lambda
+    return te.compute(x.shape, lambda *i: -x(*i))
+
+
+@tvm.te.tag_scope(tag=tag.ELEMWISE)
+def exp(x):
+    """Take exponential of input x.
+
+    Parameters
+    ----------
+    x : tvm.te.Tensor
+        Input argument.
+
+    Returns
+    -------
+    y : tvm.te.Tensor
+        The result.
+    """
+    return te.compute(x.shape, lambda *i: te.exp(x(*i)))
+
+
+@tvm.te.tag_scope(tag=tag.ELEMWISE)
+def erf(x):
+    """Take gauss error function of input x.
+
+    Parameters
+    ----------
+    x : tvm.te.Tensor
+        Input argument.
+
+    Returns
+    -------
+    y : tvm.te.Tensor
+        The result.
+    """
+    return te.compute(x.shape, lambda *i: te.erf(x(*i)))
+
+
+@tvm.te.tag_scope(tag=tag.ELEMWISE)
+def tanh(x):
+    """Take hyperbolic tanh of input x.
+
+    Parameters
+    ----------
+    x : tvm.te.Tensor
+        Input argument.
+
+    Returns
+    -------
+    y : tvm.te.Tensor
+        The result.
+    """
+    return te.compute(x.shape, lambda *i: te.tanh(x(*i)))
+
+
+@tvm.te.tag_scope(tag=tag.ELEMWISE)
+def tan(x):
+    """Take tan of input x.
+
+    Parameters
+    ----------
+    x : tvm.te.Tensor
+        Input argument.
+
+    Returns
+    -------
+    y : tvm.te.Tensor
+        The result.
+    """
+    return te.compute(x.shape, lambda *i: te.tan(x(*i)))
+
+
+@tvm.te.tag_scope(tag=tag.ELEMWISE)
+def cos(x):
+    """Take cos of input x.
+
+    Parameters
+    ----------
+    x : tvm.te.Tensor
+        Input argument.
+
+    Returns
+    -------
+    y : tvm.te.Tensor
+        The result.
+    """
+    return te.compute(x.shape, lambda *i: te.cos(x(*i)))
+
+
+@tvm.te.tag_scope(tag=tag.ELEMWISE)
+def cosh(x):
+    """Take cosh of input x.
+
+    Parameters
+    ----------
+    x : tvm.te.Tensor
+        Input argument.
+
+    Returns
+    -------
+    y : tvm.te.Tensor
+        The result.
+    """
+    return te.compute(x.shape, lambda *i: te.cosh(x(*i)))
+
+
+@tvm.te.tag_scope(tag=tag.ELEMWISE)
+def sin(x):
+    """Take sin of input x.
+
+    Parameters
+    ----------
+    x : tvm.te.Tensor
+        Input argument.
+
+    Returns
+    -------
+    y : tvm.te.Tensor
+        The result.
+    """
+    return te.compute(x.shape, lambda *i: te.sin(x(*i)))
+
+
+@tvm.te.tag_scope(tag=tag.ELEMWISE)
+def sinh(x):
+    """Take sinh of input x.
+
+    Parameters
+    ----------
+    x : tvm.te.Tensor
+        Input argument.
+
+    Returns
+    -------
+    y : tvm.te.Tensor
+        The result.
+    """
+    return te.compute(x.shape, lambda *i: te.sinh(x(*i)))
+
+
+@tvm.te.tag_scope(tag=tag.ELEMWISE)
+def acos(x):
+    """Take arc cos of input x.
+
+    Parameters
+    ----------
+    x : tvm.te.Tensor
+        Input argument.
+
+    Returns
+    -------
+    y : tvm.te.Tensor
+        The result.
+    """
+    x = _require_float_tensor("acos", x)
+    return te.compute(x.shape, lambda *i: te.acos(x(*i)))
+
+
+@tvm.te.tag_scope(tag=tag.ELEMWISE)
+def acosh(x):
+    """Take arc cosh of input x.
+
+    Parameters
+    ----------
+    x : tvm.te.Tensor
+        Input argument.
+
+    Returns
+    -------
+    y : tvm.te.Tensor
+        The result.
+    """
+    x = _require_float_tensor("acosh", x)
+    return te.compute(x.shape, lambda *i: te.acosh(x(*i)))
+
+
+@tvm.te.tag_scope(tag=tag.ELEMWISE)
+def asin(x):
+    """Take arc sin of input x.
+
+    Parameters
+    ----------
+    x : tvm.te.Tensor
+        Input argument.
+
+    Returns
+    -------
+    y : tvm.te.Tensor
+        The result.
+    """
+    x = _require_float_tensor("asin", x)
+    return te.compute(x.shape, lambda *i: te.asin(x(*i)))
+
+
+@tvm.te.tag_scope(tag=tag.ELEMWISE)
+def asinh(x):
+    """Take arc sinh of input x.
+
+    Parameters
+    ----------
+    x : tvm.te.Tensor
+        Input argument.
+
+    Returns
+    -------
+    y : tvm.te.Tensor
+        The result.
+    """
+    x = _require_float_tensor("asinh", x)
+    return te.compute(x.shape, lambda *i: te.asinh(x(*i)))
+
+
+@tvm.te.tag_scope(tag=tag.ELEMWISE)
+def atan(x):
+    """Take atan of input x.
+
+    Parameters
+    ----------
+    x : tvm.te.Tensor
+        Input argument.
+
+    Returns
+    -------
+    y : tvm.te.Tensor
+        The result.
+    """
+    return te.compute(x.shape, lambda *i: te.atan(x(*i)))
+
+
+@tvm.te.tag_scope(tag=tag.ELEMWISE)
+def atanh(x):
+    """Take atanh of input x.
+
+    Parameters
+    ----------
+    x : tvm.te.Tensor
+        Input argument.
+
+    Returns
+    -------
+    y : tvm.te.Tensor
+        The result.
+    """
+    x = _require_float_tensor("atanh", x)
+    return te.compute(x.shape, lambda *i: te.atanh(x(*i)))
+
+
+@tvm.te.tag_scope(tag=tag.ELEMWISE)
+def floor(x):
+    """Take floor of input x.
+
+    Parameters
+    ----------
+    x : tvm.te.Tensor
+        Input argument.
+
+    Returns
+    -------
+    y : tvm.te.Tensor
+        The result.
+    """
+    return te.compute(x.shape, lambda *i: te.floor(x(*i)))
+
+
+@tvm.te.tag_scope(tag=tag.ELEMWISE)
+def ceil(x):
+    """Take ceil of input x.
+
+    Parameters
+    ----------
+    x : tvm.te.Tensor
+        Input argument.
+
+    Returns
+    -------
+    y : tvm.te.Tensor
+        The result.
+    """
+    return te.compute(x.shape, lambda *i: te.ceil(x(*i)))
+
+
+def sign(x):
+    """Returns -1, 0, 1 based on sign of x.
+
+    Parameters
+    ----------
+    x : tvm.te.Tensor
+        Input argument.
+
+    Returns
+    -------
+    y : tvm.te.Tensor
+        The result.
+    """
+    return cpp.sign(x)
+
+
+@tvm.te.tag_scope(tag=tag.ELEMWISE)
+def trunc(x):
+    """Take truncated value of the input of x, element-wise.
+
+    Parameters
+    ----------
+    x : tvm.te.Tensor
+        Input argument.
+
+    Returns
+    -------
+    y : tvm.te.Tensor
+        The result.
+    """
+    return te.compute(x.shape, lambda *i: te.trunc(x(*i)))
+
+
+@tvm.te.tag_scope(tag=tag.ELEMWISE)
+def abs(x):
+    """Take absolute value of the input of x, element-wise.
+
+    Parameters
+    ----------
+    x : tvm.te.Tensor
+        Input argument.
+
+    Returns
+    -------
+    y : tvm.te.Tensor
+        The result.
+    """
+    return te.compute(x.shape, lambda *i: te.abs(x(*i)))
+
+
+@tvm.te.tag_scope(tag=tag.ELEMWISE)
+def isnan(x):
+    """Check if value of x is NaN, element-wise.
+
+    Parameters
+    ----------
+    x : tvm.te.Tensor
+        Input argument.
+
+    Returns
+    -------
+    y : tvm.te.Tensor
+        The result.
+    """
+    return te.compute(x.shape, lambda *i: te.isnan(x(*i)))
+
+
+@tvm.te.tag_scope(tag=tag.ELEMWISE)
+def isfinite(x):
+    """Check if value of x is finite, element-wise.
+
+    Parameters
+    ----------
+    x : tvm.te.Tensor
+        Input argument.
+
+    Returns
+    -------
+    y : tvm.te.Tensor
+        The result.
+    """
+    return te.compute(x.shape, lambda *i: te.isfinite(x(*i)))
+
+
+@tvm.te.tag_scope(tag=tag.ELEMWISE)
+def isinf(x):
+    """Check if value of x is infinite, element-wise.
+
+    Parameters
+    ----------
+    x : tvm.te.Tensor
+        Input argument.
+
+    Returns
+    -------
+    y : tvm.te.Tensor
+        The result.
+    """
+    return te.compute(x.shape, lambda *i: te.isinf(x(*i)))
+
+
+@tvm.te.tag_scope(tag=tag.ELEMWISE)
+def round(x):
+    """Round elements of x to nearest integer using ties-to-even (banker's rounding).
+
+    Ties are broken by rounding to the nearest even integer, matching the ONNX Round
+    specification and IEEE 754 default rounding mode.
+
+    Parameters
+    ----------
+    x : tvm.te.Tensor
+        Input argument.
+
+    Returns
+    -------
+    y : tvm.te.Tensor
+        The result.
+    """
+    return te.compute(x.shape, lambda *i: te.nearbyint(x(*i)))
+
+
+def log(x):
+    """Take logarithm of input x.
+
+    Parameters
+    ----------
+    x : tvm.te.Tensor
+        Input argument.
+
+    Returns
+    -------
+    y : tvm.te.Tensor
+        The result.
+    """
+    if x.dtype.matches_code(DataTypeCode.INT):
+        x = te.compute(x.shape, lambda *i: x(*i).astype("float32"))
+    return te.compute(x.shape, lambda *i: te.log(x(*i)), tag=tag.ELEMWISE)
+
+
+def log2(x):
+    """Take logarithm to the base 2 of input x.
+
+    Parameters
+    ----------
+    x : tvm.te.Tensor
+        Input argument.
+
+    Returns
+    -------
+    y : tvm.te.Tensor
+        The result.
+    """
+    if x.dtype.matches_code(DataTypeCode.INT):
+        x = te.compute(x.shape, lambda *i: x(*i).astype("float32"))
+    return te.compute(x.shape, lambda *i: te.log2(x(*i)), tag=tag.ELEMWISE)
+
+
+def log10(x):
+    """Take logarithm to the base 10 of input x.
+
+    Parameters
+    ----------
+    x : tvm.te.Tensor
+        Input argument.
+
+    Returns
+    -------
+    y : tvm.te.Tensor
+        The result.
+    """
+    if x.dtype.matches_code(DataTypeCode.INT):
+        x = te.compute(x.shape, lambda *i: x(*i).astype("float32"))
+    return te.compute(x.shape, lambda *i: te.log10(x(*i)), tag=tag.ELEMWISE)
+
+
+@tvm.te.tag_scope(tag=tag.ELEMWISE)
+def sqrt(x):
+    """Take square root of input x.
+
+    Parameters
+    ----------
+    x : tvm.te.Tensor
+        Input argument.
+
+    Returns
+    -------
+    y : tvm.te.Tensor
+        The result.
+    """
+    if x.dtype.matches_code(DataTypeCode.INT):
+        x = te.compute(x.shape, lambda *i: x(*i).astype("float32"))
+    return te.compute(x.shape, lambda *i: te.sqrt(x(*i)))
+
+
+@tvm.te.tag_scope(tag=tag.ELEMWISE)
+def rsqrt(x):
+    """Take inverse square root of input x.
+
+    Parameters
+    ----------
+    x : tvm.te.Tensor
+        Input argument.
+
+    Returns
+    -------
+    y : tvm.te.Tensor
+        The result.
+    """
+    if x.dtype.matches_code(DataTypeCode.INT):
+        x = te.compute(x.shape, lambda *i: x(*i).astype("float32"))
+    return te.compute(x.shape, lambda *i: te.rsqrt(x(*i)))
+
+
+@tvm.te.tag_scope(tag=tag.ELEMWISE)
+def sigmoid(x):
+    """Take sigmoid tanh of input x.
+
+    Parameters
+    ----------
+    x : tvm.te.Tensor
+        Input argument.
+
+    Returns
+    -------
+    y : tvm.te.Tensor
+        The result.
+    """
+    return te.compute(x.shape, lambda *i: te.sigmoid(x(*i)))
+
+
+@tvm.te.tag_scope(tag=tag.ELEMWISE)
+def left_shift(x, n):
+    """Take n bits left shift of input x.
+
+    Parameters
+    ----------
+    x : tvm.te.Tensor
+        Input argument.
+    n : int
+        Number of bits.
+
+    Returns
+    -------
+    y : tvm.te.Tensor
+        The result.
+    """
+    return te.compute(x.shape, lambda *i: x(*i) << n)
+
+
+@tvm.te.tag_scope(tag=tag.ELEMWISE)
+def right_shift(x, n):
+    """Take n bits right shift of input x.
+
+    Parameters
+    ----------
+    x : tvm.te.Tensor
+        Input argument.
+    n : int
+        Number of bits.
+
+    Returns
+    -------
+    y : tvm.te.Tensor
+        The result.
+    """
+    return te.compute(x.shape, lambda *i: x(*i) >> n)
+
+
+@tvm.te.tag_scope(tag=tag.ELEMWISE)
+def clip(x, a_min, a_max):
+    """Clip (limit) the values in an array. Given an interval, values
+    outside the interval are clipped to the interval edges.
+
+    Parameters
+    ----------
+    x : tvm.te.Tensor
+        Input argument.
+    a_min : tvm.tirx.Expr
+        Minimum value.
+    a_max : tvm.tirx.Expr
+        Maximum value.
+
+    Returns
+    -------
+    y : tvm.te.Tensor
+        The result.
+    """
+
+    def _compute(*indices):
+        value = x(*indices)
+        const_min = (
+            tvm.tirx.Cast(value.ty, a_min)
+            if tvm.ir.is_prim_expr(a_min)
+            else tvm.tirx.const(a_min, value.ty)
+        )
+        const_max = (
+            tvm.tirx.Cast(value.ty, a_max)
+            if tvm.ir.is_prim_expr(a_max)
+            else tvm.tirx.const(a_max, value.ty)
+        )
+        return tvm.te.max(tvm.te.min(value, const_max), const_min)
+
+    return te.compute(x.shape, _compute)
+
+
+@tvm.te.tag_scope(tag=tag.ELEMWISE)
+def fixed_point_multiply(x, multiplier, shift):
+    """Fixed point multiplication between data and a fixed point
+    constant expressed as multiplier * 2^(-shift), where multiplier
+    is a Q-number with 31 fractional bits
+
+    Parameters
+    ----------
+    x : tvm.te.Tensor or Expr
+        Input argument.
+    multiplier : int
+        Multiplier of a fixed floating point number described as multiplier*2^(-shift).
+    shift : int
+        Shift of a fixed floating point number described as multiplier*2^(-shift).
+
+    Returns
+    -------
+    y : tvm.te.Tensor
+        The result.
+    """
+
+    def _compute(*indices):
+        value = x(*indices)
+        return tvm.tirx.q_multiply_shift(
+            value,
+            tvm.tirx.const(multiplier, "int32"),
+            tvm.tirx.const(31, "int32"),
+            tvm.tirx.const(shift, "int32"),
+        )
+
+    return te.compute(x.shape, _compute)
+
+
+@tvm.te.tag_scope(tag=tag.BROADCAST)
+def fixed_point_multiply_per_axis(
+    x: te.Tensor,
+    y: te.Tensor,
+    lshift: te.Tensor,
+    rshift: te.Tensor,
+    is_lshift_required: int,
+    is_rshift_required: int,
+    axes,
+):
+    """Fixed point multiplication between data and a fixed point constant expressed as
+    multiplier * 2^(-shift), where multiplier is a Q-number with 31 fractional bits
+
+    Parameters
+    ----------
+    x : tvm.te.Tensor
+        Input argument.
+    y : tvm.te.Tensor
+        Multiplier of a fixed floating point number described as multiplier*2^(-shift).
+    lshift : tvm.te.Tensor
+        Left shifts of a fixed floating point number described as multiplier*2^(-shift).
+    rshift : tvm.te.Tensor
+        Right shifts of a fixed floating point number described as multiplier*2^(-shift).
+    is_lshift_required : int
+        Whether we need to do left shift or not.
+    is_rshift_required : int
+        Whether we need to do right shift or not.
+
+    Returns
+    -------
+    z : tvm.te.Tensor
+        The result.
+    """
+
+    def _compute(*indices):
+        elements = []
+        for element in get_const_tuple(axes):
+            elements += [indices[element]]
+        param_indices = tuple(elements)
+
+        value = x(*indices)
+        m = y(*param_indices)
+        l_shift = lshift(*param_indices)
+        r_shift = rshift(*param_indices)
+        return tvm.tirx.q_multiply_shift_per_axis(
+            value,
+            m,
+            l_shift,
+            r_shift,
+            tvm.tirx.const(31, "int32"),
+            tvm.tirx.const(is_lshift_required, "bool"),
+            tvm.tirx.const(is_rshift_required, "bool"),
+        )
+
+    return te.compute(x.shape, _compute)
+
+
+def cast(x, dtype, span=None):
+    """Cast input to specified data type.
+
+    Parameters
+    ----------
+    x : tvm.te.Tensor or Expr
+        Input argument.
+
+    dtype : str
+        Data type.
+
+    span : Optional[Span]
+        The location of the cast in the source.
+
+    Returns
+    -------
+    y : tvm.te.Tensor
+        The result.
+    """
+    if isinstance(x, te.tensor.Tensor):
+        return te.compute(x.shape, lambda *i: x(*i).astype(dtype), tag=tag.ELEMWISE)
+    # pylint: disable=import-outside-toplevel
+    from tvm.tirx import _ffi_api
+
+    return _ffi_api._cast(dtype, x, span)
+
+
+def reinterpret(x, dtype):
+    """Reinterpret input to specified data type.
+
+    Parameters
+    ----------
+    x : tvm.te.Tensor
+        Input argument.
+
+    dtype : str
+        Data type.
+
+    Returns
+    -------
+    y : tvm.te.Tensor
+        The result.
+    """
+    return cpp.reinterpret(x, dtype)
+
+
+def fast_exp(x):
+    """Take exponential of input x using fast_exp implementation
+
+    Parameters
+    ----------
+    x : tvm.te.Tensor
+        Input argument.
+
+    Returns
+    -------
+    y : tvm.te.Tensor
+        The result.
+    """
+    if _is_integer_tensor(x):
+        x = cast(x, "float32")
+    return cpp.fast_exp(x, x.dtype, tag.ELEMWISE)
+
+
+def fast_tanh(x):
+    """Take hyperbolic tangent of input x using fast_tanh implementation
+
+    Parameters
+    ----------
+    x : tvm.te.Tensor
+        Input argument.
+
+    Returns
+    -------
+    y : tvm.te.Tensor
+        The result.
+    """
+    if _is_integer_tensor(x):
+        x = cast(x, "float32")
+    return cpp.fast_tanh(x, x.dtype, tag.ELEMWISE)
+
+
+def fast_erf(x):
+    """Take gauss error function of input x using fast_erf implementation.
+
+    Parameters
+    ----------
+    x : tvm.te.Tensor
+        Input argument.
+
+    Returns
+    -------
+    y : tvm.te.Tensor
+        The result.
+    """
+    return cpp.fast_erf(x, x.dtype, tag.ELEMWISE)
+
+
+def ceil_log2(x):
+    """Compute integer ceil log2 with a special code path for vulkan
+    SPIR-V does not support log2 on fp64. Instead, we compute integer ceil_log2 via clz
+    intrinsic when the target is vulkan.
+
+    Parameters
+    ----------
+    x : tvm.te.Tensor
+        Input argument.
+
+    Returns
+    -------
+    y : tvm.te.Tensor
+        The result.
+    """
+    if not tvm.ir.is_prim_expr(x):
+        x = tvm.tirx.const(x)
+
+    if x.ty.matches_code(DataTypeCode.FLOAT, DataTypeCode.BFLOAT):
+        return tvm.tirx.ceil(tvm.tirx.log2(x))
+
+    target = tvm.target.Target.current()
+
+    if target is not None:
+        target_name = target.kind.name
+        if "vulkan" in target_name:
+            clz = tvm.tirx.clz(x)
+            bits = x.ty.dtype.bits
+            res = tvm.tirx.if_then_else(x & (x - 1) == 0, bits - clz - 1, bits - clz)
+            if res.ty != x.ty:
+                return cast(res, x.ty)
+            return res
+
+        if "adreno" in str(target.attrs.get("device", "")) or target_name in [
+            "metal",
+            "rocm",
+            "webgpu",
+        ]:
+            return cast(tvm.tirx.ceil(tvm.tirx.log2(cast(x, "float32"))), x.ty)
+
+    return cast(tvm.tirx.ceil(tvm.tirx.log2(cast(x, "float64"))), x.ty)

python/tvm/topi/nn/__init__.py

@@ -0,0 +1,57 @@
+# 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=wildcard-import
+"""Neural network operators"""
+
+from .conv1d import *
+from .conv2d import *
+from .conv3d import *
+from .correlation import *
+from .deformable_conv2d import *
+from .depthwise_conv2d import *
+from .elemwise import *
+from .dilate import *
+from .flatten import *
+from .dense import *
+from .mapping import *
+from .pooling import *
+from .softmax import *
+from .conv3d_transpose import *
+from .conv2d_transpose import *
+from .conv1d_transpose import *
+from .bnn import *
+from .qnn import *
+from .upsampling import *
+from .instance_norm import instance_norm
+from .layer_norm import layer_norm
+from .group_norm import group_norm
+from .rms_norm import rms_norm
+from .local_response_norm import *
+from .bitserial_conv2d import *
+from .bitserial_dense import *
+from .batch_matmul import *
+from .batch_norm import *
+from .pad import *
+from .fifo_buffer import *
+from .depth_to_space import *
+from .space_to_depth import *
+from .space_to_batch_nd import *
+from .batch_to_space_nd import *
+from .loss import *
+from .lstm import *

python/tvm/topi/nn/batch_matmul.py

@@ -0,0 +1,152 @@
+# 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.
+# ruff: noqa: E731
+"""Batch matrix multiplication"""
+
+# pylint: disable=invalid-name
+import logging
+
+import tvm
+from tvm import te
+
+from ..utils import get_const_tuple
+
+logger = logging.getLogger("topi")
+
+
+def batch_matmul(
+    tensor_a,
+    tensor_b,
+    oshape=None,
+    out_dtype=None,
+    transpose_a=False,
+    transpose_b=True,
+    auto_scheduler_rewritten_layout="",
+    meta_schedule_original_shape=None,
+):
+    """Compute batch matrix multiplication of `tensor_a` and `tensor_b`.
+
+    Both `tensor_a` and `tensor_b` can be transposed. For legacy reason, we use NT format
+    (transpose_a=False, transpose_b=True) by default.
+
+    Parameters
+    ----------
+    tensor_a : tvm.te.Tensor
+        3-D with shape [batch, M, K] or [batch, K, M].
+
+    tensor_b : tvm.te.Tensor
+        3-D with shape [batch, K, N] or [batch, N, K].
+
+    oshape : List[Optional]
+        Explicit intended output shape of the computation. Can be useful in cases
+        with dynamic input shapes.
+
+    out_dtype : Optional[str]
+        Specifies the output data type for mixed precision batch matmul.
+
+    transpose_a : Optional[bool] = False
+        Whether the first tensor is in transposed format.
+
+    transpose_b : Optional[bool] = True
+        Whether the second tensor is in transposed format.
+
+    auto_scheduler_rewritten_layout: Optional[str] = ""
+        The layout after auto-scheduler's layout rewrite pass.
+
+    meta_schedule_original_shape: Optional[List[Expr]] = None
+        The original shape of the tensor
+
+    Returns
+    -------
+    output : tvm.te.Tensor
+        3-D with shape [batch, M, N]
+    """
+    assert len(tensor_a.shape) == 3, "tensor_a only support 3-dim"
+    if transpose_a:
+        XB, XK, XI = get_const_tuple(tensor_a.shape)
+    else:
+        XB, XI, XK = get_const_tuple(tensor_a.shape)
+    if auto_scheduler_rewritten_layout:
+        raise RuntimeError("LEGACY-FLOW triggered, to be removed")
+    if meta_schedule_original_shape:
+        raise RuntimeError("LEGACY-FLOW triggered, to be removed")
+
+    assert len(tensor_b.shape) == 3, "tensor_b only support 3-dim"
+    if transpose_b:
+        YB, YJ, YK = get_const_tuple(tensor_b.shape)
+    else:
+        YB, YK, YJ = get_const_tuple(tensor_b.shape)
+
+    assert XK == YK or isinstance(YK, tvm.tirx.expr.Var), "shapes of x and y are inconsistent"
+    k = te.reduce_axis((0, XK), name="k")
+    if oshape is None:
+        assert XB == YB or XB == 1 or YB == 1, "batch dimension doesn't match"
+        batch = (
+            tvm.tirx.expr.Var("batch", "int32")
+            if isinstance(XB, tvm.tirx.expr.Var) or isinstance(YB, tvm.tirx.expr.Var)
+            else te.max(XB, YB)
+        )
+        oshape = (batch, XI, YJ)
+    if out_dtype is None:
+        out_dtype = tensor_a.dtype
+        if tensor_a.dtype != tensor_b.dtype:
+            logger.warning(
+                "tensor_a has different data type with tensor_b: %s, %s",
+                tensor_a.dtype,
+                tensor_b.dtype,
+            )
+
+    if (transpose_a, transpose_b) == (True, True):
+        compute_lambda = lambda b, i, j: te.sum(
+            tensor_a[b if XB != 1 else 0, k, i].astype(out_dtype)
+            * tensor_b[b if YB != 1 else 0, j, k].astype(out_dtype),
+            axis=k,
+        )
+        compute_name = "T_batch_matmul_TT"
+    elif (transpose_a, transpose_b) == (True, False):
+        compute_lambda = lambda b, i, j: te.sum(
+            tensor_a[b if XB != 1 else 0, k, i].astype(out_dtype)
+            * tensor_b[b if YB != 1 else 0, k, j].astype(out_dtype),
+            axis=k,
+        )
+        compute_name = "T_batch_matmul_TN"
+    elif (transpose_a, transpose_b) == (False, True):
+        compute_lambda = lambda b, i, j: te.sum(
+            tensor_a[b if XB != 1 else 0, i, k].astype(out_dtype)
+            * tensor_b[b if YB != 1 else 0, j, k].astype(out_dtype),
+            axis=k,
+        )
+        compute_name = "T_batch_matmul_NT"
+    else:  # (transpose_a, transpose_b) == (False, False):
+        compute_lambda = lambda b, i, j: te.sum(
+            tensor_a[b if XB != 1 else 0, i, k].astype(out_dtype)
+            * tensor_b[b if YB != 1 else 0, k, j].astype(out_dtype),
+            axis=k,
+        )
+        compute_name = "T_batch_matmul_NN"
+
+    output = te.compute(
+        oshape,
+        compute_lambda,
+        name=compute_name,
+        tag="batch_matmul",
+        attrs={"layout_free_placeholders": [tensor_b]},
+    )
+    if auto_scheduler_rewritten_layout:
+        raise RuntimeError("LEGACY-FLOW triggered, to be removed")
+
+    return output

python/tvm/topi/nn/batch_norm.py

@@ -0,0 +1,146 @@
+# 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.
+"""Batch normalization."""
+
+from functools import reduce
+
+from tvm import te, topi
+
+
+def batch_norm(
+    data: te.Tensor,
+    gamma: te.Tensor,
+    beta: te.Tensor,
+    moving_mean: te.Tensor,
+    moving_var: te.Tensor,
+    axis: int | None = None,
+    epsilon: float | None = None,
+    center: bool | None = None,
+    scale: bool | None = None,
+    training: bool | None = None,
+    momentum: float | None = None,
+) -> list[te.Tensor]:
+    """Batch normalization layer (Ioffe and Szegedy, 2014).
+
+    Normalizes the input at each batch, i.e. applies a transformation
+    that maintains the mean activation close to 0 and the activation
+    standard deviation close to 1.
+
+    Parameters
+    ----------
+    data : tvm.te.Tensor
+        Input to be batch-normalized.
+
+    gamma : tvm.te.Tensor
+        Scale factor to be applied to the normalized tensor.
+
+    beta : tvm.te.Tensor
+        Offset to be applied to the normalized tensor.
+
+    moving_mean : tvm.te.Tensor
+        Running mean of input.
+
+    moving_var : tvm.te.Tensor
+        Running variance of input.
+
+    axis : int, optional, default=1
+        Specify along which shape axis the normalization should occur.
+
+    epsilon : float, optional, default=1e-5
+        Small float added to variance to avoid dividing by zero.
+
+    center : bool, optional, default=True
+        If True, add offset of beta to normalized tensor, If False,
+        beta is ignored.
+
+    scale : bool, optional, defualt=True
+        If True, scale normalized tensor by gamma. If False, gamma
+        is ignored.
+
+    training : bool, optional, defualt=False
+        Indicating whether it is in training mode. If True, update
+        moving_mean and moving_var.
+
+    momentum : float, optional, default=0.1
+        The value used for the moving_mean and moving_var update.
+
+    Returns
+    -------
+    output : list of tvm.te.Tensor
+        Normalized data with same shape as input
+
+    moving_mean : tvm.te.Tensor
+        Running mean of input.
+
+    moving_var : tvm.te.Tensor
+        Running variance of input.
+    """
+    if axis is None:
+        axis = 1
+
+    if epsilon is None:
+        epsilon = 1e-5
+
+    if center is None:
+        center = True
+
+    if scale is None:
+        scale = True
+
+    if training is None:
+        training = False
+
+    if momentum is None:
+        momentum = 0.1
+
+    shape = [1] * len(data.shape)
+    shape[axis] = data.shape[axis]
+    data_mean = None
+    data_var = None
+
+    if training:
+        reduce_axes = list(range(len(data.shape)))
+        reduce_axes.remove(axis)
+        shape_prod = reduce(lambda x, y: x * y, [data.shape[ax] for ax in reduce_axes], 1)
+        data_mean = topi.sum(data, axis=reduce_axes) / shape_prod
+        data_mean_rs = topi.reshape(data_mean, shape)
+        data_var = (
+            topi.sum((data - data_mean_rs) * (data - data_mean_rs), axis=reduce_axes) / shape_prod
+        )
+        data_var_rs = topi.reshape(data_var, shape)
+        out = (data - data_mean_rs) / topi.math.sqrt(data_var_rs + epsilon)
+    else:
+        moving_mean_rs = topi.reshape(moving_mean, shape)
+        moving_var_rs = topi.reshape(moving_var, shape)
+        out = (data - moving_mean_rs) / topi.math.sqrt(moving_var_rs + epsilon)
+
+    if scale:
+        out = out * topi.reshape(gamma, shape)
+    if center:
+        out = out + topi.reshape(beta, shape)
+
+    if training:
+        assert 0 <= momentum <= 1, "the valid momentum range is [0, 1]."
+        return [
+            out,
+            (1 - momentum) * moving_mean + momentum * data_mean,
+            (1 - momentum) * moving_var + momentum * data_var,
+        ]
+
+    # Moving mean and var aren't updated during test. To avoid
+    # placeholder reuse, we multiply by 1 and return them.
+    return [out, moving_mean * 1, moving_var * 1]

python/tvm/topi/nn/batch_to_space_nd.py

@@ -0,0 +1,49 @@
+# 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
+"""TVM operator batch_to_space_nd compute."""
+
+from . import cpp
+
+
+def batch_to_space_nd(data, block_shape, crop_begin_list, crop_end_list):
+    """Perform space to batch transformation on the data
+
+    Parameters
+    ----------
+    data : tvm.te.Tensor
+        N-D Tensor with shape [batch, spatial_shape, remaining_shapes],
+        where spatial_shape has M dimensions.
+
+    block_shape : list of ints
+        list of size [M] where M is number of spatial dims, specifies block
+        size for each spatial dimension.
+
+    crop_begin_list : list of ints
+        list of shape [M] where M is number of spatial dims, specifies
+        begin crop size for each spatial dimension.
+
+    crop_end_list : list of ints
+        list of shape [M] where M is number of spatial dims, specifies
+        end crop size for each spatial dimension.
+
+    Returns
+    -------
+    output : tvm.te.Tensor
+    """
+
+    return cpp.nn.batch_to_space_nd(data, block_shape, crop_begin_list, crop_end_list)

python/tvm/topi/nn/bitserial_conv2d.py

@@ -0,0 +1,276 @@
+# 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-arguments
+# pylint: disable=unused-argument, redefined-builtin
+"""Bitserial Conv2D operators"""
+
+import tvm
+from tvm import te
+
+from ..utils import get_const_tuple
+from .bitserial_util import bitpack
+from .pad import pad
+from .utils import get_pad_tuple
+
+
+def bitserial_conv2d_nchw(
+    data,
+    kernel,
+    stride,
+    padding,
+    activation_bits,
+    weight_bits,
+    pack_dtype="uint32",
+    out_dtype="int16",
+    unipolar=True,
+):
+    """Bitserial Conv2D operator.
+
+    Parameters
+    ----------
+    data : tvm.te.Tensor
+        4-D with shape [batch, in_channel, in_height, in_width]
+
+    kernel : tvm.te.Tensor
+        4-D with shape [num_filter, in_channel, filter_height, filter_width]
+
+    stride : int or a list/tuple of two ints
+        stride size, or [stride_height, stride_width]
+
+    padding : int or a list/tuple of two or four ints
+        padding size, [pad_height, pad_width], [pad_top, pad_left, pad_down, pad_right]
+
+    activation_bits: int
+        number of bits used for activations/input elements
+
+    weight_bits: int
+        number of bits used for weight elements
+
+    out_dtype: str
+        return type of convolution
+
+    pack_dtype: str
+        bit packing type
+
+    unipolar: bool
+        if binarization style is in unipolar 1/0 format, instead of bipolar -1/+1 format
+
+    Returns
+    -------
+    output : tvm.te.Tensor
+        4-D with shape [batch, out_channel, out_height, out_width]
+    """
+    assert isinstance(stride, int) or len(stride) == 2
+    Input_q = bitpack(data, activation_bits, pack_axis=1, bit_axis=2, pack_type=pack_dtype)
+    if len(kernel.shape) == 4:
+        Filter_q = bitpack(kernel, weight_bits, pack_axis=1, bit_axis=4, pack_type=pack_dtype)
+    else:
+        Filter_q = kernel
+    batch, in_channel, activation_bits, in_height, in_width = Input_q.shape
+    num_filter, _, kernel_h, kernel_w, weight_bits = Filter_q.shape
+
+    if isinstance(padding, int) or (isinstance(padding, tuple | list) and len(padding) == 2):
+        TPAD, LPAD, DPAD, RPAD = get_pad_tuple(padding, kernel)
+    else:
+        TPAD, LPAD, DPAD, RPAD = padding
+    pad_before = [0, 0, 0, TPAD, LPAD]
+    pad_after = [0, 0, 0, DPAD, RPAD]
+
+    PadInput_q = pad(Input_q, pad_before, pad_after, name="pad_temp")
+    # compute the output shape
+    if isinstance(stride, int):
+        stride_h = stride_w = stride
+    else:
+        stride_h, stride_w = stride
+    out_channel = num_filter
+    out_height = (in_height - kernel_h + TPAD + DPAD) // stride_h + 1
+    out_width = (in_width - kernel_w + LPAD + RPAD) // stride_w + 1
+
+    rc = te.reduce_axis((0, in_channel), name="rc")
+    ry = te.reduce_axis((0, kernel_h), name="ry")
+    rx = te.reduce_axis((0, kernel_w), name="rx")
+    b1 = te.reduce_axis((0, activation_bits), name="b1")
+    b2 = te.reduce_axis((0, weight_bits), name="b2")
+
+    if unipolar:
+
+        def _conv(nn, ff, yy, xx):
+            b1b2 = (b1 + b2).astype(out_dtype)
+            return te.sum(
+                (
+                    (
+                        tvm.tirx.popcount(
+                            PadInput_q[nn, rc, b1, yy * stride_h + ry, xx * stride_w + rx]
+                            & Filter_q[ff, rc, ry, rx, b2]
+                        )
+                        - tvm.tirx.popcount(
+                            PadInput_q[nn, rc, b1, yy * stride_h + ry, xx * stride_w + rx]
+                            & ~Filter_q[ff, rc, ry, rx, b2]
+                        )
+                    )
+                    << (b1b2)
+                ).astype(out_dtype),
+                axis=[rc, ry, rx, b2, b1],
+            ).astype(out_dtype)
+
+    else:
+
+        def _conv(nn, ff, yy, xx):
+            b1b2 = (b1 + b2).astype(out_dtype)
+            return te.sum(
+                (
+                    tvm.tirx.popcount(
+                        PadInput_q[nn, rc, b1, yy * stride_h + ry, xx * stride_w + rx]
+                        & Filter_q[ff, rc, ry, rx, b2]
+                    )
+                    << (b1b2)
+                ).astype(out_dtype),
+                axis=[rc, ry, rx, b2, b1],
+            ).astype(out_dtype)
+
+    return te.compute(
+        (batch, out_channel, out_height, out_width),
+        _conv,
+        name="Conv2dOutput",
+        tag="bitserial_conv2d_nchw",
+    )
+
+
+def bitserial_conv2d_nhwc(
+    data,
+    kernel,
+    stride,
+    padding,
+    activation_bits,
+    weight_bits,
+    pack_dtype="uint32",
+    out_dtype="int16",
+    unipolar=True,
+):
+    """Bitserial Conv2D operator.
+
+    Parameters
+    ----------
+    data : tvm.te.Tensor
+        4-D with shape [batch, in_height, in_width, in_channel]
+
+    kernel : tvm.te.Tensor
+        4-D with shape [filter_height, filter_width, in_channel, num_filter]
+
+    stride : int or a list/tuple of two ints
+        stride size, or [stride_height, stride_width]
+
+    padding : int or a list/tuple of two or four ints
+        padding size, [pad_height, pad_width], [pad_top, pad_left, pad_down, pad_right]
+
+    activation_bits: int
+        number of bits used for activations/input elements
+
+    weight_bits: int
+        number of bits used for weight elements
+
+    out_dtype: str
+        return type of convolution
+
+    pack_dtype: str
+        bit packing type
+
+    unipolar: bool
+        if binarization style is in unipolar 1/0 format, instead of bipolar -1/+1 format
+
+    Returns
+    -------
+    output : tvm.te.Tensor
+        4-D with shape [batch, out_height, out_width, out_channel]
+    """
+    assert isinstance(stride, int) or len(stride) == 2
+    Input_q = bitpack(data, activation_bits, pack_axis=3, bit_axis=4, pack_type=pack_dtype)
+    if len(kernel.shape) == 4:
+        Filter_q = bitpack(kernel, weight_bits, pack_axis=2, bit_axis=4, pack_type=pack_dtype)
+        kernel_h, kernel_w, _, num_filter, _ = get_const_tuple(Filter_q.shape)
+    else:
+        Filter_q = kernel
+        kernel_h, kernel_w, _, _, num_filter = get_const_tuple(Filter_q.shape)
+    batch, in_height, in_width, in_channel_q, _ = get_const_tuple(Input_q.shape)
+
+    if isinstance(padding, int) or (isinstance(padding, tuple | list) and len(padding) == 2):
+        TPAD, LPAD, DPAD, RPAD = get_pad_tuple(padding, kernel)
+    else:
+        TPAD, LPAD, DPAD, RPAD = padding
+    pad_before = [0, TPAD, LPAD, 0, 0]
+    pad_after = [0, DPAD, RPAD, 0, 0]
+
+    # compute the output shape
+    if isinstance(stride, int):
+        stride_h = stride_w = stride
+    else:
+        stride_h, stride_w = stride
+    out_channel = num_filter
+    out_height = (in_height - kernel_h + TPAD + DPAD) // stride_h + 1
+    out_width = (in_width - kernel_w + LPAD + RPAD) // stride_w + 1
+    PadInput_q = pad(Input_q, pad_before, pad_after, name="PaddedInput")
+
+    rc = te.reduce_axis((0, in_channel_q), name="rc")
+    ry = te.reduce_axis((0, kernel_h), name="ry")
+    rx = te.reduce_axis((0, kernel_w), name="rx")
+    b1 = te.reduce_axis((0, activation_bits), name="b1")
+    b2 = te.reduce_axis((0, weight_bits), name="b2")
+
+    if unipolar:
+
+        def _conv(nn, yy, xx, ff):
+            b1b2 = (b1 + b2).astype(out_dtype)
+            return te.sum(
+                (
+                    (
+                        tvm.tirx.popcount(
+                            PadInput_q[nn, yy * stride_h + ry, xx * stride_w + rx, rc, b1]
+                            & Filter_q[ry, rx, rc, ff, b2]
+                        )
+                        - tvm.tirx.popcount(
+                            PadInput_q[nn, yy * stride_h + ry, xx * stride_w + rx, rc, b1]
+                            & ~Filter_q[ry, rx, rc, ff, b2]
+                        )
+                    )
+                    << b1b2
+                ).astype(out_dtype),
+                axis=[rc, ry, rx, b2, b1],
+            )
+
+    else:
+
+        def _conv(nn, yy, xx, ff):
+            b1b2 = (b1 + b2).astype(out_dtype)
+            return te.sum(
+                (
+                    tvm.tirx.popcount(
+                        PadInput_q[nn, yy * stride_h + ry, xx * stride_w + rx, rc, b1]
+                        & Filter_q[ry, rx, rc, ff, b2]
+                    )
+                    << b1b2
+                ).astype(out_dtype),
+                axis=[rc, ry, rx, b2, b1],
+            )
+
+    conv = te.compute(
+        (batch, out_height, out_width, out_channel),
+        _conv,
+        name="Conv2dOutput",
+        tag="bitserial_conv2d_nhwc",
+    )
+
+    return conv

python/tvm/topi/nn/bitserial_dense.py

@@ -0,0 +1,82 @@
+# 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-arguments
+"""Bitserial Dense operator."""
+
+import tvm
+from tvm import te
+from tvm.topi.utils import get_const_tuple
+
+from .bitserial_util import bitpack
+
+
+def bitserial_dense(
+    data, weight, data_bits, weight_bits, pack_dtype="uint32", out_dtype="int16", unipolar=True
+):
+    """The default implementation of bitserial dense in topi.
+
+    Parameters
+    ----------
+    data : tvm.te.Tensor
+        2-D with shape [batch, in_dim]
+    weight : tvm.te.Tensor
+        2-D with shape [out_dim, in_dim] or
+        3-D with shape [out_dim, weight_bits, in_dim]
+    Returns
+    -------
+    output : tvm.te.Tensor
+        2-D with shape [batch, out_dim]
+    """
+    data_packed = bitpack(data, data_bits, pack_axis=1, bit_axis=1, pack_type=pack_dtype)
+    if len(weight.shape) == 2:
+        weight_packed = bitpack(weight, weight_bits, pack_axis=1, bit_axis=1, pack_type=pack_dtype)
+    else:
+        weight_packed = weight
+    Y, DB, K = get_const_tuple(data_packed.shape)
+    X, WB, _ = get_const_tuple(weight_packed.shape)
+
+    oshape = (Y, X)
+    k = te.reduce_axis((0, K), name="k")
+    db = te.reduce_axis((0, DB), name="db")
+    wb = te.reduce_axis((0, WB), name="wb")
+
+    matmul_unipolar = te.compute(
+        oshape,
+        lambda i, j: te.sum(
+            (
+                tvm.tirx.popcount(weight_packed[j, wb, k] & data_packed[i, db, k])
+                - tvm.tirx.popcount(~weight_packed[j, wb, k] & data_packed[i, db, k])
+            ).astype(out_dtype)
+            << (db + wb).astype(out_dtype),
+            axis=[wb, db, k],
+        ),
+        tag="bitserial_dense_unipolar",
+    )
+
+    matmul = te.compute(
+        oshape,
+        lambda i, j: te.sum(
+            tvm.tirx.popcount(weight_packed[j, wb, k] & data_packed[i, db, k]).astype(out_dtype)
+            << (db + wb).astype(out_dtype),
+            axis=[wb, db, k],
+        ),
+        tag="bitserial_dense",
+    )
+
+    if unipolar:
+        return matmul_unipolar
+    return matmul

python/tvm/topi/nn/bitserial_util.py

@@ -0,0 +1,110 @@
+# 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-arguments
+"""Utility functions for bitserial operators"""
+
+import numpy as np
+
+import tvm
+from tvm import te
+from tvm.topi.transform import concatenate
+
+from ..utils import get_const_int
+
+
+def bitpack(data, bits, pack_axis, bit_axis, pack_type, name="QuantizeInput"):
+    """Packs data into format necessary for bitserial computation
+
+    Parameters
+    ----------
+    data : tvm.te.Tensor
+        The input tvm tensor
+    bits : int
+        Number of bits to use for packing
+    pack_axis : int
+       index of the axis to pack in data
+    bit_axis : int
+       index of axis to place bit axis in resulting packed data
+    pack_type : str
+        Data type for packing, must be one of: ['uint8', 'uint16', 'uint32', 'uint64']
+    name : Optional[str] = "QuantizeInput"
+        Name for the operation
+    """
+    ishape = data.shape
+    n = len(ishape)
+    if pack_type == "uint8":
+        data_width = 8
+    elif pack_type == "uint16":
+        data_width = 16
+    elif pack_type == "uint32":
+        data_width = 32
+    elif pack_type == "uint64":
+        data_width = 64
+
+    # Data must be in multiples of the data_width
+    assert get_const_int(ishape[pack_axis]) % data_width == 0, "Not a multiple of word size"
+
+    shape_vec = list(ishape)
+    shape_vec[pack_axis] = shape_vec[pack_axis] // data_width
+    shape_vec.insert(bit_axis, 1)
+    bitserial_oshape = tuple(shape_vec)
+    masks = np.array([0x1, 0x2, 0x4, 0x8, 0x10, 0x20, 0x40, 0x80])
+
+    # pack axis shifts if bit axis comes before
+    if bit_axis <= pack_axis:
+        pack_axis += 1
+
+    def _bitpack(*indices):
+        packed_data = [tvm.tirx.const(0, pack_type)] * bits
+        for k in range(data_width):
+            # Translate indices for packed data back to original
+            idx = [0] * n
+            j = 0
+            for i in range(n + 1):
+                if i == bit_axis:
+                    continue
+                if i == pack_axis:
+                    idx[j] = indices[i] * data_width + k
+                else:
+                    idx[j] = indices[i]
+                j += 1
+
+            element = data(*idx)
+            for b in range(bits):
+                extracted_bit = ((element & tvm.tirx.const(masks[b], "int32")) >> b).astype(
+                    pack_type
+                )
+                packed_data[b] = packed_data[b] | extracted_bit
+                if k < data_width - 1:
+                    packed_data[b] = packed_data[b] << 1
+
+            if k == data_width - 1:
+                return tuple(packed_data)
+        return tuple(packed_data)
+
+    output_tuple = te.compute(bitserial_oshape, _bitpack, name=name, tag="bitpack")
+
+    if bits > 1:
+        return concatenate(output_tuple, axis=bit_axis)
+    return output_tuple
+
+
+def binary_op_multiplier(pack_dtype):
+    """ "Returns number of bits packed into
+    pack_dtype: string
+        pack type for the operator (must be a uint)"""
+    return int(pack_dtype[4:])

python/tvm/topi/nn/bnn.py

@@ -0,0 +1,99 @@
+# 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.
+"""Binary Neural Network (BNN) Operators"""
+
+import tvm
+from tvm import te
+
+from .. import tag
+from ..utils import get_const_int, simplify
+
+
+def binarize_pack(data, axis=None, name="PackedInput"):
+    """Binarization and bit-packing along a certain axis.
+
+    Parameters
+    ----------
+    data : tvm.te.Tensor
+        n-D input, can be any layout.
+
+    axis : None or int
+        The axis along which to do binarization and bit-packing,
+        default is the last axis.
+
+    name : str, optional
+        The name prefix operators generate.
+
+    Returns
+    -------
+    output : tvm.te.Tensor
+        n-D, the same layout as input, dtype is uint32.
+    """
+    ishape = data.shape
+    if axis is None:
+        axis = len(ishape) - 1
+    assert get_const_int(ishape[axis]) % 32 == 0
+    n = len(ishape)
+    oshape = tuple(simplify(ishape[i] // 32) if i == axis else ishape[i] for i in range(n))
+
+    def _binarize_pack(*indices):
+        start_idx = [indices[i] * 32 if i == axis else indices[i] for i in range(n)]
+        packed = tvm.tirx.const(0, "uint32")
+        for j in range(32):
+            idx = [start_idx[i] + j if i == axis else start_idx[i] for i in range(n)]
+            sign = (data(*idx) >= 0).astype("uint32")
+            packed = packed | sign
+            if j == 31:
+                return packed
+            packed = packed << 1
+        raise RuntimeError("not resach")
+
+    return te.compute(oshape, _binarize_pack, name=name, tag="binarize_pack")
+
+
+def binary_dense(data, weight):
+    """Binary matrix multiplication using xor and bit-count.
+
+    Parameters
+    ----------
+    data : tvm.te.Tensor
+        2-D with shape [batch, in_dim], dtype is uint32.
+
+    weight : tvm.te.Tensor
+        2-D with shape [out_dim, in_dim], dtype is uint32.
+
+    Returns
+    -------
+    output : tvm.te.Tensor
+        2-D with shape [batch, out_dim], dtype is float32.
+    """
+    assert data.dtype == "uint32" and weight.dtype == "uint32", (
+        "dtype of data and weight should be uint32"
+    )
+    assert len(data.shape) == 2 and len(weight.shape) == 2, "only support 2-dim binary dense"
+    batch, in_dim = data.shape
+    out_dim, _ = weight.shape
+    k = te.reduce_axis((0, in_dim), name="k")
+    matmul = te.compute(
+        (batch, out_dim),
+        lambda i, j: te.sum(tvm.tirx.popcount(data[i, k] ^ weight[j, k]), axis=k),
+        tag="binary_dense",
+    )
+
+    return te.compute(
+        (batch, out_dim), lambda i, j: 32 * in_dim - 2.0 * matmul(i, j), tag=tag.ELEMWISE
+    )

python/tvm/topi/nn/conv1d.py

@@ -0,0 +1,141 @@
+# 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, unused-variable, unused-argument
+"""1D convolution operators."""
+
+from .conv2d import conv
+
+
+def conv1d(
+    data,
+    kernel,
+    strides=1,
+    padding="VALID",
+    dilation=1,
+    groups=1,
+    data_layout="NCW",
+    kernel_layout="",
+    out_dtype=None,
+):
+    """1D convolution forward operator.
+
+    Parameters
+    ----------
+    data : tvm.te.Tensor
+        3-D input shape [batch, in_channel, in_width] for data_layout == 'NCW'
+        and [batch, in_width, in_channel] for data_layout == 'NWC'
+
+    kernel : tvm.te.Tensor
+        3-D kernel with shape [num_filter, in_channel, filter_size] for kernel_layout == 'OIW'
+        and [filter_size, in_channel, num_filter] for kernel_layout == 'WIO'
+
+    strides : int or tuple
+        The spatial stride along width
+
+    padding : int or str
+        Padding size, or ['VALID', 'SAME']
+
+    dilation : int or tuple
+        Dilation rate if convolution should be dilated.
+
+    data_layout : str
+        How input data is laid out, must be one of ['NCW', 'NWC']
+
+    kernel_layout: Optiona[str]
+        The layout of the kernel. If unspecified, use default layout. "OIW" if data_layout == "NCW",
+        "WIO" if data_layout == "NWC".
+
+    out_dtype : str
+        The output data type. If None then output is same type as input.
+    """
+    return conv(
+        data, kernel, strides, padding, dilation, groups, data_layout, kernel_layout, out_dtype
+    )
+
+
+def conv1d_nwc(data, kernel, strides=1, padding="VALID", dilation=1, out_dtype=None):
+    """1D convolution in NWC layout. See :py:func:`conv` for details on parameters"""
+    return conv(data, kernel, strides, padding, dilation, 1, "NWC", "WIO", out_dtype=out_dtype)
+
+
+def conv1d_ncw(data, kernel, strides=1, padding="VALID", dilation=1, out_dtype=None):
+    """1D convolution in NCW layout. See :py:func:`conv` for details on parameters"""
+    return conv(data, kernel, strides, padding, dilation, 1, "NCW", "OIW", out_dtype=out_dtype)
+
+
+def group_conv1d_nwc(
+    data, kernel, strides=1, padding="VALID", dilation=1, groups=1, out_dtype=None
+):
+    """1D convolution forward operator for NWC layout.
+
+    Parameters
+    ----------
+    data : tvm.te.Tensor
+        3-D with shape [batch, in_width, in_channel]
+
+    kernel : tvm.te.Tensor
+        3-D with shape [filter_size, in_channel, num_filter]
+
+    strides : int or tuple
+        The spatial stride along width
+
+    padding : int, tuple, or str
+        Padding size can be an integer for equal padding,
+        a tuple of (left, right) or a string in ['VALID', 'SAME'].
+
+    dilation : int or tuple
+        Dilation rate if convolution should be dilated.
+
+    groups : int
+        Number of groups
+
+    out_dtype : str
+        The output data type. If None then output is same type as input.
+    """
+    return conv(data, kernel, strides, padding, dilation, groups, "NWC", "WIO", out_dtype=out_dtype)
+
+
+def group_conv1d_ncw(
+    data, kernel, strides=1, padding="VALID", dilation=1, groups=1, out_dtype=None
+):
+    """1D convolution forward operator for NCW layout.
+
+    Parameters
+    ----------
+    data : tvm.te.Tensor
+        3-D with shape [batch, in_channel, in_width]
+
+    kernel : tvm.te.Tensor
+        3-D with shape [num_filter, in_channel, filter_size]
+
+    strides : int or tuple
+        The spatial stride along width
+
+    padding : int, tuple, or str
+        Padding size can be an integer for equal padding,
+        a tuple of (left, right) or a string in ['VALID', 'SAME'].
+
+    dilation : int or tuple
+        Dilation rate if convolution should be dilated.
+
+    groups : int
+        Number of groups
+
+    out_dtype : str
+        The output data type. If None then output is same type as input.
+    """
+    return conv(data, kernel, strides, padding, dilation, groups, "NCW", "OIW", out_dtype=out_dtype)

python/tvm/topi/nn/conv1d_transpose.py

@@ -0,0 +1,217 @@
+# 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, unused-variable, unused-argument
+"""Transposed 1D convolution operators (sometimes called Deconvolution)."""
+
+from tvm import te
+
+from ..utils import simplify
+from .dilate import dilate
+from .pad import pad
+from .utils import get_pad_tuple1d
+
+
+def _conv1d_transpose_ncw_preprocess(data, kernel, stride, padding, out_dtype, output_padding):
+    """Preprocess data and kernel to make the compute pattern
+    of conv1d_transpose the same as conv1d.
+
+    Parameters
+    ----------
+    data : tvm.te.Tensor
+        3-D with shape [batch, in_channel, in_width]
+
+    kernel : tvm.te.Tensor
+        3-D with shape [in_channel, num_filter, filter_width]
+
+    stride : ints
+        The spatial stride along width
+
+    padding : int or str
+        Padding size, or ['VALID', 'SAME']
+
+    out_dtype : str
+        The output data type. This is used for mixed precision.
+
+    output_padding : ints
+        Used to recover the actual output shape in case there are more
+        than one possible shape.  Must be smaller than stride.
+
+    Returns
+    -------
+    data_pad : tvm.te.Tensor
+        Padded input data. 3-D with shape [batch, in_channel, in_width]
+
+    kernel: tvm.te.Tensor
+        Transformed kernel. 3-D with shape [num_filter, in_channel, filter_width]
+    """
+    # some pre-processing and prelimnary checks
+    if out_dtype is None:
+        out_dtype = data.dtype
+
+    # dilate and pad
+    if isinstance(stride, tuple | list):
+        stride = stride[0]
+    if isinstance(output_padding, tuple | list):
+        output_padding = output_padding[0]
+
+    _, channels_in, _ = data.shape
+    _, channels_out, kernel_width = kernel.shape
+    assert output_padding < stride
+    channels_out = simplify(channels_out)
+    data_dilate = dilate(data, [1, 1, stride], name="data_dilate")
+    pad_left, pad_right = get_pad_tuple1d(padding, (kernel_width,))
+    pad_left = kernel_width - 1 - pad_left
+    pad_right = kernel_width - 1 - pad_right + output_padding
+    data_pad = pad(data_dilate, [0, 0, pad_left], [0, 0, pad_right], name="data_pad")
+
+    # transform kernel layout from IOW to OIW, and rotate kernel by 180 degrees
+    kernel = te.compute(
+        (channels_out, channels_in, kernel_width),
+        lambda o, i, w: kernel[i][o][kernel_width - 1 - w],
+        name="kernel",
+    )
+    return data_pad, kernel
+
+
+def conv1d_transpose_ncw(data, kernel, stride, padding, out_dtype, output_padding):
+    """Transposed 1D convolution ncw forward operator.
+
+    Parameters
+    ----------
+    data : tvm.te.Tensor
+        3-D with shape [batch, in_channel, in_width]
+
+    kernel : tvm.te.Tensor
+        3-D with shape [in_channel, num_filter, filter_width]
+
+    stride : ints
+        The spatial stride along width
+
+    padding : int or str
+        Padding size, or ['VALID', 'SAME']
+
+    out_dtype : str
+        The output data type. This is used for mixed precision.
+
+    output_padding : ints
+        Used to recover the actual output shape in case there are more
+        than one possible shape.  Must be smaller than stride.
+
+    Returns
+    -------
+    output : tvm.te.Tensor
+        3-D with shape [batch, out_channel, out_width]
+
+    """
+
+    batch, channels_in, _ = data.shape
+    _, channels_out, kernel_width = kernel.shape
+
+    data_pad, transformed_kernel = _conv1d_transpose_ncw_preprocess(
+        data, kernel, stride, padding, out_dtype, output_padding
+    )
+
+    # convolution
+    _, _, data_width = data_pad.shape
+    out_w = simplify(data_width - kernel_width + 1)
+    dc = te.reduce_axis((0, channels_in), name="dc")
+    dw = te.reduce_axis((0, kernel_width), name="dw")
+    output = te.compute(
+        (batch, channels_out, out_w),
+        lambda b, c, w: te.sum(
+            data_pad[b, dc, w + dw].astype(out_dtype)
+            * transformed_kernel[c, dc, dw].astype(out_dtype),
+            axis=[dc, dw],
+        ),
+        tag="conv1d_transpose_ncw",
+    )
+
+    return output
+
+
+def group_conv1d_transpose_ncw(data, kernel, stride, padding, out_dtype, output_padding, groups):
+    """Transposed 1D group convolution ncw forward operator.
+
+    Parameters
+    ----------
+    data : tvm.te.Tensor
+        3-D with shape [batch, in_channel, in_width]
+
+    kernel : tvm.te.Tensor
+        3-D with shape [in_channel, num_filter, filter_width]
+
+    stride : ints
+        The spatial stride along width
+
+    padding : int or str
+        Padding size, or ['VALID', 'SAME']
+
+    out_dtype : str
+        The output data type. This is used for mixed precision.
+
+    output_padding : ints
+        Used to recover the actual output shape in case there are more
+        than one possible shape.  Must be smaller than stride.
+
+     groups : int
+        number of groups
+
+    Returns
+    -------
+    output : tvm.te.Tensor
+        3-D with shape [batch, out_channel, out_width]
+
+    """
+    if groups == 1:
+        return conv1d_transpose_ncw(data, kernel, stride, padding, out_dtype, output_padding)
+
+    _, in_channels, _ = data.shape
+
+    assert in_channels % groups == 0, (
+        f"input channels {in_channels} must divide group size {groups}"
+    )
+
+    data_pad, transformed_kernel = _conv1d_transpose_ncw_preprocess(
+        data, kernel, stride, padding, out_dtype, output_padding
+    )
+
+    batch, in_channels, in_w = data_pad.shape
+    out_c, _, filter_w = transformed_kernel.shape
+
+    # convolution stage
+    out_channels = simplify(out_c * groups)
+    out_w = simplify(in_w - filter_w + 1)
+    dc = te.reduce_axis((0, in_channels // groups), name="dc")
+    dw = te.reduce_axis((0, filter_w), name="dw")
+
+    # data: batch, in_channels, out_w
+    # weight: out_channels // G, in_channels, out_w
+    return te.compute(
+        (batch, out_channels, out_w),
+        lambda b, c, w: te.sum(
+            data_pad[
+                b, c // (out_channels // groups) * (in_channels // groups) + dc, w + dw
+            ].astype(out_dtype)
+            * transformed_kernel[
+                c % (out_channels // groups),
+                c // (out_channels // groups) * (in_channels // groups) + dc,
+                dw,
+            ].astype(out_dtype),
+            axis=[dc, dw],
+        ),
+        tag="group_conv1d_transpose_ncw",
+    )

python/tvm/topi/nn/conv2d_transpose.py

@@ -0,0 +1,246 @@
+# 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, unused-variable, unused-argument
+# ruff: noqa: F821
+"""Transposed 2D convolution operators (sometimes called Deconvolution)."""
+
+import collections
+
+from tvm import te
+
+from ..utils import simplify
+from .dilate import dilate
+from .pad import pad
+from .utils import get_pad_tuple
+
+
+def _ntuple(n):
+    def parse(x):
+        if isinstance(x, collections.abc.Iterable):
+            assert len(x) == n, f"Input can only have {n} elements, but got {len(x)} instead: {x}."
+            return x
+        return tuple(repeat(x, n))
+
+    return parse
+
+
+_single = _ntuple(1)
+_pair = _ntuple(2)
+_triple = _ntuple(3)
+_quadruple = _ntuple(4)
+
+
+def conv2d_transpose_nchw(Input, Filter, strides, padding, out_dtype, output_padding):
+    """Transposed 2D convolution nchw forward operator.
+
+    Parameters
+    ----------
+    Input : tvm.te.Tensor
+        4-D with shape [batch, in_channel, in_height, in_width]
+
+    Filter : tvm.te.Tensor
+        4-D with shape [in_channel, num_filter, filter_height, filter_width]
+
+    strides : tuple of two ints
+        The spatial stride along height and width
+
+    padding : int or str
+        Padding size, or ['VALID', 'SAME']
+
+    out_dtype : str
+        The output data type. This is used for mixed precision.
+
+    output_padding : tuple of ints
+        Used to get the right output shape for gradients
+
+    Returns
+    -------
+    Output : tvm.te.Tensor
+        4-D with shape [batch, out_channel, out_height, out_width]
+    """
+    return declaration_conv2d_transpose_impl(
+        Input, Filter, strides, padding, out_dtype, output_padding=output_padding
+    )
+
+
+def conv2d_transpose_nchw_preprocess(data, kernel, strides, padding, out_dtype, output_padding):
+    """Preprocess data and kernel to make the compute pattern
+    of conv2d_transpose the same as conv2d"""
+    batch, in_c, in_h, in_w = data.shape
+    _, out_c, filter_h, filter_w = kernel.shape
+    stride_h, stride_w = strides
+    opad_h, opad_w = output_padding
+    assert opad_h < stride_h and opad_w < stride_w
+    # dilate data
+    data_dilate = dilate(data, [1, 1, stride_h, stride_w], name="data_dilate")
+    # pad data
+    fpad_top, fpad_left, fpad_bottom, fpad_right = get_pad_tuple(padding, (filter_h, filter_w))
+    bpad_top = filter_h - 1 - fpad_top
+    bpad_bottom = filter_h - 1 - fpad_bottom + opad_h
+    bpad_left = filter_w - 1 - fpad_left
+    bpad_right = filter_w - 1 - fpad_right + opad_w
+    data_pad = pad(
+        data_dilate, [0, 0, bpad_top, bpad_left], [0, 0, bpad_bottom, bpad_right], name="data_pad"
+    )
+    # transform kernel layout from IOHW to OIHW, and rotate kernel by 180 degrees
+    kernel_transform = te.compute(
+        (out_c, in_c, filter_h, filter_w),
+        lambda o, i, h, w: kernel[i][o][filter_h - 1 - h][filter_w - 1 - w],
+        name="kernel_transform",
+    )
+    return data_pad, kernel_transform
+
+
+def declaration_conv2d_transpose_impl(data, kernel, strides, padding, out_dtype, output_padding):
+    """Implementation of conv2d transpose"""
+    data_pad, kernel_transform = conv2d_transpose_nchw_preprocess(
+        data, kernel, strides, padding, out_dtype, output_padding
+    )
+    batch, in_c, in_h, in_w = data_pad.shape
+    out_c, _, filter_h, filter_w = kernel_transform.shape
+
+    # convolution stage
+    out_c = simplify(out_c)
+
+    out_h = simplify(in_h - filter_h + 1)
+    out_w = simplify(in_w - filter_w + 1)
+    dc = te.reduce_axis((0, in_c), name="dc")
+    dh = te.reduce_axis((0, filter_h), name="dh")
+    dw = te.reduce_axis((0, filter_w), name="dw")
+
+    Output = te.compute(
+        (batch, out_c, out_h, out_w),
+        lambda b, c, h, w: te.sum(
+            data_pad[b, dc, h + dh, w + dw].astype(out_dtype)
+            * kernel_transform[c, dc, dh, dw].astype(out_dtype),
+            axis=[dc, dh, dw],
+        ),
+        tag="conv2d_transpose_nchw",
+    )
+
+    return Output
+
+
+def group_conv2d_transpose_nchw(data, kernel, stride, padding, out_dtype, output_padding, groups):
+    """Group convolution operator in NCHW layout.
+
+    Parameters
+    ----------
+    data : tvm.te.Tensor
+        4-D with shape [batch, in_channel, in_height, in_width]
+
+    kernel : tvm.te.Tensor
+        4-D with shape [in_channel, out_channel // groups, filter_height, filter_width]
+
+    stride : int or a list/tuple of two ints
+        Stride size, or [stride_height, stride_width]
+
+    padding : int or a list/tuple of 2 or 4 ints
+        padding size, or
+        [pad_height, pad_width] for 2 ints, or
+        [pad_top, pad_left, pad_bottom, pad_right] for 4 ints
+
+    out_dtype : str
+        The output data type. This is used for mixed precision.
+
+    output_padding : tuple of ints
+        Used to get the right output shape for gradients
+
+    groups : int
+        number of groups
+
+    out_dtype : str
+        The output type. This is used for mixed precision.
+
+    Returns
+    -------
+    Output : tvm.te.Tensor
+        4-D with shape [batch, out_channel, out_height, out_width]
+    """
+    if groups == 1:
+        return conv2d_transpose_nchw(data, kernel, stride, padding, out_dtype, output_padding)
+
+    # some pre-processing and prelimnary checks
+    if out_dtype is None:
+        out_dtype = data.dtype
+
+    batch, in_channels, in_h, in_w = data.shape
+    _, out_c, filter_h, filter_w = kernel.shape
+    assert in_channels % groups == 0, (
+        f"input channels {in_channels} must divide group size {groups}"
+    )
+    # assert out_c % groups == 0, f"output channels {in_c} must divide group size {groups}"
+
+    strides = _pair(stride)
+    # padding = _pair(padding)
+    # output_padding = _pair(output_padding)
+    # dilation = _pair(dilation)
+
+    stride_h, stride_w = strides
+    opad_h, opad_w = output_padding
+    assert opad_h < stride_h and opad_w < stride_w, (
+        f"[{output_padding}] opad_h:{opad_h} < stride_h:{stride_h} \
+        and opad_w:{opad_w} < stride_w:{stride_w} does not satisfy."
+    )
+    # dilate data
+    data_dilate = dilate(data, [1, 1, stride_h, stride_w], name="data_dilate")
+    # pad data
+    fpad_top, fpad_left, fpad_bottom, fpad_right = get_pad_tuple(padding, (filter_h, filter_w))
+    bpad_top = filter_h - 1 - fpad_top
+    bpad_bottom = filter_h - 1 - fpad_bottom + opad_h
+    bpad_left = filter_w - 1 - fpad_left
+    bpad_right = filter_w - 1 - fpad_right + opad_w
+    data_pad = pad(
+        data_dilate, [0, 0, bpad_top, bpad_left], [0, 0, bpad_bottom, bpad_right], name="data_pad"
+    )
+    # transform kernel layout from IOHW to OIHW, and rotate kernel by 180 degrees
+    kernel_transform = te.compute(
+        (out_c, in_channels, filter_h, filter_w),
+        lambda i, o, h, w: kernel[o][i][filter_h - 1 - h][filter_w - 1 - w],
+        name="kernel_transform",
+    )
+
+    batch, in_channels, in_h, in_w = data_pad.shape
+    out_c, _, filter_h, filter_w = kernel_transform.shape
+
+    # convolution stage
+    out_channels = simplify(out_c * groups)
+
+    out_h = simplify(in_h - filter_h + 1)
+    out_w = simplify(in_w - filter_w + 1)
+    dc = te.reduce_axis((0, in_channels // groups), name="dc")
+    dh = te.reduce_axis((0, filter_h), name="dh")
+    dw = te.reduce_axis((0, filter_w), name="dw")
+
+    # data: batch, in_channels, out_h, out_w
+    # weight: out_channels // G, in_channels, out_h, out_w
+    return te.compute(
+        (batch, out_channels, out_h, out_w),
+        lambda b, c, h, w: te.sum(
+            data_pad[
+                b, c // (out_channels // groups) * (in_channels // groups) + dc, h + dh, w + dw
+            ].astype(out_dtype)
+            * kernel_transform[
+                c % (out_channels // groups),
+                c // (out_channels // groups) * (in_channels // groups) + dc,
+                dh,
+                dw,
+            ].astype(out_dtype),
+            axis=[dc, dh, dw],
+        ),
+        tag="group_conv2d_transpose_nchw",
+    )

python/tvm/topi/nn/conv3d.py

@@ -0,0 +1,169 @@
+# 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, unused-variable, too-many-locals
+# pylint: disable=unused-argument, redefined-builtin, no-else-return
+"""Conv3D operators"""
+
+from tvm import te
+
+from ..utils import get_const_tuple
+from .conv2d import conv
+from .winograd_util import winograd_transform_matrices
+
+
+def conv3d_ncdhw(Input, Filter, stride, padding, dilation, groups, out_dtype=None):
+    """Conv3D operator in NCDHW layout.
+
+    Parameters
+    ----------
+    Input : tvm.te.Tensor
+        5-D with shape [batch, in_channel, in_depth, in_height, in_width]
+
+    Filter : tvm.te.Tensor
+        5-D with shape [num_filter, in_channel, filter_depth, filter_height, filter_width]
+
+    stride : int or a list/tuple of three ints
+        Stride size, or [strid_depth, stride_height, stride_width]
+
+    padding : int or str
+        Padding size, or ['VALID', 'SAME']
+
+    dilation: int or a list/tuple of three ints
+        dilation size, or [dilation_depth, dilation_height, dilation_width]
+
+    groups: int
+        Number of groups.
+
+    Returns
+    -------
+    Output : tvm.te.Tensor
+        5-D with shape [batch, out_channel, out_depth, out_height, out_width]
+    """
+    return conv(Input, Filter, stride, padding, dilation, groups, "NCDHW", "OIDHW", out_dtype)
+
+
+def conv3d_ndhwc(
+    Input,
+    Filter,
+    stride,
+    padding,
+    dilation,
+    groups,
+    out_dtype="float32",
+    auto_scheduler_rewritten_layout="",
+    meta_schedule_origin_shape=None,
+):
+    """Convolution operator in NDHWC layout.
+
+    Parameters
+    ----------
+    Input : tvm.te.Tensor
+        5-D with shape [batch, in_depth, in_height, in_width, in_channel]
+
+    Filter : tvm.te.Tensor
+        5-D with shape [filter_depth, filter_height, filter_width, in_channel, num_filter]
+
+    stride : int or a list/tuple of three ints
+        Stride size, or [stride_depth, stride_height, stride_width]
+
+    padding : int or str
+        Padding size, or ['VALID', 'SAME']
+
+    dilation: int or a list/tuple of three ints
+        dilation size, or [dilation_depth, dilation_height, dilation_width]
+
+    groups: int
+        Number of groups.
+
+    out_dtype: str = "float32",
+        The type of output tensor
+
+    auto_scheduler_rewritten_layout: str = ""
+        The layout after auto-scheduler's layout rewrite pass.
+
+    meta_schedule_origin_shape: Optional[List[Expr]] = None
+        The original shape of the input tensor.
+
+    Returns
+    -------
+    Output : tvm.te.Tensor
+        5-D with shape [batch, out_depth, out_height, out_width, out_channel]
+    """
+    return conv(
+        Input,
+        Filter,
+        stride,
+        padding,
+        dilation,
+        groups,
+        "NDHWC",
+        "DHWIO",
+        out_dtype,
+        auto_scheduler_rewritten_layout,
+        meta_schedule_origin_shape,
+    )
+
+
+def conv3d_winograd_weight_transform(kernel, tile_size):
+    """Weight transformation for 3D winograd
+
+    Parameters
+    ----------
+    kernel: Tensor
+        The raw kernel tensor with layout "NCDHW".
+    tile_size: int
+        Tile size of winograd transform. e.g. 2 for F(2x2, 3x3) and 4 for F(4x4, 3x3)
+
+    Returns
+    -------
+    output : tvm.te.Tensor
+        5-D with shape [alpha, alpha, alpha, CO, CI]
+    """
+    CO, CI, KD, KH, KW = get_const_tuple(kernel.shape)
+
+    depth_transform = 2 < KD < 8 and KD == KH
+
+    if depth_transform:
+        assert KD == KH == KW, "Only support NxNxN kernel"
+    else:
+        assert KH == KW, "Only supports DxNxN kernel"
+
+    r = tile_size + KH - 1
+
+    r_kh = te.reduce_axis((0, KH), name="r_kh")
+    r_kw = te.reduce_axis((0, KW), name="r_kw")
+    _, _, G = winograd_transform_matrices(tile_size, KH, kernel.dtype)
+    if depth_transform:
+        shape = (r, r, r, CO, CI)
+        r_kd = te.reduce_axis((0, KD), name="r_kd")
+        return te.compute(
+            shape,
+            lambda omg, eps, nu, co, ci: te.sum(
+                kernel[co][ci][r_kd][r_kh][r_kw] * G[omg][r_kd] * G[eps][r_kh] * G[nu][r_kw],
+                axis=[r_kd, r_kh, r_kw],
+            ),
+            name="transform_weight",
+        )
+    else:
+        shape = (r, r, KD, CO, CI)
+        return te.compute(
+            shape,
+            lambda eps, nu, d, co, ci: te.sum(
+                kernel[co][ci][d][r_kh][r_kw] * G[eps][r_kh] * G[nu][r_kw], axis=[r_kh, r_kw]
+            ),
+            name="transform_weight",
+        )

python/tvm/topi/nn/conv3d_transpose.py

@@ -0,0 +1,200 @@
+# 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, unused-variable, unused-argument
+"""Transposed 3D convolution operators (sometimes called Deconvolution)."""
+
+from tvm import te
+
+from ..utils import simplify
+from .dilate import dilate
+from .pad import pad
+from .utils import get_pad_tuple3d
+
+
+def conv3d_transpose_ncdhw(Input, Filter, strides, padding, out_dtype, output_padding):
+    """Transposed 3D convolution ncdhw forward operator.
+
+    Parameters
+    ----------
+    Input : tvm.te.Tensor
+        5-D with shape [batch, in_channel, in_depth, in_height, in_width]
+
+    Filter : tvm.te.Tensor
+        5-D with shape [in_channel, num_filter, filter_depth, filter_height, filter_width]
+
+    strides : int or a list/tuple of three ints
+        The spatial stride along depth,height and width
+
+    padding : int or str
+        Padding size, or ['VALID', 'SAME']
+
+    out_dtype : str
+        The output data type. This is used for mixed precision.
+
+    output_padding : tuple of ints
+        Used to get the right output shape for gradients
+
+    Returns
+    -------
+    Output : tvm.te.Tensor
+        5-D with shape [batch, out_channel, out_depth, out_height, out_width]
+    """
+    return declaration_conv3d_transpose_impl(
+        Input, Filter, strides, padding, out_dtype, output_padding
+    )
+
+
+def conv3d_transpose_ncdhw_preprocess(data, kernel, strides, padding, out_dtype, output_padding):
+    """Preprocess data and kernel to make the compute pattern
+    of conv3d_transpose the same as conv3d"""
+    batch, in_c, in_d, in_h, in_w = data.shape
+    _, out_c, filter_d, filter_h, filter_w = kernel.shape
+    stride_d, stride_h, stride_w = strides
+    opad_d, opad_h, opad_w = output_padding
+    assert opad_d < stride_d and opad_h < stride_h and opad_w < stride_w
+    # dilate data
+    data_dilate = dilate(data, [1, 1, stride_d, stride_h, stride_w], name="data_dilate")
+    # pad data
+    fpad_front, fpad_top, fpad_left, fpad_back, fpad_bottom, fpad_right = get_pad_tuple3d(
+        padding, (filter_d, filter_h, filter_w)
+    )
+    bpad_front = filter_d - 1 - fpad_front
+    bpad_back = filter_d - 1 - fpad_back + opad_d
+    bpad_top = filter_h - 1 - fpad_top
+    bpad_bottom = filter_h - 1 - fpad_bottom + opad_h
+    bpad_left = filter_w - 1 - fpad_left
+    bpad_right = filter_w - 1 - fpad_right + opad_w
+    data_pad = pad(
+        data_dilate,
+        [0, 0, bpad_front, bpad_top, bpad_left],
+        [0, 0, bpad_back, bpad_bottom, bpad_right],
+        name="data_pad",
+    )
+    # transform kernel layout from IODHW to OIDHW, and rotate kernel by 180 degrees
+    kernel_transform = te.compute(
+        (out_c, in_c, filter_d, filter_h, filter_w),
+        lambda o, i, d, h, w: kernel[i][o][filter_d - 1 - d][filter_h - 1 - h][filter_w - 1 - w],
+        name="kernel_transform",
+    )
+    return data_pad, kernel_transform
+
+
+def declaration_conv3d_transpose_impl(data, kernel, strides, padding, out_dtype, output_padding):
+    """Implementation of conv3d transpose"""
+    data_pad, kernel_transform = conv3d_transpose_ncdhw_preprocess(
+        data, kernel, strides, padding, out_dtype, output_padding
+    )
+    batch, in_c, in_d, in_h, in_w = data_pad.shape
+    out_c, _, filter_d, filter_h, filter_w = kernel_transform.shape
+    stride_d, stride_h, stride_w = strides
+
+    # convolution stage
+    out_c = simplify(out_c)
+    out_d = simplify(in_d - filter_d + 1)
+    out_h = simplify(in_h - filter_h + 1)
+    out_w = simplify(in_w - filter_w + 1)
+    dc = te.reduce_axis((0, in_c), name="dc")
+    dd = te.reduce_axis((0, filter_d), name="dd")
+    dh = te.reduce_axis((0, filter_h), name="dh")
+    dw = te.reduce_axis((0, filter_w), name="dw")
+
+    Output = te.compute(
+        (batch, out_c, out_d, out_h, out_w),
+        lambda b, c, d, h, w: te.sum(
+            data_pad[b, dc, d + dd, h + dh, w + dw].astype(out_dtype)
+            * kernel_transform[c, dc, dd, dh, dw].astype(out_dtype),
+            axis=[dc, dd, dh, dw],
+        ),
+        tag="conv3d_transpose_ncdhw",
+    )
+
+    return Output
+
+
+def group_conv3d_transpose_ncdhw(data, kernel, strides, padding, out_dtype, output_padding, groups):
+    """Transposed group 3D convolution ncdhw forward operator.
+
+    Parameters
+    ----------
+    data : tvm.te.Tensor
+        5-D with shape [batch, in_channel, in_depth, in_height, in_width]
+
+    kernel : tvm.te.Tensor
+        5-D with shape [in_channel, num_filter, filter_depth, filter_height, filter_width]
+
+    strides : int or a list/tuple of three ints
+        The spatial stride along depth,height and width
+
+    padding : int or str
+        Padding size, or ['VALID', 'SAME']
+
+    out_dtype : str
+        The output data type. This is used for mixed precision.
+
+    output_padding : tuple of ints
+        Used to get the right output shape for gradients
+
+    groups : int
+        number of groups
+
+    Returns
+    -------
+    Output : tvm.te.Tensor
+        5-D with shape [batch, out_channel, out_depth, out_height, out_width]
+    """
+    if not isinstance(strides, tuple | list):
+        strides = (strides, strides, strides)
+
+    if groups == 1:
+        return conv3d_transpose_ncdhw(data, kernel, strides, padding, out_dtype, output_padding)
+
+    data_pad, kernel_transform = conv3d_transpose_ncdhw_preprocess(
+        data, kernel, strides, padding, out_dtype, output_padding
+    )
+    batch, in_c, in_d, in_h, in_w = data_pad.shape
+    out_c, _, filter_d, filter_h, filter_w = kernel_transform.shape
+    assert in_c % groups == 0, f"input channels {in_c} must divide group size {groups}"
+
+    # convolution stage
+    out_c = simplify(out_c * groups)
+    out_d = simplify(in_d - filter_d + 1)
+    out_h = simplify(in_h - filter_h + 1)
+    out_w = simplify(in_w - filter_w + 1)
+    dc = te.reduce_axis((0, in_c // groups), name="dc")
+    dd = te.reduce_axis((0, filter_d), name="dd")
+    dh = te.reduce_axis((0, filter_h), name="dh")
+    dw = te.reduce_axis((0, filter_w), name="dw")
+
+    # data: batch, in_channels, out_d, out_h, out_w
+    # weight: out_channels // G, in_channels, out_d, out_h, out_w
+    return te.compute(
+        (batch, out_c, out_d, out_h, out_w),
+        lambda b, c, d, h, w: te.sum(
+            data_pad[
+                b, c // (out_c // groups) * (in_c // groups) + dc, d + dd, h + dh, w + dw
+            ].astype(out_dtype)
+            * kernel_transform[
+                c % (out_c // groups),
+                c // (out_c // groups) * (in_c // groups) + dc,
+                dd,
+                dh,
+                dw,
+            ].astype(out_dtype),
+            axis=[dc, dd, dh, dw],
+        ),
+        tag="group_conv3d_transpose_ncdhw",
+    )

python/tvm/topi/nn/correlation.py

@@ -0,0 +1,124 @@
+# 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.
+# ruff: noqa: E731
+"""Correlation operators"""
+
+from tvm import te
+
+from ..utils import get_const_tuple
+from .pad import pad
+
+
+def correlation_nchw(
+    data1, data2, kernel_size, max_displacement, stride1, stride2, padding, is_multiply
+):
+    """Correlation operator in NCHW layout.
+
+    Parameters
+    ----------
+    data1 : tvm.te.Tensor
+        4-D with shape [batch, channel, height, width]
+
+    data2 : tvm.te.Tensor
+        4-D with shape [batch, channel, height, width]
+
+    kernel_size: int
+        Kernel size for correlation, must be an odd number
+
+    max_displacement: int
+        Max displacement of Correlation
+
+    stride1: int
+        Stride for data1
+
+    stride2: int
+        Stride for data2 within the neightborhood centered around data1
+
+    padding : int or a list/tuple of 2 or 4 ints
+        Padding size, or
+        [pad_height, pad_width] for 2 ints, or
+        [pad_top, pad_left, pad_bottom, pad_right] for 4 ints
+
+    is_multiply: bool
+        operation type is either multiplication or substraction
+
+    Returns
+    -------
+    Output : tvm.te.Tensor
+        4-D with shape [batch, out_channel, out_height, out_width]
+    """
+    # pylint: disable=unnecessary-lambda, invalid-name
+    data_shape = get_const_tuple(data1.shape)
+    assert get_const_tuple(data2.shape) == data_shape, "data1 and data2 should have the same shape"
+    assert kernel_size > 0 and kernel_size % 2, "kernel_size should be non-negative odd number"
+    if isinstance(padding, tuple | list):
+        if len(padding) == 2:
+            pad_before_h = pad_after_h = padding[0]
+            pad_before_w = pad_after_w = padding[1]
+        elif len(padding) == 4:
+            pad_before_h, pad_before_w, pad_after_h, pad_after_w = padding
+        else:
+            raise ValueError("invalid padding")
+    elif isinstance(padding, int):
+        pad_before_h = pad_after_h = pad_before_w = pad_after_w = padding
+    else:
+        raise ValueError("invalid padding")
+    pad_before = [0, 0, pad_before_h, pad_before_w]
+    pad_after = [0, 0, pad_after_h, pad_after_w]
+    padded_data1 = pad(data1, pad_before, pad_after)
+    padded_data2 = pad(data2, pad_before, pad_after)
+
+    batch, channel, height, width = data_shape
+
+    kernel_radius = (kernel_size - 1) // 2
+    border_size = max_displacement + kernel_radius
+    displacement_radius = max_displacement // stride2
+    displacement_size = 2 * displacement_radius + 1
+
+    padded_width = width + pad_before_w + pad_after_w
+    padded_height = height + pad_before_h + pad_after_h
+    out_channel = displacement_size * displacement_size
+    out_height = (padded_height - 2 * border_size + stride1 - 1) // stride1
+    out_width = (padded_width - 2 * border_size + stride1 - 1) // stride1
+
+    rc = te.reduce_axis((0, channel), name="rc")
+    ry = te.reduce_axis((0, kernel_size), name="ry")
+    rx = te.reduce_axis((0, kernel_size), name="rx")
+
+    if is_multiply:
+        corr_func = lambda x, y: x * y
+    else:
+        corr_func = lambda x, y: te.abs(x - y)
+
+    def _compute_correlation(n, q, i, j):
+        # location in data1
+        y1 = i * stride1 + max_displacement
+        x1 = j * stride1 + max_displacement
+        # location in data2
+        y2 = y1 + (te.indexdiv(q, displacement_size) - displacement_radius) * stride2
+        x2 = x1 + (te.indexmod(q, displacement_size) - displacement_radius) * stride2
+        return te.sum(
+            corr_func(padded_data1[n, rc, y1 + ry, x1 + rx], padded_data2[n, rc, y2 + ry, x2 + rx]),
+            axis=[rc, ry, rx],
+        )
+
+    correlation = te.compute(
+        (batch, out_channel, out_height, out_width),
+        lambda n, q, i, j: _compute_correlation(n, q, i, j),
+        tag="correlation_nchw",
+    )
+    return correlation / (kernel_size * kernel_size * channel)

python/tvm/topi/nn/deformable_conv2d.py

@@ -0,0 +1,241 @@
+# 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-arguments
+"""Deformable Conv2D operators"""
+
+import tvm
+from tvm import te
+
+from ..cpp.utils import bilinear_sample_nchw, bilinear_sample_nhwc
+from ..utils import get_const_tuple
+from .utils import get_pad_tuple
+
+
+def deformable_conv2d_nchw(
+    data, offset, kernel, strides, padding, dilation, deformable_groups, groups, out_dtype
+):
+    """Deformable conv2D operator in NCHW layout.
+
+    The deformable convolution operation is described in https://arxiv.org/abs/1703.06211
+
+    Parameters
+    ----------
+    data : tvm.te.Tensor
+        4-D with shape [batch, in_channel, in_height, in_width]
+
+    offset : tvm.te.Tensor
+        4-D with shape [batch, deformable_groups * filter_height * filter_width * 2,
+        out_height, out_width].
+
+    kernel : tvm.te.Tensor
+        4-D with shape [num_filter, in_channel, filter_height, filter_width]
+
+    strides : int or a list/tuple of two ints
+        stride size, or [stride_height, stride_width]
+
+    padding : int or a list/tuple of two ints
+        padding size, or [pad_height, pad_width]
+
+    dilation : int or a list/tuple of two ints
+        dilation size, or [dilation_height, dilation_width]
+
+    deformable_groups : int
+        number of deformable groups
+
+    groups : int
+        number of groups
+
+    Returns
+    -------
+    output : tvm.te.Tensor
+        4-D with shape [batch, out_channel, out_height, out_width]
+    """
+    if out_dtype is None:
+        out_dtype = data.dtype
+
+    if isinstance(strides, int):
+        stride_h = stride_w = strides
+    else:
+        stride_h, stride_w = strides
+
+    if isinstance(dilation, int):
+        dilation_h = dilation_w = dilation
+    else:
+        dilation_h, dilation_w = dilation
+
+    batch, in_channel, in_height, in_width = get_const_tuple(data.shape)
+    out_channel, channel, kernel_h, kernel_w = get_const_tuple(kernel.shape)
+    _, _, out_height, out_width = get_const_tuple(offset.shape)
+    assert in_channel % deformable_groups == 0, "Input cahnnels must divide deformable group size"
+    assert groups == 1, "deformable_conv2d_nchw does not support groups > 1"
+
+    ic_per_dgroup = channel // deformable_groups
+
+    dilated_kernel_h = (kernel_h - 1) * dilation_h + 1
+    dilated_kernel_w = (kernel_w - 1) * dilation_w + 1
+    pad_top, pad_left, _, _ = get_pad_tuple(padding, (dilated_kernel_h, dilated_kernel_w))
+    rc = te.reduce_axis((0, in_channel), name="rc")
+    ry = te.reduce_axis((0, kernel_h), name="ry")
+    rx = te.reduce_axis((0, kernel_w), name="rx")
+
+    zero = tvm.tirx.const(0.0, data.dtype)
+
+    def _bilinear(n, c, h, w):
+        outside = tvm.tirx.any(h < 0, w < 0, h >= in_height, w >= in_width)
+        val = bilinear_sample_nchw(data, (n, c, h, w), in_height - 1, in_width - 1)
+        return tvm.tirx.if_then_else(outside, zero, val)
+
+    data_deform = te.compute(
+        (batch, in_channel, kernel_h, kernel_w, out_height, out_width),
+        lambda n, c, kh, kw, y, x: _bilinear(
+            n,
+            c,
+            y * stride_h
+            - pad_top
+            + kh * dilation_h
+            + offset[
+                n, c // ic_per_dgroup * (kernel_w * kernel_h * 2) + (kh * kernel_w + kw) * 2, y, x
+            ],
+            x * stride_w
+            - pad_left
+            + kw * dilation_w
+            + offset[
+                n,
+                c // ic_per_dgroup * (kernel_w * kernel_h * 2) + (kh * kernel_w + kw) * 2 + 1,
+                y,
+                x,
+            ],
+        ),
+        tag="data_deform",
+    )
+    return te.compute(
+        (batch, out_channel, out_height, out_width),
+        lambda n, f, y, x: te.sum(
+            data_deform[n, rc, ry, rx, y, x].astype(out_dtype)
+            * kernel[f, rc, ry, rx].astype(out_dtype),
+            axis=[rc, ry, rx],
+        ),
+        tag="deformable_conv2d_nchw",
+    )
+
+
+def deformable_conv2d_nhwc(
+    data, offset, kernel, strides, padding, dilation, deformable_groups, groups, out_dtype
+):
+    """Deformable conv2D operator in NHWC layout.
+
+    The deformable convolution operation is described in https://arxiv.org/abs/1703.06211
+
+    Parameters
+    ----------
+    data : tvm.te.Tensor
+        4-D with shape [batch, in_height, in_width, in_channel]
+
+    offset : tvm.te.Tensor
+        4-D with shape [batch, out_height, out_width,
+                        deformable_groups * filter_height * filter_width * 2].
+
+    kernel : tvm.te.Tensor
+        4-D with shape [filter_height, filter_width, in_channel, num_filter]
+
+    strides : int or a list/tuple of two ints
+        stride size, or [stride_height, stride_width]
+
+    padding : int or a list/tuple of two ints
+        padding size, or [pad_height, pad_width]
+
+    dilation : int or a list/tuple of two ints
+        dilation size, or [dilation_height, dilation_width]
+
+    deformable_groups : int
+        number of deformable groups
+
+    groups : int
+        number of groups
+
+    Returns
+    -------
+    output : tvm.te.Tensor
+        4-D with shape [batch, out_height, out_width, out_channel]
+    """
+    if out_dtype is None:
+        out_dtype = data.dtype
+
+    if isinstance(strides, int):
+        stride_h = stride_w = strides
+    else:
+        stride_h, stride_w = strides
+
+    if isinstance(dilation, int):
+        dilation_h = dilation_w = dilation
+    else:
+        dilation_h, dilation_w = dilation
+
+    batch, in_height, in_width, in_channel = get_const_tuple(data.shape)
+    kernel_h, kernel_w, channel, out_channel = get_const_tuple(kernel.shape)
+    _, out_height, out_width, _ = get_const_tuple(offset.shape)
+    assert in_channel % deformable_groups == 0, "Input cahnnels must divide deformable group size"
+    assert groups == 1, "deformable_conv2d_nchw does not support groups > 1"
+
+    ic_per_dgroup = channel // deformable_groups
+
+    dilated_kernel_h = (kernel_h - 1) * dilation_h + 1
+    dilated_kernel_w = (kernel_w - 1) * dilation_w + 1
+    pad_top, pad_left, _, _ = get_pad_tuple(padding, (dilated_kernel_h, dilated_kernel_w))
+    rc = te.reduce_axis((0, in_channel), name="rc")
+    ry = te.reduce_axis((0, kernel_h), name="ry")
+    rx = te.reduce_axis((0, kernel_w), name="rx")
+
+    zero = tvm.tirx.const(0.0, data.dtype)
+
+    def _bilinear(n, h, w, c):
+        outside = tvm.tirx.any(h < 0, w < 0, h >= in_height, w >= in_width)
+        val = bilinear_sample_nhwc(data, (n, h, w, c), in_height - 1, in_width - 1)
+        return tvm.tirx.if_then_else(outside, zero, val)
+
+    data_deform = te.compute(
+        (batch, kernel_h, kernel_w, in_channel, out_height, out_width),
+        lambda n, kh, kw, c, y, x: _bilinear(
+            n,
+            y * stride_h
+            - pad_top
+            + kh * dilation_h
+            + offset[
+                n, y, x, c // ic_per_dgroup * (kernel_w * kernel_h * 2) + (kh * kernel_w + kw) * 2
+            ],
+            x * stride_w
+            - pad_left
+            + kw * dilation_w
+            + offset[
+                n,
+                y,
+                x,
+                c // ic_per_dgroup * (kernel_w * kernel_h * 2) + (kh * kernel_w + kw) * 2 + 1,
+            ],
+            c,
+        ),
+        tag="data_deform",
+    )
+    return te.compute(
+        (batch, out_height, out_width, out_channel),
+        lambda n, y, x, f: te.sum(
+            data_deform[n, ry, rx, rc, y, x].astype(out_dtype)
+            * kernel[ry, rx, rc, f].astype(out_dtype),
+            axis=[ry, rx, rc],
+        ),
+        tag="deformable_conv2d_nhwc",
+    )

python/tvm/topi/nn/dense.py

@@ -0,0 +1,261 @@
+# 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,unused-argument
+# ruff: noqa: E741, F821
+"""TVM operator fully connected compute."""
+
+import tvm
+from tvm import te
+
+from .. import add, tag
+
+
+def matmul(
+    tensor_a,
+    tensor_b,
+    bias=None,
+    out_dtype=None,
+    transpose_a=False,
+    transpose_b=False,
+    auto_scheduler_rewritten_layout="",
+    meta_schedule_original_shape=None,
+):
+    """The default implementation of matmul in topi.
+
+    Parameters
+    ----------
+    tensor_a : tvm.te.Tensor
+        2-D with shape [batch, in_dim]
+
+    tensor_b : tvm.te.Tensor
+        2-D with shape [out_dim, in_dim]
+
+    bias : Optional[tvm.te.Tensor]
+        1-D with shape [out_dim]
+
+    out_dtype : Optional[str]
+        The output type. This is used for mixed precision.
+
+    transpose_a : Optional[bool] = False
+        Whether the tensor_a is in transposed format.
+
+    transpose_b : Optional[bool] = False
+        Whether the tensor_b is in transposed format.
+
+    auto_scheduler_rewritten_layout: Optional[str] = ""
+        The layout after auto-scheduler's layout rewrite pass.
+
+    meta_schedule_original_shape: Optional[List[Expr]] = None
+        The original shape of the input tensor.
+
+    Returns
+    -------
+    output : tvm.te.Tensor
+        2-D with shape [batch, out_dim]
+    """
+    # TODO(yixin): support cases for 1-dim input
+    # TODO(yixin): adding support and further check for >2-dim input in autotvm template
+    assert len(tensor_a.shape) >= 2 and len(tensor_b.shape) >= 2, (
+        "1-dim matmul is not supported yet."
+    )
+
+    if bias is not None:
+        assert len(bias.shape) == 1
+    if out_dtype is None:
+        out_dtype = tensor_a.dtype
+    if transpose_a:
+        reduce_dim_a, in_dim = tensor_a.shape[-2:]
+    else:
+        in_dim, reduce_dim_a = tensor_a.shape[-2:]
+    batch_dims_a = tensor_a.shape[:-2]
+
+    if auto_scheduler_rewritten_layout:
+        # Infer shape for the rewritten layout
+        raise RuntimeError("LEGACY-FLOW triggered, to be removed")
+    if meta_schedule_original_shape:
+        raise RuntimeError("LEGACY-FLOW triggered, to be removed")
+
+    if transpose_b:
+        out_dim, reduce_dim_b = tensor_b.shape[-2:]
+    else:
+        reduce_dim_b, out_dim = tensor_b.shape[-2:]
+    batch_dims_b = tensor_b.shape[:-2]
+
+    if not isinstance(reduce_dim_a, tvm.tirx.Var) and not isinstance(reduce_dim_b, tvm.tirx.Var):
+        assert int(reduce_dim_a) == int(reduce_dim_b), (
+            f"Reduction dimensions of dense do not match. {reduce_dim_a} vs {reduce_dim_b}."
+        )
+
+    result_ndim = max(len(batch_dims_a), len(batch_dims_b))
+    batch_dims_a = [1] * (result_ndim - len(batch_dims_a)) + batch_dims_a
+    batch_dims_b = [1] * (result_ndim - len(batch_dims_b)) + batch_dims_b
+
+    for idx, (l, r) in enumerate(zip(batch_dims_a, batch_dims_b)):
+        if (
+            not isinstance(l, tvm.tirx.Var)
+            and not isinstance(r, tvm.tirx.Var)
+            and int(l) != 1
+            and int(r) != 1
+        ):
+            assert int(l) == int(r), (
+                "Batch dimensions of dense do not match: "
+                f"{tensor_a.shape[:-2]} vs {tensor_b.shape[:-2]}."
+            )
+        if not isinstance(l, tvm.tirx.Var) and int(l) == 1:
+            batch_dims_a[idx] = batch_dims_b[idx]
+
+    k = te.reduce_axis((0, reduce_dim_a), name="k")
+
+    def compute(*indices):
+        batch_indices_a = indices[-len(tensor_a.shape) : -2]
+        batch_indices_a = [
+            i if isinstance(dim, tvm.tirx.Var) or int(dim) != 1 else 0
+            for i, dim in zip(batch_indices_a, tensor_a.shape[:-2])
+        ]
+        batch_indices_b = indices[-len(tensor_b.shape) : -2]
+        batch_indices_b = [
+            i if isinstance(dim, tvm.tirx.Var) or int(dim) != 1 else 0
+            for i, dim in zip(batch_indices_b, tensor_b.shape[:-2])
+        ]
+        i, j = indices[-2:]
+        a_indices = (*batch_indices_a, k, i) if transpose_a else (*batch_indices_a, i, k)
+        b_indices = (*batch_indices_b, j, k) if transpose_b else (*batch_indices_b, k, j)
+        return te.sum(
+            tensor_a[a_indices].astype(out_dtype) * tensor_b[b_indices].astype(out_dtype), axis=k
+        )
+
+    compute_name = {
+        (True, True): "T_matmul_TT",
+        (True, False): "T_matmul_TN",
+        (False, True): "T_matmul_NT",
+        (False, False): "T_matmul_NN",
+    }[(transpose_a, transpose_b)]
+
+    # TODO(jcf94): Remove `dense` when `matmul` is finally ready
+    compute_tag = "dense" if (transpose_a, transpose_b) == (False, True) else "matmul"
+
+    mat = te.compute(
+        (*batch_dims_a, in_dim, out_dim),
+        compute,
+        name=compute_name,
+        tag=compute_tag,
+        attrs={"layout_free_placeholders": [tensor_b]},
+    )
+
+    if bias is not None:
+        mat = add(mat, bias.astype(out_dtype))
+
+    if auto_scheduler_rewritten_layout:
+        raise RuntimeError("LEGACY-FLOW triggered, to be removed")
+
+    return mat
+
+
+def dense(
+    data,
+    weight,
+    bias=None,
+    out_dtype=None,
+    auto_scheduler_rewritten_layout="",
+    meta_schedule_original_shape=None,
+):
+    """The default implementation of dense in topi.
+    This is an alias of matmul_nt operator for data tensor in non-transposed format and weight
+    tensor in transposed format.
+
+    Parameters
+    ----------
+    data : tvm.te.Tensor
+        2-D with shape [batch, in_dim]
+
+    weight : tvm.te.Tensor
+        2-D with shape [out_dim, in_dim]
+
+    bias : Optional[tvm.te.Tensor]
+        1-D with shape [out_dim]
+
+    out_dtype : Optional[str]
+        The output type. This is used for mixed precision.
+
+    auto_scheduler_rewritten_layout: str = ""
+        The layout after auto-scheduler's layout rewrite pass.
+
+    meta_schedule_original_shape: Optional[List[Expr]] = None
+        The original shape of the input tensor.
+
+    Returns
+    -------
+    output : tvm.te.Tensor
+        2-D with shape [batch, out_dim]
+    """
+
+    return matmul(
+        data,
+        weight,
+        bias,
+        out_dtype,
+        False,
+        True,
+        auto_scheduler_rewritten_layout,
+        meta_schedule_original_shape,
+    )
+
+
+def dense_pack(data, weight, bias=None, out_dtype=None):
+    """The default implementation of dense_pack in topi.
+
+    Parameters
+    ----------
+    data : tvm.te.Tensor
+        2-D with shape [batch, in_dim]
+
+    weight : tvm.te.Tensor
+        2-D with shape [out_dim, in_dim]
+
+    bias : Optional[tvm.te.Tensor]
+        1-D with shape [out_dim]
+
+    out_dtype : Optional[str]
+        The output type. This is used for mixed precision.
+
+    Returns
+    -------
+    output : tvm.te.Tensor
+        2-D with shape [batch, out_dim]
+    """
+    if out_dtype is None:
+        out_dtype = data.dtype
+    M, K = get_const_tuple(data.shape)  # batch, in_dim
+    N, _, packw_bn = get_const_tuple(weight.shape)  # out_dim
+    N = N * packw_bn
+
+    idxdiv = tvm.tirx.indexdiv
+    idxmod = tvm.tirx.indexmod
+    k = te.reduce_axis((0, K), name="k")
+    C = te.compute(
+        (M, N),
+        lambda y, x: te.sum(
+            data[y, k].astype(out_dtype)
+            * weight[idxdiv(x, packw_bn), k, idxmod(x, packw_bn)].astype(out_dtype),
+            axis=k,
+        ),
+        name="T_dense_pack",
+        tag="dense_pack",
+    )
+    if bias is not None:
+        C = te.compute((M, N), lambda i, j: C[i, j] + bias[j].astype(out_dtype), tag=tag.BROADCAST)
+    return C

python/tvm/topi/nn/depth_to_space.py

@@ -0,0 +1,88 @@
+# 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
+"""TVM operator depth_to_space compute."""
+
+import tvm
+from tvm import te
+
+from .. import tag
+
+
+def depth_to_space(data, block_size, layout="NCHW", mode="DCR"):
+    """Perform depth to space transformation on the data
+
+    Parameters
+    ----------
+    data : tvm.te.Tensor
+        4-D tensor in either NCHW or NHWC layout.
+
+    block_size : int
+        Size of blocks to compose from channel dimension.
+
+    layout : string
+        Either NCHW or NHWC, indicating data layout.
+
+    mode : string
+        Either DCR or CDR, indicates how channels should be accessed.
+        In DCR, channels are interwoven in the Tensorflow style while
+        in CDR channels are accessed sequentially as in Pytorch.
+
+    Returns
+    -------
+    output : tvm.te.Tensor
+        Output of shape [N, C / block_size**2, H * block_size, W * block_size]
+    """
+    if layout == "NCHW":
+        in_n, in_c, in_h, in_w = data.shape
+        channel_factor = tvm.tirx.truncdiv(in_c, (block_size * block_size))
+        output_shape = [in_n, channel_factor, in_h * block_size, in_w * block_size]
+    elif layout == "NHWC":
+        in_n, in_h, in_w, in_c = data.shape
+        channel_factor = tvm.tirx.truncdiv(in_c, (block_size * block_size))
+        output_shape = [in_n, in_h * block_size, in_w * block_size, channel_factor]
+    else:
+        raise ValueError("Only NCHW and NHWC layouts are currently supported.")
+
+    def _get_indices(*indices):
+        if layout == "NCHW":
+            n, c, y, x = indices
+        elif layout == "NHWC":
+            n, y, x, c = indices
+        return n, c, y, x
+
+    def _get_pixel(n, c, y, x):
+        block_x = tvm.tirx.truncdiv(x, block_size)
+        block_y = tvm.tirx.truncdiv(y, block_size)
+        idx_x = tvm.tirx.truncmod(x, block_size)
+        idx_y = tvm.tirx.truncmod(y, block_size)
+        if mode == "DCR":
+            channel_idx = channel_factor * ((block_size * idx_y) + idx_x) + c
+        else:
+            channel_idx = (c * block_size * block_size) + ((block_size * idx_y) + idx_x)
+
+        if layout == "NCHW":
+            output = data(n, channel_idx, block_y, block_x)
+        else:
+            output = data(n, block_y, block_x, channel_idx)
+        return output
+
+    def _compute(*indices):
+        n, c, y, x = _get_indices(*indices)
+        return _get_pixel(n, c, y, x)
+
+    return te.compute(output_shape, _compute, name="depth_to_space", tag=tag.INJECTIVE)

python/tvm/topi/nn/depthwise_conv2d.py

@@ -0,0 +1,462 @@
+# 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, unused-variable, too-many-locals, unused-argument
+# ruff: noqa: F841
+"""Depthwise convolution operators"""
+
+from collections import namedtuple
+
+import numpy as np
+
+import tvm
+from tvm import te
+
+from ..utils import get_const_tuple, simplify
+from .dilate import dilate
+from .pad import pad
+from .utils import get_pad_tuple
+
+# workload description of depthwise-conv2d
+Workload = namedtuple(
+    "Workload",
+    [
+        "in_dtype",
+        "out_dtype",
+        "height",
+        "width",
+        "in_filter",
+        "out_filter",
+        "kernel_h",
+        "kernel_w",
+        "padt",
+        "padl",
+        "padb",
+        "padr",
+        "dilation_h",
+        "dilation_w",
+        "stride_h",
+        "stride_w",
+    ],
+)
+
+
+def _get_workload(data, kernel, stride, padding, dilation, out_dtype, data_layout="NCHW"):
+    """Get the workload structure for a depthwise conv2d.
+
+    Input data and filter should use NCHW layout.
+    """
+    if data_layout == "NCHW":
+        _, in_channel, height, width = get_const_tuple(data.shape)
+        filter_channel, channel_multiplier, kh, kw = get_const_tuple(kernel.shape)
+    elif data_layout == "NHWC":
+        _, height, width, in_channel = get_const_tuple(data.shape)
+        kh, kw, filter_channel, channel_multiplier = get_const_tuple(kernel.shape)
+    elif data_layout == "NCHWc":
+        _, in_channel_chunk, height, width, in_channel_block = get_const_tuple(data.shape)
+        in_channel = in_channel_chunk * in_channel_block
+        (filter_channel_chunk, cm_chunk, kh, kw, cm_block, filter_channel_block) = get_const_tuple(
+            kernel.shape
+        )
+        filter_channel = filter_channel_chunk * filter_channel_block
+        channel_multiplier = cm_chunk * cm_block
+
+        assert in_channel_block == filter_channel_block, (
+            f"Incorrect dimensions, data has block size {in_channel_block}, but filter has "
+            f"block size {filter_channel_block}"
+        )
+
+    else:
+        raise ValueError(f"Data layout {data_layout} not supported")
+
+    assert in_channel == filter_channel, (
+        f"Incorrect dimensions, data has {in_channel} channels but filter expects "
+        f"{filter_channel} channels"
+    )
+
+    out_channel = filter_channel * channel_multiplier
+    dilation_h, dilation_w = (
+        dilation if isinstance(dilation, tuple | list) else (dilation, dilation)
+    )
+    if isinstance(stride, tuple | list):
+        HSTR, WSTR = stride
+    else:
+        HSTR, WSTR = stride, stride
+    assert (data.dtype == kernel.dtype) or (data.dtype == "uint8" and kernel.dtype == "int8"), (
+        f"Do not support inputs with different data types now. {data.dtype} vs. {kernel.dtype}"
+    )
+    dilated_kernel_h = (kh - 1) * dilation_h + 1
+    dilated_kernel_w = (kw - 1) * dilation_w + 1
+    pt, pl, pb, pr = get_pad_tuple(padding, (dilated_kernel_h, dilated_kernel_w))
+    return Workload(
+        data.dtype,
+        out_dtype,
+        height,
+        width,
+        in_channel,
+        out_channel,
+        kh,
+        kw,
+        pt,
+        pl,
+        pb,
+        pr,
+        dilation_h,
+        dilation_w,
+        HSTR,
+        WSTR,
+    )
+
+
+def depthwise_conv2d_nchw(Input, Filter, stride, padding, dilation, out_dtype=None):
+    """Depthwise convolution nchw forward operator.
+
+    Parameters
+    ----------
+    Input : tvm.te.Tensor
+        4-D with shape [batch, in_channel, in_height, in_width]
+
+    Filter : tvm.te.Tensor
+        4-D with shape [in_channel, channel_multiplier, filter_height, filter_width]
+
+    stride : int or a list/tuple of two ints
+        The spatial stride, or (stride_height, stride_width).
+
+    padding : int or str
+        Padding size, or ['VALID', 'SAME']
+
+    dilation: int or a list/tuple of two ints
+        dilation size, or [dilation_height, dilation_width]
+
+    out_dtype: str, optional
+        Output data type
+
+    Returns
+    -------
+    Output : tvm.te.Tensor
+        4-D with shape [batch, out_channel, out_height, out_width]
+    """
+    out_dtype = Input.dtype if out_dtype is None else out_dtype
+
+    if isinstance(stride, int):
+        stride_h = stride_w = stride
+    else:
+        stride_h, stride_w = stride
+
+    if isinstance(dilation, int):
+        dilation_h = dilation_w = dilation
+    else:
+        dilation_h, dilation_w = dilation
+
+    batch, in_channel, in_height, in_width = Input.shape
+    # shape of dilated kernel
+    filter_channel, channel_multiplier, filter_height, filter_width = Filter.shape
+
+    dilated_kernel_h = (filter_height - 1) * dilation_h + 1
+    dilated_kernel_w = (filter_width - 1) * dilation_w + 1
+    pad_top, pad_left, pad_down, pad_right = get_pad_tuple(
+        padding, (dilated_kernel_h, dilated_kernel_w)
+    )
+    out_channel = simplify(in_channel * channel_multiplier)
+    out_height = simplify((in_height - dilated_kernel_h + pad_top + pad_down) // stride_h + 1)
+    out_width = simplify((in_width - dilated_kernel_w + pad_left + pad_right) // stride_w + 1)
+
+    # padding stage
+    pad_before = [0, 0, pad_top, pad_left]
+    pad_after = [0, 0, pad_down, pad_right]
+    PaddedInput = pad(Input, pad_before, pad_after, name="PaddedInput")
+    # depthconv stage
+    idxdiv = tvm.tirx.indexdiv
+    idxmod = tvm.tirx.indexmod
+    di = te.reduce_axis((0, filter_height), name="di")
+    dj = te.reduce_axis((0, filter_width), name="dj")
+    Output = te.compute(
+        (batch, out_channel, out_height, out_width),
+        lambda b, c, i, j: te.sum(
+            (
+                PaddedInput[
+                    b,
+                    idxdiv(c, channel_multiplier),
+                    i * stride_h + di * dilation_h,
+                    j * stride_w + dj * dilation_w,
+                ].astype(out_dtype)
+                * Filter[
+                    idxdiv(c, channel_multiplier), idxmod(c, channel_multiplier), di, dj
+                ].astype(out_dtype)
+            ),
+            axis=[di, dj],
+        ),
+        name="DepthwiseConv2d",
+        tag="depthwise_conv2d_nchw",
+    )
+    return Output
+
+
+def depthwise_conv2d_nhwc(
+    Input, Filter, stride, padding, dilation, kernel_layout="HWOI", out_dtype=None
+):
+    """Depthwise convolution nhwc forward operator.
+
+    Parameters
+    ----------
+    Input : tvm.te.Tensor
+        4-D with shape [batch, in_height, in_width, in_channel]
+
+    Filter : tvm.te.Tensor
+        4-D with shape [filter_height, filter_width, in_channel, channel_multiplier]
+
+    stride : tuple of two ints
+        The spatial stride along height and width
+
+    padding : int or str
+        Padding size, or ['VALID', 'SAME']
+
+    dilation: int or a list/tuple of two ints
+        dilation size, or [dilation_height, dilation_width]
+
+    out_dtype: str, optional
+        Output data type
+
+    Returns
+    -------
+    Output : tvm.te.Tensor
+        4-D with shape [batch, out_height, out_width, out_channel]
+    """
+    out_dtype = Input.dtype if out_dtype is None else out_dtype
+
+    if isinstance(stride, int):
+        stride_h = stride_w = stride
+    else:
+        stride_h, stride_w = stride
+
+    if isinstance(dilation, int):
+        dilation_h = dilation_w = dilation
+    else:
+        dilation_h, dilation_w = dilation
+
+    batch, in_height, in_width, in_channel = Input.shape
+
+    # shape of dilated kernel
+    if kernel_layout == "HWIO":
+        filter_height, filter_width, channel_multiplier, filter_channel = Filter.shape
+        kernel_permutation = [0, 1, 3, 2]
+    else:
+        filter_height, filter_width, filter_channel, channel_multiplier = Filter.shape
+        kernel_permutation = [0, 1, 2, 3]
+
+    dilated_kernel_h = (filter_height - 1) * dilation_h + 1
+    dilated_kernel_w = (filter_width - 1) * dilation_w + 1
+    pad_top, pad_left, pad_down, pad_right = get_pad_tuple(
+        padding, (dilated_kernel_h, dilated_kernel_w)
+    )
+    out_channel = simplify(in_channel * channel_multiplier)
+    out_height = simplify((in_height - dilated_kernel_h + pad_top + pad_down) // stride_h + 1)
+    out_width = simplify((in_width - dilated_kernel_w + pad_left + pad_right) // stride_w + 1)
+
+    # padding stage
+    pad_before = [0, pad_top, pad_left, 0]
+    pad_after = [0, pad_down, pad_right, 0]
+    PaddedInput = pad(Input, pad_before, pad_after, name="PaddedInput")
+    # depthconv stage
+    idxdiv = tvm.tirx.indexdiv
+    idxmod = tvm.tirx.indexmod
+
+    di = te.reduce_axis((0, filter_height), name="di")
+    dj = te.reduce_axis((0, filter_width), name="dj")
+    Output = te.compute(
+        (batch, out_height, out_width, out_channel),
+        lambda b, i, j, c: te.sum(
+            (
+                PaddedInput[
+                    b,
+                    i * stride_h + di * dilation_h,
+                    j * stride_w + dj * dilation_w,
+                    idxdiv(c, channel_multiplier),
+                ].astype(out_dtype)
+                * Filter[
+                    tuple(
+                        np.array(
+                            [di, dj, idxdiv(c, channel_multiplier), idxmod(c, channel_multiplier)]
+                        )[kernel_permutation]
+                    )
+                ].astype(out_dtype)
+            ),
+            axis=[di, dj],
+        ),
+        name="DepthwiseConv2d",
+        tag="depthwise_conv2d_nhwc",
+    )
+    return Output
+
+
+def depthwise_conv2d_backward_input_nhwc(Filter, Out_grad, oshape, ishape, stride, padding):
+    """Depthwise convolution nhwc backward wrt input operator.
+
+    Parameters
+    ----------
+    Filter : tvm.te.Tensor
+        4-D with shape [filter_height, filter_width, in_channel, channel_multiplier]
+
+    Out_grad : tvm.te.Tensor
+        4-D with shape [batch, out_height, out_width, out_channel]
+
+    stride : tuple of two ints
+        The spatial stride along height and width
+
+    padding : int or str
+        Padding size, or ['VALID', 'SAME']
+
+    Returns
+    -------
+    Output : tvm.te.Tensor
+        4-D with shape [batch, in_height, in_width, in_channel]
+    """
+    batch, in_h, in_w, in_c = ishape
+    _, out_h, out_w, out_c = oshape
+    filter_h, filter_w, _, channel_multiplier = Filter.shape
+    if isinstance(stride, int):
+        stride_h = stride_w = stride
+    else:
+        stride_h, stride_w = stride
+
+    dilated_out_grad = dilate(Out_grad, [1, stride_h, stride_w, 1], name="dilated_out_grad")
+
+    # padding params in forward propagation
+    fpad_top, fpad_left, fpad_bottom, fpad_right = get_pad_tuple(padding, (filter_h, filter_w))
+    # padding params in backward propagation
+    bpad_top = filter_h - 1 - fpad_top
+    bpad_bottom = (filter_h - 1 - fpad_bottom) + (stride_h - 1)
+    bpad_left = filter_w - 1 - fpad_left
+    bpad_right = (filter_w - 1 - fpad_right) + (stride_w - 1)
+
+    padded_out_grad = pad(
+        dilated_out_grad,
+        [0, bpad_top, bpad_left, 0],
+        [0, bpad_bottom, bpad_right, 0],
+        name="padded_out_grad",
+    )
+
+    dh = te.reduce_axis((0, filter_h), name="dh")
+    dw = te.reduce_axis((0, filter_w), name="dw")
+    dc = te.reduce_axis((0, channel_multiplier), name="dc")
+
+    In_grad = te.compute(
+        (batch, in_h, in_w, in_c),
+        lambda b, h, w, c: te.sum(
+            padded_out_grad[b, h + dh, w + dw, c * channel_multiplier + dc]
+            * Filter[filter_h - 1 - dh, filter_w - 1 - dw, c, dc],
+            axis=[dh, dw, dc],
+        ),
+        tag="depthwise_conv2d_backward_input_nhwc",
+    )
+
+    return In_grad
+
+
+def depthwise_conv2d_backward_weight_nhwc(Input, Out_grad, oshape, fshape, stride, padding):
+    """Depthwise convolution nhwc backward wrt weight operator.
+
+    Parameters
+    ----------
+    Input : tvm.te.Tensor
+        4-D with shape [batch, in_height, in_width, in_channel]
+
+    Out_grad : tvm.te.Tensor
+        4-D with shape [batch, out_height, out_width, out_channel]
+
+    stride : tuple of two ints
+        The spatial stride along height and width
+
+    padding : int or str
+        Padding size, or ['VALID', 'SAME']
+
+    Returns
+    -------
+    Output : tvm.te.Tensor
+        4-D with shape [filter_height, filter_width, in_channel, channel_multiplier]
+    """
+    batch, out_h, out_w, out_c = oshape
+    filter_h, filter_w, _, channel_multiplier = fshape
+    in_c = Input.shape[3].value
+    if isinstance(stride, int):
+        stride_h = stride_w = stride
+    else:
+        stride_h, stride_w = stride
+
+    pad_top, pad_left, pad_bottom, pad_right = get_pad_tuple(padding, (filter_h, filter_w))
+
+    padded_in = pad(
+        Input, [0, pad_top, pad_left, 0], [0, pad_bottom, pad_right, 0], name="padded_in"
+    )
+
+    dh = te.reduce_axis((0, Out_grad.shape[1].value), name="dh")
+    dw = te.reduce_axis((0, Out_grad.shape[2].value), name="dw")
+    db = te.reduce_axis((0, batch), name="db")
+    idxdiv = tvm.tirx.indexdiv
+    idxmod = tvm.tirx.indexmod
+
+    Weight_grad = te.compute(
+        (filter_h, filter_w, in_c, channel_multiplier),
+        lambda fh, fw, c, m: te.sum(
+            Out_grad[db, dh, dw, c * channel_multiplier + idxmod(m, channel_multiplier)]
+            * padded_in[db, fh + dh * stride_h, fw + dw * stride_w, c],
+            axis=[db, dh, dw],
+        ),
+        tag="depthwise_conv2d_backward_weight_nhwc",
+    )
+
+    return Weight_grad
+
+
+def depthwise_conv2d_NCHWc(
+    Input, Filter, stride, padding, dilation, layout, out_layout, out_dtype=None
+):
+    """Depthwise convolution NCHW[x]c forward operator.
+
+    Parameters
+    ----------
+    Input : tvm.te.Tensor
+        5-D with shape [batch, in_channel_chunk, in_height, in_width, in_channel_block]
+
+    Filter : tvm.te.Tensor
+        6-D with shape [out_channel_chunk, 1, filter_height, filter_width, 1, out_channel_block]
+        In NCHWc depthwise convolution,
+        we group kernel's in_channel and channel_multiplier together then do the tiling.
+
+    stride : tuple of two ints
+        The spatial stride along height and width
+
+    padding : int or str
+        Padding size, or ['VALID', 'SAME']
+
+    dilation: int or a list/tuple of two ints
+         dilation size, or [dilation_height, dilation_width]
+
+    layout : str
+        Input data layout
+
+    out_layout : str
+        Output data layout
+
+    out_dtype: str, optional
+        Output data type
+
+    Returns
+    -------
+    Output : tvm.te.Tensor
+        5-D with shape [batch, out_channel_chunk, out_height, out_width, out_channel_block]
+    """
+    raise ValueError("missing register for topi.nn.depthwise_conv2d_NCHWc")

python/tvm/topi/nn/dilate.py

@@ -0,0 +1,73 @@
+# 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
+"""Dilation operators"""
+
+import tvm
+from tvm import te
+
+from .. import tag, utils
+
+
+@te.tag_scope(tag=tag.INJECTIVE + ",dilate")
+def dilate(data, strides, dilation_value=0.0, name="DilatedInput"):
+    """Dilate data with given dilation value (0 by default).
+
+    Parameters
+    ----------
+    data : tvm.te.Tensor
+        n-D, can be any layout.
+
+    strides : list / tuple of n ints
+        Dilation stride on each dimension, 1 means no dilation.
+
+    dilation_value : int/float, optional
+        Value used to dilate the input.
+
+    name : str, optional
+        The name prefix operators generated
+
+    Returns
+    -------
+    Output : tvm.te.Tensor
+        n-D, the same layout as data.
+    """
+    n = len(data.shape)
+    if len(strides) != n:
+        raise ValueError(f"data dimension and strides size dismatch : {n} vs {len(strides)}")
+    ana = tvm.arith.Analyzer()
+    out_shape = tuple(ana.simplify((data.shape[i] - 1) * strides[i] + 1) for i in range(n))
+
+    def _dilate(*indices):
+        not_zero = []
+        index_tuple = []
+        idxdiv = tvm.tirx.indexdiv
+        idxmod = tvm.tirx.indexmod
+        for i in range(n):
+            if not utils.equal_const_int(strides[i], 1):
+                index_tuple.append(idxdiv(indices[i], strides[i]))
+                not_zero.append(idxmod(indices[i], strides[i]).equal(0))
+            else:
+                index_tuple.append(indices[i])
+        if not_zero:
+            not_zero = tvm.tirx.all(*not_zero)
+            return tvm.tirx.if_then_else(
+                not_zero, data(*index_tuple), tvm.tirx.const(dilation_value, data.dtype)
+            )
+        return data(*index_tuple)
+
+    return te.compute(out_shape, _dilate, name=name)

python/tvm/topi/nn/elemwise.py

@@ -0,0 +1,143 @@
+# 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.
+"""Elementwise operators"""
+
+import tvm
+from tvm import te
+
+from .. import tag
+from ..utils import get_const_int
+
+
+@tvm.te.tag_scope(tag=tag.ELEMWISE)
+def relu(x):
+    """Take relu of input x.
+
+    Parameters
+    ----------
+    x : tvm.te.Tensor
+        Input argument.
+
+    Returns
+    -------
+    y : tvm.te.Tensor
+        The result.
+    """
+    return te.compute(x.shape, lambda *i: tvm.te.max(x(*i), tvm.tirx.const(0, x.dtype)))
+
+
+@tvm.te.tag_scope(tag=tag.ELEMWISE)
+def leaky_relu(x, alpha):
+    """Take leaky relu of input x.
+
+    Parameters
+    ----------
+    x : tvm.te.Tensor
+        Input argument.
+
+    alpha : float
+        The slope for the small gradient when x < 0
+
+    Returns
+    -------
+    y : tvm.te.Tensor
+        The result.
+    """
+
+    def _compute(*indices):
+        value = x(*indices)
+        calpha = tvm.tirx.const(alpha, value.ty)
+        return tvm.tirx.Select(value > 0, value, value * calpha)
+
+    return te.compute(x.shape, _compute)
+
+
+@tvm.te.tag_scope(tag=tag.ELEMWISE)
+def softplus(x, beta=1.0, threshold=20.0):
+    """Compute Softplus activation for input x with numerical stability.
+
+    Parameters
+    ----------
+    x : tvm.te.Tensor
+        Input tensor.
+
+    beta : float, optional
+        The scaling factor β in the Softplus formula (default is 1.0).
+
+    threshold : float, optional
+        The threshold value for numerical stability (default is 20.0).
+
+    Returns
+    -------
+    y : tvm.te.Tensor
+        The result.
+    """
+
+    def _compute(*indices):
+        value = x(*indices)
+        b = tvm.tirx.const(beta, value.ty)
+        t = tvm.tirx.const(threshold, value.ty)
+
+        return tvm.tirx.Select(
+            b * value > t, value, (1 / b) * tvm.tirx.log(1 + tvm.tirx.exp(b * value))
+        )
+
+    return te.compute(x.shape, _compute)
+
+
+@tvm.te.tag_scope(tag=tag.BROADCAST)
+def prelu(x, slope, axis=1):
+    """PReLU.
+    It accepts two arguments: an input ``x`` and a weight array ``W``
+    and computes the output as :math:`PReLU(x) y = x > 0 ? x : W * x`,
+    where :math:`*` is an elementwise multiplication for each sample in the
+    batch.
+
+    Parameters
+    ----------
+    x : tvm.te.Tensor
+        Input argument.
+
+    slope : tvm.te.Tensor
+        Channelised slope tensor for prelu
+
+    axis : int
+        The axis where the channel data needs to be applied
+
+    Returns
+    -------
+    y : tvm.te.Tensor
+        The result.
+
+    Links
+    -----
+    [http://arxiv.org/pdf/1502.01852v1.pdf]
+    """
+
+    assert len(slope.shape) == 1
+    assert axis < len(x.shape)
+    if slope.shape[0] == 1:
+        slope = te.compute(
+            (get_const_int(x.shape[axis]),), lambda c: slope[0], name="slope_broadcasted"
+        )
+    assert get_const_int(slope.shape[0]) == get_const_int(x.shape[axis])
+
+    def _compute_channelwise(*indices):
+        xval = x(*indices)
+        return tvm.tirx.Select(xval > 0, xval, xval * slope(indices[axis]))
+
+    return te.compute(x.shape, _compute_channelwise)

python/tvm/topi/nn/fifo_buffer.py

@@ -0,0 +1,188 @@
+# 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.
+# ruff: noqa: E741
+
+"""FIFO buffer op"""
+
+import tvm
+from tvm import te
+
+from .. import tag
+from ..transform import concatenate, strided_slice
+
+
+@tvm.te.tag_scope(tag=tag.INJECTIVE + ",fifo_buffer")
+def fifo_buffer(data, buffer, axis):
+    """
+    FIFO buffer to enable computation reuse in CNNs with sliding indow input
+
+    Compute equivalent of
+
+    .. code-block:: python
+
+        concat(buffer, data, axis=axis)
+        .slice_axis(axis=axis,
+                    begin=data.shape[axis],
+                    end=data.shape[axis]+buffer.shape[axis])
+
+    Useful for
+
+    * Encoding explicit re-use of computation in convolution ops operated on a sliding window input
+    * Implementing a FIFO queue to cache intermediate results, e.g. as in Fast WaveNet.
+
+    Parameters
+    ----------
+    data : tvm.te.Tensor
+        The input data
+    buffer : tvm.te.Tensor
+        Previous value of the FIFO buffer
+    axis : int
+        Specify which axis should be used for buffering
+
+    Returns
+    -------
+    result : tvm.te.Tensor
+        Updated value for the buffer
+    """
+    assert len(data.shape) == len(buffer.shape), (
+        f"buffer and data must have same number of dimensions, "
+        f"buffer.shape = {buffer.shape}, data.shape = {data.shape}"
+    )
+    assert len(buffer.shape) >= 1, "Zero-dimension tensor not supported"
+    assert 0 <= axis < len(buffer.shape), "buffer axis out of range"
+    for i in range(len(data.shape)):
+        if i == axis:
+            assert int(str(data.shape[i])) <= int(str(buffer.shape[i]))
+        else:
+            assert int(str(data.shape[i])) == int(str(buffer.shape[i]))
+
+    buflen = buffer.shape[axis]
+    data_size = data.shape[axis]
+
+    # Explicitly write out formula up to 4D, and then use concat+slice combo for 5D and higher
+    if len(buffer.shape) == 1:
+        return te.compute(
+            buffer.shape,
+            lambda i: tvm.tirx.if_then_else(
+                i < buflen - data_size, buffer[i + data_size], data[i - buflen + data_size]
+            ),
+            name="new_buffer",
+        )
+    if len(buffer.shape) == 2:
+        if axis == 0:
+            return te.compute(
+                buffer.shape,
+                lambda i, j: tvm.tirx.if_then_else(
+                    i < buflen - data_size,
+                    buffer[i + data_size, j],
+                    data[i - buflen + data_size, j],
+                ),
+                name="new_buffer",
+            )
+        if axis == 1:
+            return te.compute(
+                buffer.shape,
+                lambda i, j: tvm.tirx.if_then_else(
+                    j < buflen - data_size,
+                    buffer[i, j + data_size],
+                    data[i, j - buflen + data_size],
+                ),
+                name="new_buffer",
+            )
+        assert False, f"Invalid value for axis; it should be at most {len(buffer.shape)}"
+    elif len(buffer.shape) == 3:
+        if axis == 0:
+            return te.compute(
+                buffer.shape,
+                lambda i, j, k: tvm.tirx.if_then_else(
+                    i < buflen - data_size,
+                    buffer[i + data_size, j, k],
+                    data[i - buflen + data_size, j, k],
+                ),
+                name="new_buffer",
+            )
+        if axis == 1:
+            return te.compute(
+                buffer.shape,
+                lambda i, j, k: tvm.tirx.if_then_else(
+                    j < buflen - data_size,
+                    buffer[i, j + data_size, k],
+                    data[i, j - buflen + data_size, k],
+                ),
+                name="new_buffer",
+            )
+        if axis == 2:
+            return te.compute(
+                buffer.shape,
+                lambda i, j, k: tvm.tirx.if_then_else(
+                    k < buflen - data_size,
+                    buffer[i, j, k + data_size],
+                    data[i, j, k - buflen + data_size],
+                ),
+                name="new_buffer",
+            )
+        assert False, f"Invalid value for axis; it should be at most {len(buffer.shape)}"
+    elif len(buffer.shape) == 4:
+        if axis == 0:
+            return te.compute(
+                buffer.shape,
+                lambda i, j, k, l: tvm.tirx.if_then_else(
+                    i < buflen - data_size,
+                    buffer[i + data_size, j, k, l],
+                    data[i - buflen + data_size, j, k, l],
+                ),
+                name="new_buffer",
+            )
+        if axis == 1:
+            return te.compute(
+                buffer.shape,
+                lambda i, j, k, l: tvm.tirx.if_then_else(
+                    j < buflen - data_size,
+                    buffer[i, j + data_size, k, l],
+                    data[i, j - buflen + data_size, k, l],
+                ),
+                name="new_buffer",
+            )
+        if axis == 2:
+            return te.compute(
+                buffer.shape,
+                lambda i, j, k, l: tvm.tirx.if_then_else(
+                    k < buflen - data_size,
+                    buffer[i, j, k + data_size, l],
+                    data[i, j, k - buflen + data_size, l],
+                ),
+                name="new_buffer",
+            )
+        if axis == 3:
+            return te.compute(
+                buffer.shape,
+                lambda i, j, k, l: tvm.tirx.if_then_else(
+                    l < buflen - data_size,
+                    buffer[i, j, k, l + data_size],
+                    data[i, j, k, l - buflen + data_size],
+                ),
+                name="new_buffer",
+            )
+        assert False, f"Invalid value for axis; it should be at most {len(buffer.shape)}"
+    else:
+        # Implement FIFO buffer as combination of concat and slice
+        begin = [0] * len(buffer.shape)
+        begin[axis] = data.shape[axis]
+        end = list(buffer.shape[:])
+        end[axis] += data.shape[axis]
+        return strided_slice(concatenate((buffer, data), axis=axis), begin=begin, end=end)
+    return None

python/tvm/topi/nn/flatten.py

@@ -0,0 +1,54 @@
+# 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.
+"""TVM operator flatten compute."""
+
+import tvm
+from tvm import te
+
+from .. import tag
+
+
+@tvm.te.tag_scope(tag=tag.INJECTIVE)
+def flatten(data):
+    """Flattens the input array into a 2-D array by collapsing the higher dimensions.
+
+    Parameters
+    ----------
+    data : tvm.te.Tensor
+        Input array.
+
+    Returns
+    -------
+    output : tvm.te.Tensor
+        2-D array with collapsed higher dimensions.
+    """
+    ishape = data.shape
+    dim = 1
+    for i in range(1, len(ishape)):
+        dim = dim * ishape[i]
+    oshape = [ishape[0], dim]
+    idxdiv = tvm.tirx.indexdiv
+    idxmod = tvm.tirx.indexmod
+
+    def unwrap(idx, shape):
+        index = []
+        for s in reversed(shape):
+            index.append(idxmod(idx, s))
+            idx = idxdiv(idx, s)
+        return list(reversed(index))
+
+    return te.compute(oshape, lambda i, j: data(i, *unwrap(j, ishape[1:])))

python/tvm/topi/nn/group_norm.py

@@ -0,0 +1,55 @@
+# 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.
+"""Layer normalization operator."""
+
+from .. import cpp
+
+
+def group_norm(data, gamma, beta, num_groups, channel_axis, axes, epsilon=1e-5):
+    """Group normalization operator.
+    It accepts fp16 and fp32 as input data type. It will cast the input to fp32
+    to perform the computation. The output will have the same data type as input.
+
+    Parameters
+    ----------
+    data : tvm.te.Tensor
+        N-D with shape (d_0, d_1, ..., d_{N-1})
+
+    gamma: tvm.te.Tensor
+        1-D with shape (r_0) where r_0 == d_{channel_axis}
+
+    beta: tvm.te.Tensor
+        Optional, 1-D with shape (r_0) where r_0 == d_{channel_axis}
+
+    num_groups : int
+        The number of groups
+
+    channel_axis : int
+        The channel axis
+
+    axes : list of int
+        Axis over the normalization applied, excluding the channel axis
+
+    epsilon : float
+        The epsilon value to avoid division by zero.
+
+    Returns
+    -------
+    result : tvm.te.Tensor
+        N-D with shape (d_0, d_1, ..., d_{N-1})
+    """
+    return cpp.nn.group_norm(data, gamma, beta, num_groups, channel_axis, axes, epsilon)

python/tvm/topi/nn/instance_norm.py

@@ -0,0 +1,48 @@
+# 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.
+"""Instance normalization operator."""
+
+from .. import cpp
+
+
+def instance_norm(data, gamma, beta, channel_axis, axis, epsilon=1e-5):
+    """Instance normalization operator.
+
+    Parameters
+    ----------
+    data : tvm.te.Tensor
+        N-D with shape (d_0, d_1, ..., d_{N-1})
+
+    gamma: tvm.te.Tensor
+        K-D with shape (r_0, r_1, ..., r_{K-1}) where K == len(axis) and d_{axis_k} == r_k
+
+    beta: tvm.te.Tensor
+        Optional, K-D with shape (r_0, r_1, ..., r_{K-1}) where K == len(axis) and d_{axis_k} == r_k
+
+    axis : list of int
+        Axis over the normalization applied (the axis along which the mean and variance are
+        computed)
+
+    epsilon : float
+        The epsilon value to avoid division by zero.
+
+    Returns
+    -------
+    result : tvm.te.Tensor
+        N-D with shape (d_0, d_1, ..., d_{N-1})
+    """
+    return cpp.nn.instance_norm(data, gamma, beta, channel_axis, axis, epsilon)

python/tvm/topi/nn/layer_norm.py

@@ -0,0 +1,49 @@
+# 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.
+"""Layer normalization operator."""
+
+from .. import cpp
+
+
+def layer_norm(data, gamma, beta, axis, epsilon=1e-5):
+    """Layer normalization operator.
+    It accepts fp16 and fp32 as input data type. It will cast the input to fp32
+    to perform the computation. The output will have the same data type as input.
+
+    Parameters
+    ----------
+    data : tvm.te.Tensor
+        N-D with shape (d_0, d_1, ..., d_{N-1})
+
+    gamma: tvm.te.Tensor
+        K-D with shape (r_0, r_1, ..., r_{K-1}) where K == len(axis) and d_{axis_k} == r_k
+
+    beta: tvm.te.Tensor
+        Optional, K-D with shape (r_0, r_1, ..., r_{K-1}) where K == len(axis) and d_{axis_k} == r_k
+
+    axis : list of int
+        Axis over the normalization applied
+
+    epsilon : float
+        The epsilon value to avoid division by zero.
+
+    Returns
+    -------
+    result : tvm.te.Tensor
+        N-D with shape (d_0, d_1, ..., d_{N-1})
+    """
+    return cpp.nn.layer_norm(data, gamma, beta, axis, epsilon)

python/tvm/topi/nn/local_response_norm.py

@@ -0,0 +1,59 @@
+# 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
+"""TVM operator for local response norm compute."""
+
+from .. import cpp
+
+
+def lrn(data, size, axis=1, alpha=0.0001, beta=0.75, bias=2):
+    """Perform the across channels local response normalisation
+    on the input data.
+
+    sum_sqr_up^i{x, y} = (bias+((alpha/size)* \
+                                {sum_{j=max(0, i-size/2)}^{min(N-1,i+size/2)} \
+                                     (data^j{x,y})^2}))^beta
+    output^i{x, y} = data^i{x, y}/sum_sqr_up^i{x, y}
+    N is the number for input channels
+
+    Parameters
+    ----------
+    data : tvm.te.Tensor
+        4-D with shape [batch, channel, height, width]
+
+    size : int
+        normalisation window size
+
+    axis : int
+        input data layout channel axis
+        default value is 1 for NCHW format
+
+    bias : float
+        offset to avoid dividing by 0
+
+    alpha : float
+        to be divided
+
+    beta : float
+        exponent
+
+    Returns
+    -------
+    output : tvm.te.Tensor
+        4-D output with same shape
+    """
+    return cpp.nn.lrn(data, size, axis, alpha, beta, bias)

python/tvm/topi/nn/loss.py

@@ -0,0 +1,60 @@
+# 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,unused-argument
+"""Loss functions definitions."""
+
+from . import cpp
+
+
+def nll_loss(predictions, targets, weights, reduction, ignore_index):
+    """Negative log likelihood loss on the input data.
+
+    output{n, i_1, i_2, ..., i_k} = -p * w
+      where t = target{n, i_1, i_2, ..., i_k}
+            p = predictions{n, t, i_1, i_2, i_k}
+            w = weights{n, i_1, i_2, ..., i_k} if t != ignore_index else 0
+
+    result = reduction(output)
+
+    Parameters
+    ----------
+    predictions : tvm.te.Tensor
+        (k+2)-D with shape (N, C, d_1, d_2, ..., d_k),
+        where C is the number of target classes
+
+    targets : tvm.te.Tensor
+        (k+1)-D with shape (N, d_1, d_2, ..., d_k)
+        The target value of the input.
+
+    weights : tvm.te.Tensor
+        1-D with shape (C,)
+        The weight of each target value.
+
+    reduction : string
+        The reduction method to apply to output.
+        Can be "mean", "sum" or "none".
+
+    ignore_index : int
+        The target value to ignore.
+
+    Returns
+    -------
+    output : tvm.te.Tensor
+        a scalar if the reduction type is "mean" or "sum",
+        otherwise the same shape as `target`.
+    """
+    return cpp.nn.nll_loss(predictions, targets, weights, reduction, ignore_index)

python/tvm/topi/nn/lstm.py

@@ -0,0 +1,238 @@
+# 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: E731
+"""General LSTM implementation using TE scan."""
+
+from tvm import te, tirx
+from tvm.topi import tag
+
+
+def lstm(
+    Xs,
+    Wi,
+    Wh,
+    Bi=None,
+    Bh=None,
+    h_init=None,
+    c_init=None,
+    proj=None,
+    p_i=None,
+    p_f=None,
+    p_o=None,
+    f_act=tirx.sigmoid,
+    g_act=tirx.tanh,
+    h_act=tirx.tanh,
+    reverse=False,
+    weight_layout: str = "IFGO",
+):
+    """General LSTM implemented using TE scan.
+
+    Parameters
+    ----------
+    Xs : te.Tensor
+        Input sequence with shape `(seq_len, batch_size, in_dim)`
+    Wi : te.Tensor
+        Input weight matrix with shape `(4 * hidden_dim, in_dim)`. The weights are packed according
+        to `weight_layout`.
+    Wh : te.Tensor
+        Hidden weight matrix with shape `(4 * hidden_dim, hidden_dim or proj_dim)`. Packed as `Wh`.
+    Bi : te.Tensor, optional
+        Input bias with shape `(4 * hidden_dim,)`, by default None. Packed as `Wh`.
+    Bh : te.Tensor, optional
+        Hidden bias with shape as `Bi`, by default None. Packed as `Wh`.
+    h_init : te.Tensor, optional
+        Initial hidden state with shape `(batch_size, hidden_dim or proj_dim)`, zero if None
+    c_init : te.Tensor, optional
+        Initial cell state with same shape as `h_init`, zero if None
+    proj : te.Tensor, optional
+        Projection matrix with shape `(proj_dim, hidden_dim)`, by default None
+    p_i, p_f, p_o : te.Tensor, optional
+        Peephole LSTM matrices with shape `(batch_size, hidden_dim)`, by default None
+    f_act, g_act, h_act : F, optional
+        Gate activation functions
+    reverse : bool, optional
+        Whether to process `Xs` in reverse, by default False
+    weight_layout : str, optional
+        The packed weight layout for gates, by default "IFGO". Note: I = input, F = forget,
+        G = cell, O = output.
+
+    Returns
+    -------
+    result : te.Tensor, te.Tensor
+        Tuple of hidden states (with shape `(seq_len, batch_size, hidden_dim or proj_dim)`), and
+        cell states (with shape `(seq_len, batch_size, hidden_dim)`).
+    """
+    assert len(weight_layout) == 4 and sorted(weight_layout) == sorted("IFGO"), (
+        f'given weight layout "{weight_layout}" is not a permutation of "IFGO"'
+    )
+
+    i_gate_idx = weight_layout.find("I")
+    f_gate_idx = weight_layout.find("F")
+    g_gate_idx = weight_layout.find("G")
+    o_gate_idx = weight_layout.find("O")
+
+    seq_len, batch_size, in_dim = Xs.shape
+    assert Wi.shape[0] % 4 == 0, (
+        f"dim 0 of input weight should be 4 * hidden_dim, but {Wi.shape[0]} is not divisible by 4"
+    )
+    hidden_dim = Wi.shape[0] // 4
+    proj_dim = hidden_dim
+    if proj is not None:
+        proj_dim = proj.shape[0]
+
+    # te.scan uses up 1 element for the initial value
+    scan_len = seq_len + 1
+
+    # precompute input-hidden matmul outside the scan
+    ki = te.reduce_axis((0, in_dim), name="ki2h")
+    Xi2h = te.compute(
+        (seq_len * batch_size, 4 * hidden_dim),
+        lambda tb, ij: te.sum(Xs[(tb // batch_size), tb % batch_size, ki] * Wi[ij, ki], axis=ki),
+        name="Xi2h",
+    )
+    if Bi is not None:
+        Xi2h = te.compute(
+            Xi2h.shape, lambda tb, ij: Xi2h[tb, ij] + Bi[ij], name="Xi2h_bias", tag=tag.INJECTIVE
+        )
+
+    h_state = te.placeholder((scan_len, batch_size, proj_dim), name="h_state")
+    c_state = te.placeholder((scan_len, batch_size, hidden_dim), name="c_state")
+    h_init = te.compute(
+        (1, batch_size, proj_dim),
+        lambda _, b, i: h_init[b, i] if h_init is not None else 0.0,
+        name="h_init",
+    )
+    c_init = te.compute(
+        (1, batch_size, hidden_dim),
+        lambda _, b, i: c_init[b, i] if c_init is not None else 0.0,
+        name="c_init",
+    )
+
+    # begin scan computations, first the (batched) hidden-hidden dense
+    kh = te.reduce_axis((0, proj_dim), name="kh2h")
+    s_h2h = te.compute(
+        (scan_len, batch_size, 4, hidden_dim),
+        lambda t, b, i, j: te.sum(h_state[t - 1, b, kh] * Wh[i * hidden_dim + j, kh], axis=kh),
+        name="s_h2h",
+    )
+    if Bh is not None:
+        s_h2h = te.compute(
+            s_h2h.shape,
+            lambda t, b, i, j: s_h2h[t, b, i, j] + Bh[i * hidden_dim + j],
+            name="s_h2h_bias",
+            tag=tag.INJECTIVE,
+        )
+
+    # helper to reverse time if scanning backwards
+    get_x_t = lambda t: seq_len - t if reverse else t - 1
+
+    gates = te.compute(
+        (scan_len, batch_size, 4, hidden_dim),
+        lambda t, b, i, j: (
+            Xi2h[get_x_t(t) * batch_size + b, i * hidden_dim + j] + s_h2h[t, b, i, j]
+        ),
+        name="gates",
+        tag=tag.INJECTIVE,
+    )
+
+    # helper to correctly read each gate dense from the batched output
+    read_gate = lambda t, b, j, idx: gates[t, b, idx, j]
+
+    gate_shape = (scan_len, batch_size, hidden_dim)
+
+    # compute the activated gates (and do some extra stuff if peephole weights are present)
+    if p_i is not None and p_f is not None:
+        i_gate = te.compute(
+            gate_shape,
+            lambda t, b, j: f_act(
+                read_gate(t, b, j, i_gate_idx) + p_i[b, j] * c_state[t - 1, b, j]
+            ),
+            name="i_gate_p",
+            tag=tag.INJECTIVE,
+        )
+        f_gate = te.compute(
+            gate_shape,
+            lambda t, b, j: f_act(
+                read_gate(t, b, j, f_gate_idx) + p_f[b, j] * c_state[t - 1, b, j]
+            ),
+            name="f_gate_p",
+            tag=tag.INJECTIVE,
+        )
+    else:
+        i_gate = te.compute(
+            gate_shape,
+            lambda *i: f_act(read_gate(*i, i_gate_idx)),
+            name="i_gate",
+            tag=tag.INJECTIVE,
+        )
+        f_gate = te.compute(
+            gate_shape,
+            lambda *i: f_act(read_gate(*i, f_gate_idx)),
+            name="f_gate",
+            tag=tag.INJECTIVE,
+        )
+
+    g_gate = te.compute(
+        gate_shape, lambda *i: g_act(read_gate(*i, g_gate_idx)), name="g_gate", tag=tag.INJECTIVE
+    )
+
+    next_c = te.compute(
+        gate_shape,
+        lambda t, b, j: f_gate[t, b, j] * c_state[t - 1, b, j] + i_gate[t, b, j] * g_gate[t, b, j],
+        name="next_c",
+    )
+
+    if p_o is not None:
+        o_gate = te.compute(
+            gate_shape,
+            lambda t, b, j: f_act(read_gate(t, b, j, o_gate_idx) + p_o[b, j] * next_c[t, b, j]),
+            name="o_gate_p",
+            tag=tag.INJECTIVE,
+        )
+    else:
+        o_gate = te.compute(
+            gate_shape,
+            lambda *i: f_act(read_gate(*i, o_gate_idx)),
+            name="o_gate",
+            tag=tag.INJECTIVE,
+        )
+
+    next_h = te.compute(gate_shape, lambda *i: o_gate(*i) * h_act(next_c(*i)), name="next_h")
+
+    # project hidden state back to proj_dim if projection matrix is present
+    if proj is not None:
+        kr = te.reduce_axis((0, hidden_dim), name="kh2p")
+        next_h = te.compute(
+            (scan_len, batch_size, proj_dim),
+            lambda t, b, j: te.sum(next_h[t, b, kr] * proj[j, kr], axis=kr),
+            name="next_h_proj",
+        )
+
+    scan_h, scan_c = te.scan(
+        [h_init, c_init], [next_h, next_c], [h_state, c_state], name="lstm_scan"
+    )
+
+    # drop the initial values, TODO(@altanh): is there a better way?
+    scan_h = te.compute(
+        (seq_len, batch_size, proj_dim), lambda t, b, j: scan_h[t + 1, b, j], name="hidden_states"
+    )
+    scan_c = te.compute(
+        (seq_len, batch_size, hidden_dim), lambda t, b, j: scan_c[t + 1, b, j], name="cell_states"
+    )
+
+    return scan_h, scan_c

python/tvm/topi/nn/mapping.py

@@ -0,0 +1,100 @@
+# 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, line-too-long
+"""Operators of one-to-one-mapping on the first input"""
+
+import tvm
+from tvm import te
+
+from .. import tag
+
+
+@tvm.te.tag_scope(tag=tag.BROADCAST)
+def scale_shift_nchw(Input, Scale, Shift):
+    """Batch normalization operator in inference.
+
+    Parameters
+    ----------
+    Input : tvm.te.Tensor
+        4-D input tensor, NCHW layout [batch, channel, height, width]
+
+    Scale : tvm.te.Tensor
+        Scale tensor, 1-D of size channel number
+
+    Shift : tvm.te.Tensor
+        Shift tensor, 1-D of size channel number
+
+    Returns
+    -------
+    Output : tvm.te.Tensor
+        Output tensor, layout is NCHW
+    """
+    return te.compute(
+        Input.shape, lambda b, c, i, j: Input[b, c, i, j] * Scale[c] + Shift[c], name="ScaleShift"
+    )
+
+
+@tvm.te.tag_scope(tag=tag.BROADCAST)
+def scale_shift_nhwc(Input, Scale, Shift):
+    """Batch normalization operator in inference.
+
+    Parameters
+    ----------
+    Input : tvm.te.Tensor
+        4-D input tensor, NHWC layout [batch, height, width, channel]
+
+    Scale : tvm.te.Tensor
+        Scale tensor, 1-D of size channel number
+
+    Shift : tvm.te.Tensor
+        Shift tensor, 1-D of size channel number
+
+    Returns
+    -------
+    Output : tvm.te.Tensor
+        Output tensor, layout is NHWC
+    """
+    return te.compute(
+        Input.shape, lambda b, i, j, c: Input[b, i, j, c] * Scale[c] + Shift[c], name="ScaleShift"
+    )
+
+
+@tvm.te.tag_scope(tag=tag.BROADCAST)
+def scale_shift_nchwc(Input, Scale, Shift):
+    """Batch normalization operator in inference.
+
+    Parameters
+    ----------
+    Input : tvm.te.Tensor
+        5-D input tensor, NCHWc layout [batch, channel_chunk, height, width, channel_block]
+
+    Scale : tvm.te.Tensor
+        Scale tensor, 2-D of size [channel_chunk, channel_block]
+
+    Shift : tvm.te.Tensor
+        Shift tensor, 2-D of size [channel_chunk, channel_block]
+
+    Returns
+    -------
+    Output : tvm.te.Tensor
+        Output tensor, layout is NHWC
+    """
+    return te.compute(
+        Input.shape,
+        lambda b, cc, i, j, cb: Input[b, cc, i, j, cb] * Scale[cc, cb] + Shift[cc, cb],
+        name="ScaleShift",
+    )

python/tvm/topi/nn/pad.py

@@ -0,0 +1,316 @@
+# 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.
+"""Pad the data by constant value"""
+
+import tvm
+from tvm import te
+from tvm.tirx import if_then_else
+
+from .. import tag
+from ..utils import equal_const_int
+
+
+def get_padded_shape(data, pad_before, pad_after=None):
+    """
+    Calculates the output shape of a tensor after applying padding.
+
+    Args:
+        data (tvm.te.Tensor): The input tensor to which padding is applied.
+        pad_before : list / tuple of n ints
+            Pad width on each dimension to pad the before the axis begin.
+        pad_after : list / tuple of n ints, optional
+            Pad width each dimension to pad the after the axis end.
+
+    Raises:
+        ValueError: If `pad_before` or `pad_after` lengths mismatch with `data` dimensions.
+
+    Returns:
+        tuple: A tuple representing the padded shape of the tensor.
+    """
+    n = data.ndim
+    pad_after = pad_after if pad_after else pad_before
+
+    if len(pad_before) != n:
+        raise ValueError(f"pad_before length {len(pad_before)} != input dims {n}")
+    if len(pad_after) != n:
+        raise ValueError(f"pad_after length {len(pad_after)} != input dims {n}")
+
+    ana = tvm.arith.Analyzer()
+    out_shape = tuple(ana.simplify(data.shape[i] + pad_before[i] + pad_after[i]) for i in range(n))
+
+    return out_shape
+
+
+@tvm.te.tag_scope(tag=tag.INJECTIVE + ",pad")
+def pad(data, pad_before, pad_after=None, pad_value=0.0, name="PadInput", attrs=None):
+    """Pad Input with using pad values.
+
+    Parameters
+    ----------
+    data : tvm.te.Tensor
+        n-D input, can be any layout.
+
+    pad_before : list / tuple of n ints
+        Pad width on each dimension to pad the before the axis begin.
+
+    pad_after : list / tuple of n ints, optional
+        Pad width each dimension to pad the after the axis end.
+
+    pad_value : float, optional
+        The value to be padded.
+
+    name : str, optional
+        The name prefix operators generated
+
+    Returns
+    -------
+    Output : tvm.te.Tensor
+        n-D, the same layout as Input.
+    """
+    n = len(data.shape)
+    pad_after = pad_after if pad_after else pad_before
+    if len(pad_before) != n:
+        raise ValueError(f"Input dimension and pad_before dismatch : {n} vs {len(pad_before)}")
+    if len(pad_after) != n:
+        raise ValueError(f"Input dimension and pad_after dismatch : {n} vs {len(pad_after)}")
+    ana = tvm.arith.Analyzer()
+    dshape = []
+    for dim in data.shape:
+        dshape.append(dim)
+    out_shape = tuple(ana.simplify(dshape[i] + pad_before[i] + pad_after[i]) for i in range(n))
+    pad_value = (
+        pad_value if tvm.ir.is_prim_expr(pad_value) else tvm.tirx.const(pad_value, data.dtype)
+    )
+
+    def _pad(*indices):
+        not_zero = []
+        index_tuple = []
+        for i in range(n):
+            if equal_const_int(pad_before[i], 0) and equal_const_int(pad_after[i], 0):
+                index_tuple.append(indices[i])
+            else:
+                index_tuple.append(indices[i] - pad_before[i])
+                not_zero.append(indices[i] >= pad_before[i])
+                not_zero.append(indices[i] < data.shape[i] + pad_before[i])
+        if not_zero:
+            not_zero = tvm.tirx.all(*not_zero)
+            return tvm.tirx.if_then_else(not_zero, data(*index_tuple), pad_value)
+        return data(*index_tuple)
+
+    return te.compute(out_shape, _pad, name=name, attrs=attrs)
+
+
+@tvm.te.tag_scope(tag=tag.INJECTIVE + ",pad")
+def mirror_pad(data, pad_before, pad_after=None, mode="SYMMETRIC", name="MirrorPadInput"):
+    """Pad Input with mirroring either symmetric or reflected.
+
+    Parameters
+    ----------
+    data : tvm.te.Tensor
+        n-D input, can be any layout.
+
+    pad_before : list / tuple of n ints
+        Pad width on each dimension to pad the before the axis begin.
+
+    pad_after : list / tuple of n ints, optional
+        Pad width each dimension to pad the after the axis end.
+
+    mode: str, optional
+        Type of mirror padding to apply. Must be SYMMETRIC or REFLECT
+
+    name : str, optional
+        The name prefix operators generated
+
+    Returns
+    -------
+    Output : tvm.te.Tensor
+        n-D, the same layout as Input.
+    """
+    n = len(data.shape)
+    pad_after = pad_after if pad_after else pad_before
+    if len(pad_before) != n:
+        raise ValueError(f"Input dimension and pad_before dismatch : {n} vs {len(pad_before)}")
+    if len(pad_after) != n:
+        raise ValueError(f"Input dimension and pad_after dismatch : {n} vs {len(pad_after)}")
+    ana = tvm.arith.Analyzer()
+    out_shape = tuple(ana.simplify(data.shape[i] + pad_before[i] + pad_after[i]) for i in range(n))
+    assert mode in ("SYMMETRIC", "REFLECT")
+    mode = int(mode == "SYMMETRIC")
+
+    def _pad(*indices):
+        index_tuple = []
+        above = []
+        below = []
+        for i in range(n):
+            if equal_const_int(pad_before[i], 0) and equal_const_int(pad_after[i], 0):
+                index_tuple.append(indices[i])
+                above.append(False)
+                below.append(False)
+            else:
+                index_tuple.append(indices[i] - pad_before[i])
+                above.append(indices[i] >= data.shape[i] + pad_before[i])
+                below.append(indices[i] < pad_before[i])
+        mapped_tuple = []
+        for i, axis in enumerate(index_tuple):
+            mapped_axis = tvm.tirx.if_then_else(below[i], -axis - mode, axis)
+            mapped_axis = tvm.tirx.if_then_else(
+                above[i], (2 * (data.shape[i] - 1)) - axis + mode, mapped_axis
+            )
+            mapped_tuple.append(mapped_axis)
+        return data(*mapped_tuple)
+
+    return te.compute(out_shape, _pad, name=name)
+
+
+@tvm.te.tag_scope(tag=tag.INJECTIVE + ",pad")
+def reflect_pad(data, pad_before, pad_after=None, name="ReflectPadInput"):
+    """
+    Apply reflect padding to the input tensor.
+
+    Parameters
+    ----------
+    data : tvm.te.Tensor
+        Input tensor.
+
+    pad_before : List[int]
+        Amount to pad before each dimension.
+
+    pad_after : List[int], optional
+        Amount to pad after each dimension. If None, defaults to pad_before.
+
+    name : str
+        Name of the resulting tensor.
+
+    Returns
+    -------
+    out : tvm.te.Tensor
+        Reflect-padded tensor.
+    """
+    out_shape = get_padded_shape(data, pad_before, pad_after)
+
+    def _pad(*indices):
+        index_tuple = []
+        for i in range(data.ndim):
+            idx = indices[i]
+            size = data.shape[i]
+            before = pad_before[i]
+
+            orig_idx = idx - before
+
+            reflected_idx = if_then_else(
+                orig_idx < 0,
+                -orig_idx,  # reflect from start (no repeat)
+                if_then_else(
+                    orig_idx >= size,
+                    (2 * size - 2) - orig_idx,  # reflect from end
+                    orig_idx,
+                ),
+            )
+            index_tuple.append(reflected_idx)
+        return data(*index_tuple)
+
+    return te.compute(out_shape, _pad, name=name)
+
+
+@tvm.te.tag_scope(tag=tag.INJECTIVE + ",pad")
+def replicate_pad(data, pad_before, pad_after=None, name="ReplicatePadInput"):
+    """
+    Apply replicate padding (edge padding) to the input tensor.
+
+    Parameters
+    ----------
+    data : tvm.te.Tensor
+        Input tensor.
+
+    pad_before : List[int]
+        Amount to pad before each dimension.
+
+    pad_after : List[int], optional
+        Amount to pad after each dimension. If None, defaults to pad_before.
+
+    name : str
+        Name of the resulting tensor.
+
+    Returns
+    -------
+    out : tvm.te.Tensor
+        Replicate-padded tensor.
+    """
+    out_shape = get_padded_shape(data, pad_before, pad_after)
+
+    def _pad(*indices):
+        index_tuple = []
+        for i in range(data.ndim):
+            idx = indices[i]
+            size = data.shape[i]
+            before = pad_before[i]
+
+            orig_idx = idx - before
+            clamped_idx = if_then_else(
+                orig_idx < 0,
+                tvm.tirx.const(0, "int32"),  # replicate first element
+                if_then_else(
+                    orig_idx >= size,
+                    size - 1,  # replicate last element
+                    orig_idx,
+                ),
+            )
+            index_tuple.append(clamped_idx)
+        return data(*index_tuple)
+
+    return te.compute(out_shape, _pad, name=name)
+
+
+@tvm.te.tag_scope(tag=tag.INJECTIVE + ",pad")
+def circular_pad(data, pad_before, pad_after=None, name="CircularPadInput"):
+    """
+    Apply circular padding (wrap around) to the input tensor.
+
+    Parameters
+    ----------
+    data : tvm.te.Tensor
+        Input tensor.
+
+    pad_before : List[int]
+        Amount to pad before each dimension.
+
+    pad_after : List[int], optional
+        Amount to pad after each dimension. If None, defaults to pad_before.
+
+    name : str
+        Name of the resulting tensor.
+
+    Returns
+    -------
+    out : tvm.te.Tensor
+        Circular-padded tensor.
+    """
+    out_shape = get_padded_shape(data, pad_before, pad_after)
+
+    def _pad(*indices):
+        index_tuple = []
+        for i in range(data.ndim):
+            idx = indices[i]
+            size = data.shape[i]
+            before = pad_before[i]
+
+            orig_idx = idx - before
+            wrapped_idx = tvm.tirx.indexmod(orig_idx + size, size)
+            index_tuple.append(wrapped_idx)
+        return data(*index_tuple)
+
+    return te.compute(out_shape, _pad, name=name)

python/tvm/topi/nn/pixel_shuffle.py

@@ -0,0 +1,75 @@
+# 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.
+# ruff: noqa: RUF005
+"""TVM operator pixel shuffle compute."""
+
+import tvm
+
+
+def pixel_shuffle(data, upscale_factor, name="PixelShuffle"):
+    """PixelShuffle operator that rearranges elements in a tensor of shape
+    [..., C * r * r, H, W] to [..., C, H * r, W * r].
+
+    Parameters
+    ----------
+    data : tvm.te.Tensor
+        N-D input tensor with at least 3 dimensions. Channel must be at index -3.
+
+    upscale_factor : int
+        The upscale factor (r).
+
+    name : str
+        Name of the output tensor.
+
+    Returns
+    -------
+    output : tvm.te.Tensor
+        Pixel shuffled tensor with shape [..., C, H*r, W*r]
+    """
+    assert isinstance(upscale_factor, int) and upscale_factor > 0
+    ndim = len(data.shape)
+    assert ndim >= 3, "Input must be at least 3D"
+
+    upscale_factor_const = tvm.tirx.const(upscale_factor, "int32")
+    c_in, h_in, w_in = data.shape[-3], data.shape[-2], data.shape[-1]
+
+    c_out = tvm.tirx.floordiv(c_in, upscale_factor_const * upscale_factor_const)
+    h_out = h_in * upscale_factor_const
+    w_out = w_in * upscale_factor_const
+
+    out_shape = list(data.shape[:-3]) + [c_out, h_out, w_out]
+
+    def _compute(*indices):
+        batch_indices = indices[:-3]
+        c_out_idx, h_out_idx, w_out_idx = indices[-3], indices[-2], indices[-1]
+
+        h_idx = tvm.tirx.floordiv(h_out_idx, upscale_factor_const)
+        h_offset = h_out_idx % upscale_factor_const
+
+        w_idx = tvm.tirx.floordiv(w_out_idx, upscale_factor_const)
+        w_offset = w_out_idx % upscale_factor_const
+
+        c_in_idx = (
+            (c_out_idx * upscale_factor_const * upscale_factor_const)
+            + (h_offset * upscale_factor_const)
+            + w_offset
+        )
+
+        index_tuple = batch_indices + (c_in_idx, h_idx, w_idx)
+        return data[index_tuple]
+
+    return tvm.te.compute(out_shape, _compute, name=name)

python/tvm/topi/nn/pooling.py

@@ -0,0 +1,406 @@
+# 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.
+"""TVM operator pooling compute."""
+
+from .. import cpp
+
+POOL_TYPE_CODE = {"avg": 0, "max": 1}
+
+
+def global_pool(data, pool_type, layout="NCHW"):
+    """Perform global pooling on height and width dimension of data.
+       It decides the height and width dimension according to the layout string,
+       in which 'W' and 'H' means width and height respectively.
+       Width and height dimension cannot be split.
+       For example, NCHW, NCHW16c, etc. are valid for pool,
+       while NCHW16w, NCHW16h are not.
+       See parameter `layout` for more information of the layout string convention.
+
+    Parameters
+    ----------
+    data : tvm.te.Tensor
+        n-D with shape of layout
+
+    pool_type : str
+        Pool type, 'max' or 'avg'
+
+    layout : str
+        Layout of the input data.
+        The layout is supposed to be composed of upper cases, lower cases and numbers,
+        where upper case indicates a dimension and
+        the corresponding lower case with factor size indicates the split dimension.
+        For example, NCHW16c can describe a 5-D tensor of
+        [batch_size, channel, height, width, channel_block],
+        in which channel_block=16 is a split of dimension channel.
+
+    Returns
+    -------
+    output : tvm.te.Tensor
+        n-D in same layout with height and width dimension size of 1.
+        e.g., for NCHW, the output shape will be [batch, channel, 1, 1]
+    """
+    return cpp.nn.global_pool(data, POOL_TYPE_CODE[pool_type], layout)
+
+
+def pool_grad(
+    grads,
+    data,
+    kernel,
+    stride,
+    padding,
+    pool_type,
+    ceil_mode=False,
+    count_include_pad=True,
+    layout="NCHW",
+):
+    """Gradient of pooling on height and width dimension of data.
+       It decides the height and width dimension according to the layout string,
+       in which 'W' and 'H' means width and height respectively.
+       Width and height dimension cannot be split.
+       For example, NCHW, NCHW16c, etc. are valid for pool,
+       while NCHW16w, NCHW16h are not.
+       See parameter `layout` for more information of the layout string convention.
+
+    Parameters
+    ----------
+    grads : tvm.te.Tensor
+        n-D with shape of layout
+
+    data : tvm.te.Tensor
+        n-D with shape of layout
+
+    kernel : list/tuple of two ints
+        Kernel size, [kernel_height, kernel_width]
+
+    stride : list/tuple of two ints
+        Stride size, [stride_height, stride_width]
+
+    padding : list/tuple of four ints
+        Pad size, [pad_top, pad_left, pad_bottom, pad_right]]
+
+    pool_type : str
+        Pool type, 'max' or 'avg'
+
+    ceil_mode : bool
+        Whether to use ceil when calculating output size.
+
+    count_include_pad: bool
+        Whether include padding in the calculation when pool_type is 'avg'
+
+    layout: string
+        Layout of the input data.
+        The layout is supposed to be composed of upper cases, lower cases and numbers,
+        where upper case indicates a dimension and
+        the corresponding lower case with factor size indicates the split dimension.
+        For example, NCHW16c can describe a 5-D tensor of
+        [batch_size, channel, height, width, channel_block],
+        in which channel_block=16 is a split of dimension channel.
+
+
+    Returns
+    -------
+    output : tvm.te.Tensor
+        n-D in the same layout
+    """
+    return cpp.nn.pool_grad(
+        grads,
+        data,
+        kernel,
+        stride,
+        padding,
+        POOL_TYPE_CODE[pool_type],
+        ceil_mode,
+        layout,
+        count_include_pad,
+    )
+
+
+def adaptive_pool(data, output_size, pool_type, layout="NCHW"):
+    """Perform pooling on height and width dimension of data.
+       The pooling kernel and stride sizes are automatically chosen for desired
+       output sizes.
+       It decides the height and width dimension according to the layout string,
+       in which 'W' and 'H' means width and height respectively.
+       Width and height dimension cannot be split.
+       For example, NCHW, NCHW16c, etc. are valid for pool,
+       while NCHW16w, NCHW16h are not.
+       See parameter `layout` for more information of the layout string convention.
+
+    Parameters
+    ----------
+    data : tvm.te.Tensor
+        n-D with shape of layout
+
+    output_size : tuple of int
+        output height and width.
+
+    pool_type : str
+        Pool type, 'max' or 'avg'
+
+    layout: string
+        Layout of the input data.
+        The layout is supposed to be composed of upper cases, lower cases and numbers,
+        where upper case indicates a dimension and
+        the corresponding lower case with factor size indicates the split dimension.
+        For example, NCHW16c can describe a 5-D tensor of
+        [batch_size, channel, height, width, channel_block],
+        in which channel_block=16 is a split of dimension channel.
+
+    Returns
+    -------
+    output : tvm.te.Tensor
+        n-D in the same layout
+    """
+    return cpp.nn.adaptive_pool(data, output_size, POOL_TYPE_CODE[pool_type], layout)
+
+
+def adaptive_pool1d(data, output_size, pool_type, layout="NCW"):
+    """Perform pooling on three dimensional data.
+    See the two dimensional version above for details.
+    """
+    return cpp.nn.adaptive_pool1d(data, output_size, POOL_TYPE_CODE[pool_type], layout)
+
+
+def adaptive_pool3d(data, output_size, pool_type, layout="NCDHW"):
+    """Perform pooling on three dimensional data.
+    See the two dimensional version above for details.
+    """
+    return cpp.nn.adaptive_pool3d(data, output_size, POOL_TYPE_CODE[pool_type], layout)
+
+
+def pool1d(
+    data,
+    kernel,
+    stride,
+    dilation,
+    padding,
+    pool_type,
+    ceil_mode=False,
+    layout="NCW",
+    count_include_pad=True,
+):
+    """Perform pooling on width dimension of data.
+       Width axis is determined according to the layout string.
+       in which 'w' means width.
+       Width dimension cannot be split.
+       For example, NCW, NCW16c, etc. are valid for pool,
+       while NCW16w is not.
+       See parameter `layout` for more information of the layout string convention.
+
+    Parameters
+    ----------
+    data : tvm.te.Tensor
+        n-D with shape of layout
+
+    kernel : list/tuple of one int or int
+        Kernel size, [kernel_width]
+
+    stride : list/tuple of one int or int
+        Stride size, [stride_width]
+
+    dilation: list/tuple of two ints
+        Dilation size, [dilation_height, dilation_width]
+
+    padding : list/tuple of two ints
+        Pad size, [pad_left, pad_right]
+
+    pool_type : str
+        Pool type, 'max' or 'avg'
+
+    ceil_mode : bool
+        Whether to use ceil when calculating output size.
+
+    layout: string
+        Layout of the input data.
+        The layout is supposed to be composed of upper cases, lower cases and numbers,
+        where upper case indicates a dimension and
+        the corresponding lower case with factor size indicates the split dimension.
+        For example, NCW16c can describe a 4-D tensor of
+        [batch_size, channel, width, channel_block],
+        in which channel_block=16 is a split of dimension channel.
+
+    count_include_pad: bool
+        Whether include padding in the calculation when pool_type is 'avg'
+
+    Returns
+    -------
+    output : tvm.te.Tensor
+        n-D in the same layout
+    """
+    if isinstance(kernel, int):
+        kernel = [
+            kernel,
+        ]
+    if isinstance(stride, int):
+        stride = [
+            stride,
+        ]
+    return cpp.nn.pool1d(
+        data,
+        kernel,
+        stride,
+        dilation,
+        padding,
+        POOL_TYPE_CODE[pool_type],
+        ceil_mode,
+        layout,
+        count_include_pad,
+    )
+
+
+def pool2d(
+    data,
+    kernel,
+    stride,
+    dilation,
+    padding,
+    pool_type,
+    ceil_mode=False,
+    layout="NCHW",
+    count_include_pad=True,
+):
+    """Perform pooling on height and width dimension of data.
+       It decides the height and width dimension according to the layout string,
+       in which 'W' and 'H' means width and height respectively.
+       Width and height dimension cannot be split.
+       For example, NCHW, NCHW16c, etc. are valid for pool,
+       while NCHW16w, NCHW16h are not.
+       See parameter `layout` for more information of the layout string convention.
+
+    Parameters
+    ----------
+    data : tvm.te.Tensor
+        n-D with shape of layout
+
+    kernel : list/tuple of two ints
+        Kernel size, [kernel_height, kernel_width]
+
+    stride : list/tuple of two ints
+        Stride size, [stride_height, stride_width]
+
+    dilation: list/tuple of two ints
+        Dilation size, [dilation_height, dilation_width]
+
+    padding : list/tuple of four ints
+        Pad size, [pad_top, pad_left, pad_bottom, pad_right]]
+
+    pool_type : str
+        Pool type, 'max' or 'avg'
+
+    ceil_mode : bool
+        Whether to use ceil when calculating output size.
+
+    layout: string
+        Layout of the input data.
+        The layout is supposed to be composed of upper cases, lower cases and numbers,
+        where upper case indicates a dimension and
+        the corresponding lower case with factor size indicates the split dimension.
+        For example, NCHW16c can describe a 5-D tensor of
+        [batch_size, channel, height, width, channel_block],
+        in which channel_block=16 is a split of dimension channel.
+
+    count_include_pad: bool
+        Whether include padding in the calculation when pool_type is 'avg'
+
+    Returns
+    -------
+    output : tvm.te.Tensor
+        n-D in the same layout
+    """
+    return cpp.nn.pool2d(
+        data,
+        kernel,
+        stride,
+        dilation,
+        padding,
+        POOL_TYPE_CODE[pool_type],
+        ceil_mode,
+        layout,
+        count_include_pad,
+    )
+
+
+def pool3d(
+    data,
+    kernel,
+    stride,
+    dilation,
+    padding,
+    pool_type,
+    ceil_mode=False,
+    layout="NCDHW",
+    count_include_pad=True,
+):
+    """Perform pooling on depth, height and width dimension of data.
+       It decides the depth, height and width dimension according to the layout string,
+       in which 'D', 'W' and 'H' means depth, width and height respectively.
+       Depth, width and height dimension cannot be split.
+       For example, NCDHW, NCDHW16c, etc. are valid for pool,
+       while NCDHW16d, NCDHW16w, NCDHW16h are not.
+       See parameter `layout` for more information of the layout string convention.
+
+    Parameters
+    ----------
+    data : tvm.te.Tensor
+        n-D with shape of layout
+
+    kernel : list/tuple of three ints
+        Kernel size, [kernel_depth, kernel_height, kernel_width]
+
+    stride : list/tuple of three ints
+        Stride size, [stride_depth, stride_height, stride_width]
+
+    dilation: list/tuple of two ints
+        Dilation size, [dilation_height, dilation_width]
+
+    padding : list/tuple of six ints
+        Pad size, [pad_front, pad_top, pad_left, pad_back, pad_bottom, pad_right]
+
+    pool_type : str
+        Pool type, 'max' or 'avg'
+
+    ceil_mode : bool
+        Whether to use ceil when calculating output size.
+
+    layout: string
+        Layout of the input data.
+        The layout is supposed to be composed of upper cases, lower cases and numbers,
+        where upper case indicates a dimension and
+        the corresponding lower case with factor size indicates the split dimension.
+        For example, NCDHW16c can describe a 6-D tensor of
+        [batch_size, channel, depth, height, width, channel_block],
+        in which channel_block=16 is a split of dimension channel.
+
+    count_include_pad: bool
+        Whether include padding in the calculation when pool_type is 'avg'
+
+    Returns
+    -------
+    output : tvm.te.Tensor
+        n-D in the same layout
+    """
+    return cpp.nn.pool3d(
+        data,
+        kernel,
+        stride,
+        dilation,
+        padding,
+        POOL_TYPE_CODE[pool_type],
+        ceil_mode,
+        layout,
+        count_include_pad,
+    )

python/tvm/topi/nn/qnn.py

@@ -0,0 +1,193 @@
+# 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.
+"""Quantized Neural Network (QNN) Operators"""
+
+import tvm
+from tvm import te, tirx, topi
+
+SQNN_DISABLE = 0
+SQNN_INT8 = 1
+SQNN_UINT8 = 2
+SQNN_INT32 = 3
+
+SQNN_DTYPE_TO_CODE = {
+    "disable": SQNN_DISABLE,
+    "int8": SQNN_INT8,
+    "uint8": SQNN_UINT8,
+    "int32": SQNN_INT32,
+}
+
+SQNN_CODE_TO_DTYPE = {v: k for k, v in SQNN_DTYPE_TO_CODE.items()}
+
+
+@tvm.te.tag_scope(tag=topi.tag.ELEMWISE)
+def simulated_quantize(data, out_dtype, output_scale=None, output_zero_point=None, axis=-1):
+    """Simulated QNN quantize operator that mimics QNN outputs without changing datatype.
+    The benefit of this operator over true QNN quantize is that this operator allows dynamic
+    datatype selection and can operate on both per-channel and scalar scales and zero points while
+    QNN quantize requires both of these to be fixed at compile time.
+
+    Parameters
+    ----------
+    data: tvm.te.Tensor
+        An N-D input tensor to the operator.
+
+    out_dtype: tvm.te.Tensor
+        A scalar variable that indicates which datatype to simulate quantization with. Use
+        SQNN_DTYPE_TO_CODE to convert a dtype string into the corresponding variable
+        value.
+
+    output_scale: tvm.te.Tensor, optional
+        A scalar tensor representing the scale to use when quantizing to integer datatypes.
+        When it contains more than a single value, N must match the number of channels in data.
+
+    output_zero_point: tvm.te.Tensor, optional
+        A 1-D tensor representing the zero point to use when quantizing to integer datatypes.
+        When it contains more than a single value, N must match the number of channels in data.
+
+    axis: int, optional
+        The channel axis for quantization. Default value is -1 which corresponds to the last axis.
+
+    """
+
+    # When disabled, just pass through the input values.
+    def _compute_pass_through(value, *indices):
+        return value[indices]
+
+    # Simulate quantization for arbitrary integer datatypes. The computation for all datatypes is:
+    # Q_output = clip((round(input_tensor/output_scale) + output_zero_point),
+    #                 out_dtype::min,
+    #                 out_dtype::max)
+    def _compute_intn(dtype, value, *indices):
+        assert output_scale is not None and output_zero_point is not None
+        const_min = tvm.tirx.min_value(dtype)
+        const_max = tvm.tirx.max_value(dtype)
+        # Use indexmod to handle both scalar and per-channel QNN parameters.
+        scale_idx = tirx.indexmod(indices[axis], topi.shape(output_scale)[0])
+        zp_idx = tirx.indexmod(indices[axis], topi.shape(output_zero_point)[0])
+        return te.max(
+            te.min(
+                te.round(value[indices] / output_scale[scale_idx]) + output_zero_point[zp_idx],
+                const_max,
+            ),
+            const_min,
+        )
+
+    # Use an if chain to dynamically return the proper quantization based on the input datatype.
+    # This allows the op to compile once but apply different quantization approaches
+    # using a variable datatype input.
+    def _dispatch_sim_quantize(value):
+        pass_through_value = te.compute(
+            data.shape, lambda *indices: _compute_pass_through(value, *indices)
+        )
+        int8_value = te.compute(
+            data.shape,
+            lambda *indices: tirx.if_then_else(
+                out_dtype.equal(SQNN_DTYPE_TO_CODE["int8"]),
+                _compute_intn("int8", value, *indices),
+                pass_through_value[indices],
+            ),
+        )
+        uint8_value = te.compute(
+            data.shape,
+            lambda *indices: tirx.if_then_else(
+                out_dtype.equal(SQNN_DTYPE_TO_CODE["uint8"]),
+                _compute_intn("uint8", value, *indices),
+                int8_value[indices],
+            ),
+        )
+        int32_value = te.compute(
+            data.shape,
+            lambda *indices: tirx.if_then_else(
+                out_dtype.equal(SQNN_DTYPE_TO_CODE["int32"]),
+                _compute_intn("int32", value, *indices),
+                uint8_value[indices],
+            ),
+        )
+
+        return int32_value
+
+    return te.compute(data.shape, lambda *indices: _dispatch_sim_quantize(data)[indices])
+
+
+@tvm.te.tag_scope(tag=topi.tag.ELEMWISE)
+def simulated_dequantize(data, in_dtype, input_scale=None, input_zero_point=None, axis=-1):
+    """Simulated QNN dequantize operator that mimics QNN outputs without changing datatype.
+    The benefit of this operator over true QNN dequantize is that this operator allows dynamic
+    datatype selection and can operate on both per-channel and scalar scales and zero points while
+    QNN dequantize requires both of these to be fixed at compile time.
+
+    Parameters
+    ----------
+    data: tvm.te.Tensor
+        An N-D input tensor to the operator.
+
+    in_dtype: tvm.te.Tensor
+        A scalar variable that indicates which datatype to simulate dequantization with. Use
+        SQNN_DTYPE_TO_CODE to convert a dtype string into the corresponding variable
+        value.
+
+    input_scale: tvm.te.Tensor, optional
+        A scalar tensor representing the scale to use when dequantizing from integer datatypes.
+        When it contains more than a single value, N must match the number of channels in data.
+
+    input_zero_point: tvm.te.Tensor, optional
+        A 1-D tensor representing the zero point to use when dequantizing from integer datatypes.
+        When it contains more than a single value, N must match the number of channels in data.
+
+    axis: int, optional
+        The channel axis for quantization. Default value is -1 which corresponds to the last axis.
+
+    """
+
+    # When disabled simply return the input tensor.
+    def _compute_pass_through(value, *indices):
+        return value[indices]
+
+    # Simulate dequantization for arbitrary integer datatypes. The computation for all datatypes is:
+    # DQ_output = (input - zero_point) * scale
+    def _compute_intn(value, *indices):
+        assert input_scale is not None and input_zero_point is not None
+        # Use indexmod to handle both scalar and per-channel QNN parameters.
+        scale_idx = tirx.indexmod(indices[axis], topi.shape(input_scale)[0])
+        zp_idx = tirx.indexmod(indices[axis], topi.shape(input_zero_point)[0])
+        return (value[indices] - input_zero_point[zp_idx]) * input_scale[scale_idx]
+
+    # Use an if chain to dynamically return the proper dequantization based on the input datatype.
+    # This allows the op to compile once but apply different quantization approaches
+    # using a variable datatype input.
+    def _dispatch_sim_dequantize(value):
+        pass_through_value = te.compute(
+            data.shape, lambda *indices: _compute_pass_through(value, *indices)
+        )
+        intn_condition = tvm.te.any(
+            in_dtype.equal(SQNN_DTYPE_TO_CODE["int8"]),
+            in_dtype.equal(SQNN_DTYPE_TO_CODE["uint8"]),
+            in_dtype.equal(SQNN_DTYPE_TO_CODE["int32"]),
+        )
+        intn_value = te.compute(
+            data.shape,
+            lambda *indices: tirx.if_then_else(
+                intn_condition,
+                _compute_intn(value, *indices),
+                pass_through_value[indices],
+            ),
+        )
+
+        return intn_value
+
+    return te.compute(data.shape, lambda *indices: _dispatch_sim_dequantize(data)[indices])

python/tvm/topi/nn/rms_norm.py

@@ -0,0 +1,44 @@
+# 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.
+"""Root mean square normalization operator."""
+
+from .. import cpp
+
+
+def rms_norm(data, weight, axis, epsilon=1e-5):
+    """Root mean square normalization operator. The output will have the same data type as input.
+
+    Parameters
+    ----------
+    data : tvm.te.Tensor
+        N-D with shape (d_0, d_1, ..., d_{N-1})
+
+    weight: tvm.te.Tensor
+        K-D with shape (r_0, r_1, ..., r_{K-1}) where K == len(axis) and d_{axis_k} == r_k
+
+    axis : list of int
+        Axis over the normalization applied
+
+    epsilon : float
+        The epsilon value to avoid division by zero.
+
+    Returns
+    -------
+    result : tvm.te.Tensor
+        N-D with shape (d_0, d_1, ..., d_{N-1})
+    """
+    return cpp.nn.rms_norm(data, weight, axis, epsilon)

python/tvm/topi/nn/softmax.py

@@ -0,0 +1,177 @@
+# 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, pointless-exception-statement
+# ruff: noqa: RUF005
+"""TVM operator for softmax and log_softmax compute."""
+
+import tvm
+from tvm import te, topi
+
+
+@tvm.te.tag_scope(tag="softmax_output")
+def softmax(x, axis=-1):
+    """Perform softmax activation on the data.
+
+    Parameters
+    ----------
+    x : tvm.te.Tensor
+        can be any dimension
+
+    axis : int
+        channel axis
+
+    Returns
+    -------
+    output : tvm.te.Tensor
+        output shape is the same as input
+    """
+    return softmax_common(x, axis, False)
+
+
+@tvm.te.tag_scope(tag="fast_softmax_output")
+def fast_softmax(x, axis=-1):
+    """Perform softmax activation on the data.
+    Use approximation to compute exponent for faster speed.
+
+    Parameters
+    ----------
+    x : tvm.te.Tensor
+        can be any dimension
+
+    axis : int
+        channel axis
+
+    Returns
+    -------
+    output : tvm.te.Tensor
+        output shape is the same as input
+    """
+    return softmax_common(x, axis, True)
+
+
+def softmax_common(x, axis, use_fast_exp):
+    """The common part of softmax and fast_softmax"""
+    shape = x.shape
+    if axis < 0:
+        axis = len(shape) + axis
+    if axis >= len(shape):
+        ValueError("axis parameter should be less than input dim")
+
+    k1 = te.reduce_axis((0, shape[axis]), name="k")
+    k2 = te.reduce_axis((0, shape[axis]), name="k")
+
+    def insert_reduce_index(indices, reduce_index):
+        return indices[:axis] + (reduce_index,) + indices[axis:]
+
+    def get_non_reduce_indices(indices):
+        return tuple([var for (i, var) in enumerate(indices) if i != axis])
+
+    def _compute_max(*indices):
+        eval_range = insert_reduce_index(indices, k1)
+        return tvm.te.max(x[eval_range], axis=k1)
+
+    def _compute_delta(max_elem, *indices):
+        non_reduce_indices = get_non_reduce_indices(indices)
+        return x[indices] - max_elem[non_reduce_indices]
+
+    def _compute_exp(max_elem, *indices):
+        non_reduce_indices = get_non_reduce_indices(indices)
+        return te.exp(x[indices] - max_elem[non_reduce_indices])
+
+    def _compute_expsum(exp, *indices):
+        eval_range = insert_reduce_index(indices, k2)
+        return te.sum(exp[eval_range], axis=k2)
+
+    def _normalize(exp, expsum, *indices):
+        non_reduce_indices = get_non_reduce_indices(indices)
+        return exp[indices] / expsum[non_reduce_indices]
+
+    reduced_shape = tuple([dim for (i, dim) in enumerate(shape) if i != axis])
+    max_elem = te.compute(reduced_shape, _compute_max, name="T_softmax_maxelem")
+
+    if use_fast_exp:
+        delta = te.compute(
+            shape, lambda *indices: _compute_delta(max_elem, *indices), name="T_softmax_delta"
+        )
+        exp = topi.math.fast_exp(delta)
+    else:
+        exp = te.compute(
+            shape, lambda *indices: _compute_exp(max_elem, *indices), name="T_softmax_exp"
+        )
+    expsum = te.compute(
+        reduced_shape, lambda *indices: _compute_expsum(exp, *indices), name="T_softmax_expsum"
+    )
+    return te.compute(
+        shape,
+        lambda *indices: _normalize(exp, expsum, *indices),
+        name="T_softmax_norm",
+        attrs={"axis": axis},
+    )
+
+
+@tvm.te.tag_scope(tag="log_softmax_output")
+def log_softmax(x, axis=-1):
+    """Perform log softmax activation on the data
+
+    Parameters
+    ----------
+    x : tvm.te.Tensor
+        N-D input data
+
+    axis : int
+        channel axis
+
+    Returns
+    -------
+    output : tvm.te.Tensor
+        N-D output with same shape
+    """
+    shape = x.shape
+    if axis < 0:
+        axis = len(shape) + axis
+    if axis >= len(shape):
+        ValueError("axis parameter should be less than input dim")
+
+    k1 = te.reduce_axis((0, shape[axis]), name="k")
+    k2 = te.reduce_axis((0, shape[axis]), name="k")
+
+    def insert_reduce_index(indices, reduce_index):
+        return indices[:axis] + (reduce_index,) + indices[axis:]
+
+    def get_non_reduce_indices(indices):
+        return tuple([var for (i, var) in enumerate(indices) if i != axis])
+
+    def _compute_max(*indices):
+        eval_range = insert_reduce_index(indices, k1)
+        return tvm.te.max(x[eval_range], axis=k1)
+
+    def _compute_expsum(max_elem, *indices):
+        eval_range = insert_reduce_index(indices, k2)
+        return te.sum(te.exp(x[eval_range] - max_elem[indices]), axis=k2)
+
+    def _normalize(max_elem, expsum, *indices):
+        non_reduce_indices = get_non_reduce_indices(indices)
+        return x[indices] - max_elem[non_reduce_indices] - te.log(expsum[non_reduce_indices])
+
+    reduced_shape = tuple([dim for (i, dim) in enumerate(shape) if i != axis])
+    max_elem = te.compute(reduced_shape, _compute_max, name="T_softmax_maxelem")
+    expsum = te.compute(reduced_shape, lambda *indices: _compute_expsum(max_elem, *indices))
+    return te.compute(
+        shape,
+        lambda *indices: _normalize(max_elem, expsum, *indices),
+        attrs={"axis": axis},
+    )

python/tvm/topi/nn/space_to_batch_nd.py

@@ -0,0 +1,52 @@
+# 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
+"""TVM operator space_to_batch_nd compute."""
+
+from . import cpp
+
+
+def space_to_batch_nd(data, block_shape, pad_before, pad_after, pad_value=0.0):
+    """Perform batch to space transformation on the data
+
+    Parameters
+    ----------
+    data : tvm.te.Tensor
+        N-D Tensor with shape [batch, spatial_shape, remaining_shapes],
+        where spatial_shape has M dimensions.
+
+    block_shape : list of ints
+        list of size [M] where M is number of spatial dims, specifies block
+        size for each spatial dimension.
+
+    pad_before : list of ints
+        list of shape [M] where M is number of spatial dims, specifies
+        zero-padding size before each spatial dimension.
+
+    pad_after : list of ints
+        list of shape [M] where M is number of spatial dims, specifies
+        zero-padding size after each spatial dimension.
+
+    pad_value : float, optional
+        The value used for padding.
+
+    Returns
+    -------
+    output : tvm.te.Tensor
+    """
+
+    return cpp.nn.space_to_batch_nd(data, block_shape, pad_before, pad_after, pad_value)

python/tvm/topi/nn/space_to_depth.py

@@ -0,0 +1,88 @@
+# 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
+"""TVM operator space_to_depth compute."""
+
+import tvm
+from tvm import te
+
+from .. import tag
+
+
+def space_to_depth(data, block_size, layout="NCHW"):
+    """Perform space to depth transformation on the data
+
+    Parameters
+    ----------
+    data : tvm.te.Tensor
+        4-D tensor in either NCHW or NHWC layout.
+
+    block_size : int
+        Size of blocks to decompose into channel dimension.
+
+    layout : string
+        Either NCHW or NHWC, indicating data layout.
+
+    Returns
+    -------
+    output : tvm.te.Tensor
+        Output of shape [N, C * block_size**2, H / block_size, W / block_size]
+    """
+
+    if layout == "NCHW":
+        in_n, in_c, in_h, in_w = data.shape
+        output_shape = [
+            in_n,
+            in_c * block_size * block_size,
+            tvm.tirx.truncdiv(in_h, block_size),
+            tvm.tirx.truncdiv(in_w, block_size),
+        ]
+    elif layout == "NHWC":
+        in_n, in_h, in_w, in_c = data.shape
+        output_shape = [
+            in_n,
+            tvm.tirx.truncdiv(in_h, block_size),
+            tvm.tirx.truncdiv(in_w, block_size),
+            in_c * block_size * block_size,
+        ]
+    else:
+        raise ValueError("Only NCHW and NHWC layouts are currently supported.")
+
+    def _get_indices(*indices):
+        if layout == "NCHW":
+            n, c, y, x = indices
+        elif layout == "NHWC":
+            n, y, x, c = indices
+        return n, c, y, x
+
+    def _get_pixel(n, c, y, x):
+        block_offset = tvm.tirx.truncdiv(c, in_c)
+        channel_idx = tvm.tirx.truncmod(c, in_c)
+        x_idx = tvm.tirx.truncmod(block_offset, block_size)
+        y_idx = tvm.tirx.truncdiv(block_offset, block_size)
+
+        if layout == "NCHW":
+            output = data(n, channel_idx, y_idx + (y * block_size), x_idx + (x * block_size))
+        else:
+            output = data(n, y_idx + (y * block_size), x_idx + (x * block_size), channel_idx)
+        return output
+
+    def _compute(*indices):
+        n, c, y, x = _get_indices(*indices)
+        return _get_pixel(n, c, y, x)
+
+    return te.compute(output_shape, _compute, name="space_to_depth", tag=tag.INJECTIVE)

python/tvm/topi/nn/upsampling.py

@@ -0,0 +1,204 @@
+# 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.
+"""TVM operator upsampling compute."""
+
+from tvm import te, topi
+
+from ..utils import simplify
+
+
+def upsampling(
+    data,
+    scale_h,
+    scale_w,
+    layout="NCHW",
+    method="nearest_neighbor",
+    align_corners=False,
+    output_shape=None,
+):
+    """Perform upsampling on the data.
+       Nearest neighbor and bilinear upsampling are supported.
+
+    Parameters
+    ----------
+    data : tvm.te.Tensor
+        inputs is a 4-D tensor with shape
+        [batch, channel, in_height, in_width]
+        or  [batch, in_height, in_width, channel]
+
+    scale_h : float
+        Scaling factor for height
+
+    scale_w : float
+        Scaling factor for width
+
+    layout : string, optional
+        either "NCHW" or "NHWC"
+
+    method : {"bilinear", "nearest_neighbor", "bicubic"}
+        Method to be used for upsampling.
+
+    output_shape: tvm_ffi.Array, optional
+        Shape to return. If left None will be inferred
+        (If shape is determined dynamically, pass out_dtype.shape as output_shape)
+
+    Returns
+    -------
+    output : tvm.te.Tensor
+        4-D with shape [batch, channel, in_height*scale_h, in_width*scale_w]
+        or [batch, in_height*scale, in_width*scale, channel]
+    """
+    base_layout = layout[0:4]
+    if base_layout == "NCHW":
+        if not output_shape:  # static case
+            scaled_h = data.shape[2] * scale_h
+            scaled_w = data.shape[3] * scale_w
+            reshape_size = (
+                simplify(topi.cast(te.round(scaled_h), data.shape[2].ty)),
+                simplify(topi.cast(te.round(scaled_w), data.shape[3].ty)),
+            )
+        else:  # dynamic case -- we don't need to scale; already done in shape func
+            reshape_size = (
+                simplify(topi.cast(te.round(output_shape[2]), output_shape[2].ty)),
+                simplify(topi.cast(te.round(output_shape[3]), output_shape[3].ty)),
+            )
+    elif layout == "NHWC":
+        if not output_shape:  # static case
+            scaled_h = data.shape[1] * scale_h
+            scaled_w = data.shape[2] * scale_w
+            reshape_size = (
+                simplify(topi.cast(te.round(scaled_h), data.shape[1].ty)),
+                simplify(topi.cast(te.round(scaled_w), data.shape[2].ty)),
+            )
+        else:  # dynamic case
+            reshape_size = (
+                simplify(topi.cast(te.round(output_shape[1]), output_shape[1].ty)),
+                simplify(topi.cast(te.round(output_shape[2]), output_shape[2].ty)),
+            )
+
+    else:
+        raise ValueError(f"not support this layout {layout} yet")
+    coord_trans = "align_corners" if align_corners else "asymmetric"
+    if method[0:2] == "bi":
+        method = method[2:]
+    return topi.image.resize2d(
+        data,
+        [0.0] * 4,
+        reshape_size,
+        layout=layout,
+        method=method,
+        coordinate_transformation_mode=coord_trans,
+        output_shape=output_shape,
+    )
+
+
+def upsampling3d(
+    data,
+    scale_d,
+    scale_h,
+    scale_w,
+    layout="NCDHW",
+    method="nearest_neighbor",
+    coordinate_transformation_mode="half_pixel",
+    output_shape=None,
+):
+    """Perform upsampling on the data.
+       Nearest neighbor and bilinear upsampling are supported.
+
+    Parameters
+    ----------
+    data : tvm.te.Tensor
+        inputs is a 5-D tensor with shape
+        [batch, channel, in_depth, in_height, in_width]
+        or  [batch, in_depth, in_height, in_width, channel]
+
+    scale_d : float
+        Scaling factor for depth
+
+    scale_h : float
+        Scaling factor for height
+
+    scale_w : float
+        Scaling factor for width
+
+    layout : string, optional
+        either "NCDHW" or "NDHWC"
+
+    method : {"trilinear", "nearest_neighbor"}
+        Method to be used for upsampling.
+
+    coordinate_transformation_mode: string, optional
+        Describes how to transform the coordinate in the resized tensor
+        to the coordinate in the original tensor.
+        Refer to the ONNX Resize operator specification for details.
+        Available options are "half_pixel", "align_corners" and "asymmetric".
+
+    output_shape: tvm_ffi.Array, optional
+        Shape to return. If left None will be inferred
+        (If shape is determined dynamically, pass out_dtype.shape as output_shape)
+
+    Returns
+    -------
+    output : tvm.te.Tensor
+        5-D with shape [batch, channel, in_depth*scale, in_height*scale, in_width*scale]
+        or [batch, in_depth*scale, in_height*scale, in_width*scale, channel]
+    """
+    base_layout = layout[0:5]
+    if base_layout == "NCDHW":
+        if not output_shape:  # static case
+            scaled_d = data.shape[2] * scale_d
+            scaled_h = data.shape[3] * scale_h
+            scaled_w = data.shape[4] * scale_w
+            resize_shape = (
+                simplify(topi.cast(te.round(scaled_d), data.shape[2].ty)),
+                simplify(topi.cast(te.round(scaled_h), data.shape[3].ty)),
+                simplify(topi.cast(te.round(scaled_w), data.shape[4].ty)),
+            )
+        else:  # dynamic case -- don't need to scale; already done in shape func
+            resize_shape = (
+                simplify(topi.cast(te.round(output_shape[2]), data.shape[2].ty)),
+                simplify(topi.cast(te.round(output_shape[3]), data.shape[3].ty)),
+                simplify(topi.cast(te.round(output_shape[4]), data.shape[4].ty)),
+            )
+    elif layout == "NDHWC":
+        if not output_shape:  # static case
+            scaled_d = data.shape[1] * scale_d
+            scaled_h = data.shape[2] * scale_h
+            scaled_w = data.shape[3] * scale_w
+            resize_shape = (
+                simplify(topi.cast(te.round(scaled_d), data.shape[1].ty)),
+                simplify(topi.cast(te.round(scaled_h), data.shape[2].ty)),
+                simplify(topi.cast(te.round(scaled_w), data.shape[3].ty)),
+            )
+        else:  # dynamic case
+            resize_shape = (
+                simplify(topi.cast(te.round(output_shape[1]), data.shape[1].ty)),
+                simplify(topi.cast(te.round(output_shape[2]), data.shape[2].ty)),
+                simplify(topi.cast(te.round(output_shape[3]), data.shape[3].ty)),
+            )
+    else:
+        raise ValueError(f"not support this layout {layout} yet")
+    if method[0:3] == "tri":
+        method = method[3:]
+    return topi.image.resize3d(
+        data,
+        [0.0] * 6,
+        resize_shape,
+        layout=layout,
+        method=method,
+        coordinate_transformation_mode=coordinate_transformation_mode,
+    )

python/tvm/topi/nn/utils.py

@@ -0,0 +1,310 @@
+# 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, unused-variable
+"""NN operator common utilities"""
+
+import tvm
+
+from ..utils import get_const_int
+
+
+def infer_pad(data, data_pad):
+    """Infer the padding from stages in reverse.
+
+    Parameters
+    ----------
+    data : Tensor
+        data stage.
+
+    data_pad : Tensor
+        pad stage.
+
+    Returns
+    -------
+    hpad : int
+        padding size on height
+    wpad : int
+        padding size on width
+    """
+    if data_pad is None:
+        return 0, 0
+    _, _, IH, IW = data.shape
+    _, _, TH, TW = data_pad.shape
+    hpad = (TH - IH) // 2
+    wpad = (TW - IW) // 2
+    return get_const_int(hpad), get_const_int(wpad)
+
+
+def infer_pad3d(data, data_pad, layout):
+    """Infer the padding from stages in reverse.
+
+    Parameters
+    ----------
+    data : Tensor
+        data stage.
+
+    data_pad : Tensor
+        pad stage.
+
+    Returns
+    -------
+    dpad : int
+        padding depth
+    hpad : int
+        padding height
+    wpad : int
+        padding width
+    """
+    if data_pad is None:
+        return 0, 0, 0
+
+    if layout == "NDHWC":
+        _, ID, IH, IW, _ = data.shape
+        _, TD, TH, TW, _ = data_pad.shape
+    elif layout == "NCDHW":
+        _, _, ID, IH, IW = data.shape
+        _, _, TD, TH, TW = data_pad.shape
+    else:
+        raise ValueError(f"Layout {layout} is not supported")
+    dpad = TD - ID
+    hpad = TH - IH
+    wpad = TW - IW
+    return get_const_int(dpad), get_const_int(hpad), get_const_int(wpad)
+
+
+def infer_stride(data, kernel, out):
+    """Infer the stride from stages in reverse.
+
+    Parameters
+    ----------
+    data : Tensor
+        data stage.
+
+    kernel : Tensor
+        kernel stage.
+
+    out : Tensor
+        output stage.
+
+    Returns
+    -------
+    hstride : int
+        stride size on height
+    wstride : int
+        stride size on width
+    """
+    _, _, IH, IW = data.shape
+    _, _, KH, KW = kernel.shape
+    _, _, OH, OW = out.shape
+    hstride = (IH - KH) // tvm.te.max(OH - 1, 1) + tvm.tirx.Select(OH == 1, 1, 0)
+    wstride = (IW - KW) // tvm.te.max(OW - 1, 1) + tvm.tirx.Select(OW == 1, 1, 0)
+    return get_const_int(hstride), get_const_int(wstride)
+
+
+def get_pad_tuple(padding, kernel):
+    """Common code to get the pad option
+
+    Parameters
+    ----------
+    padding : int or str
+        Padding size, or ['VALID', 'SAME']
+
+    kernel : tuple of int
+        Conv kernel size
+
+    Returns
+    -------
+    pad_top : int
+        Padding size on top
+
+    pad_left : int
+        Padding size on left
+
+    pad_down : int
+        Padding size on down.
+
+    pad_right : int
+        Padding size on right.
+    """
+    # compute the padding size
+    if isinstance(padding, tuple | list):
+        if len(padding) == 2:
+            pad_h = padding[0] * 2
+            pad_w = padding[1] * 2
+        elif len(padding) == 4:
+            return padding[0], padding[1], padding[2], padding[3]
+        else:
+            raise ValueError("Size of padding can only be 2 or 4")
+    elif isinstance(padding, int):
+        pad_h = pad_w = padding * 2
+    elif padding == "VALID":
+        pad_h = 0
+        pad_w = 0
+    elif padding == "SAME":
+        pad_h = kernel[0] - 1
+        pad_w = kernel[1] - 1
+    else:
+        raise ValueError(f"Unknown padding option {padding}")
+    pad_top = (pad_h + 1) // 2
+    pad_left = (pad_w + 1) // 2
+    return pad_top, pad_left, pad_h - pad_top, pad_w - pad_left
+
+
+def get_pad_tuple_generic(padding, kernel):
+    """Common code to get the pad option
+
+    Parameters
+    ----------
+    padding : int or str
+        Padding size, or ['VALID', 'SAME']
+
+    kernel : tuple of int
+        Conv kernel size
+
+    Returns
+    -------
+    pad_top : int
+        Padding size on top
+
+    pad_down : int
+        Padding size on down.
+
+    pad_left : int
+        Padding size on left
+
+    pad_right : int
+        Padding size on right.
+    """
+    # compute the padding size
+    if isinstance(padding, tuple | list):
+        if len(padding) == len(kernel):
+            pad_dimensions = [p * 2 for p in padding]
+        elif len(padding) == len(kernel) * 2:
+            return (
+                [padding[i] for i in range(len(kernel))],
+                [padding[len(kernel) + i] for i in range(len(kernel))],
+            )
+        else:
+            raise ValueError("Size of padding can only be len(kernel) or len(kernel) * 2")
+    elif isinstance(padding, int):
+        pad_dimensions = [padding * 2 for _ in range(len(kernel))]
+    elif padding == "VALID":
+        pad_dimensions = [0 for _ in range(len(kernel))]
+    elif padding == "SAME":
+        pad_dimensions = [k - 1 for k in kernel]
+    else:
+        raise ValueError(f"Unknown padding option {padding}")
+    pad_begin = [(p + 1) // 2 for p in pad_dimensions]
+    return [pad_begin, [pd - pb for pb, pd in zip(pad_begin, pad_dimensions)]]
+
+
+def get_pad_tuple3d(padding, kernel):
+    """Common code to get the pad option
+
+    Parameters
+    ----------
+    padding : int or str
+        Padding size, or ['VALID', 'SAME']
+
+    kernel : tuple of int
+        Conv kernel size
+
+    Returns
+    -------
+    pad_front : int
+        Padding size on front.
+
+    pad_top : int
+        Padding size on top
+
+    pad_left : int
+        Padding size on left
+
+    pad_back : int
+        Padding size on back.
+
+    pad_down : int
+        Padding size on down.
+
+    pad_right : int
+        Padding size on right.
+    """
+    # compute the padding size
+    if isinstance(padding, tuple | list):
+        if len(padding) == 3:
+            pad_d = padding[0] * 2
+            pad_h = padding[1] * 2
+            pad_w = padding[2] * 2
+        elif len(padding) == 6:
+            return padding[0], padding[1], padding[2], padding[3], padding[4], padding[5]
+        else:
+            raise ValueError("Size of padding can only be 3 or 6")
+    elif isinstance(padding, int):
+        pad_d = pad_w = pad_h = padding * 2
+    elif padding == "VALID":
+        pad_h = 0
+        pad_w = 0
+        pad_d = 0
+    elif padding == "SAME":
+        pad_d = kernel[0] - 1
+        pad_h = kernel[1] - 1
+        pad_w = kernel[2] - 1
+    else:
+        raise ValueError(f"Unknown padding option {padding}")
+    pad_top = (pad_h + 1) // 2
+    pad_left = (pad_w + 1) // 2
+    pad_front = (pad_d + 1) // 2
+    return pad_front, pad_top, pad_left, pad_d - pad_front, pad_h - pad_top, pad_w - pad_left
+
+
+def get_pad_tuple1d(padding, kernel):
+    """Common code to get the pad option
+
+    Parameters
+    ----------
+    padding : int or str
+        Padding size, or ['VALID', 'SAME']
+
+    kernel : tuple of int
+        Conv kernel size
+
+    Returns
+    -------
+    pad_left : int
+        Padding size on left
+
+    pad_right : int
+        Padding size on right.
+    """
+    # compute the padding size
+    if isinstance(padding, tuple | list):
+        if len(padding) == 1:
+            pad_w = padding[0] * 2
+        elif len(padding) == 2:
+            return padding[0], padding[1]
+        else:
+            raise ValueError("Size of padding can only be 2 or 4")
+    elif isinstance(padding, int):
+        pad_w = padding * 2
+    elif padding == "VALID":
+        pad_w = 0
+    elif padding == "SAME":
+        pad_w = kernel[0] - 1
+    else:
+        raise ValueError(f"Unknown padding option {padding}")
+    pad_left = (pad_w + 1) // 2
+    return pad_left, pad_w - pad_left

python/tvm/topi/nn/winograd_util.py

@@ -0,0 +1,181 @@
+# 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.
+#
+# ruff: noqa: E731
+"""Utility functions for implementing Winograd convolutions
+[*] Fast Algorithms for Convolutional Neural Networks
+    Andrew Lavin, Scott Gray
+    https://arxiv.org/abs/1509.09308
+    https://github.com/andravin/wincnn
+"""
+
+from functools import reduce
+from operator import mul
+
+import numpy as np
+
+from tvm.contrib.pickle_memoize import memoize
+
+from ..utils import const_matrix
+
+
+# pylint: disable=invalid-name
+def _cook_toom_convolution(a, n, r):
+    """Compute Cook-Toom convolution A,B,G matrices"""
+
+    def _F_m(a, n):
+        f = lambda j, i: reduce(mul, ((a[i] - a[k] if k != i else 1) for k in range(0, n - 1)), 1)
+        F = np.fromfunction(np.vectorize(f), (1, n - 1), dtype=int)
+        F = np.diagflat(F)
+        F = np.append(F, np.zeros((n - 1, 1), dtype=int), axis=1)
+        f = lambda i, j: 1 if j == (n - 1) else 0
+        z = np.fromfunction(np.vectorize(f), (1, n), dtype=int)
+
+        return np.append(F, z, axis=0)
+
+    def _A_m(a, m, n):
+        f = lambda i, j: a[i] ** j
+        A = np.fromfunction(np.vectorize(f), (m - 1, n), dtype=int)
+        f = lambda i, j: 1 if j == (n - 1) else 0
+        z = np.fromfunction(np.vectorize(f), (1, n), dtype=int)
+
+        return np.append(A, z, axis=0)
+
+    def _B_m(a, n):
+        f = lambda j, i: reduce(mul, ((a[i] - a[k] if k != i else 1) for k in range(0, n - 1)), 1)
+        Ff = np.fromfunction(np.vectorize(f), (1, n - 1), dtype=int)
+        f = lambda i, nth: (
+            (
+                reduce(mul, [(np.poly1d([1, -a[k]]) if k != i else 1) for k in range(0, n - 1)], 1)
+            ).coef[n - 1 - nth - 1]
+            / Ff[0, i]
+        )
+        F = np.fromfunction(np.vectorize(f), (n - 1, n - 1), dtype=int)
+        f = lambda i, j: -(a[i] ** (n - 1))
+        t = np.fromfunction(np.vectorize(f), (n - 1, 1), dtype=int)
+        T = np.append(np.eye(n - 1), t, axis=1)
+
+        return np.append(F.T.dot(T), np.array([np.eye(n)[n - 1]]), axis=0)
+
+    alpha = n + r - 1
+
+    f = _F_m(a, alpha)
+
+    if f[0, 0] < 0:
+        f[0, :] *= -1
+
+    A = _A_m(a, alpha, n)
+
+    G = _A_m(a, alpha, r).T
+    G = G.dot(np.linalg.inv(f)).T
+
+    B = _B_m(a, alpha)
+    B = B.dot(f.T)
+
+    return (A, B, G)
+
+
+def _interpolation_points(degree):
+    """Propose filter points"""
+
+    assert 2 < degree < 18
+
+    # Default interpolation lookup table
+    #
+    # [1] Error Analysis and Improving the Accuracy of Winograd Convolution for Deep Neural Networks
+    #     Barbara Barabasz, Andrew Anderson, Kirk M. Soodhalter, David Gregg
+    #     https://arxiv.org/abs/1803.10986
+    #
+
+    # pylint: disable=bad-whitespace,line-too-long
+    in_pts = [
+        #   {invalid}
+        [],
+        # 01 {E=4.63E-08 on conv2d  [1]}
+        [],
+        # 02 {E=7.65E-08 on F( 2,3) [1]}
+        [0, -1, 1],
+        # 03 {E=2.35E-07 on F( 3,3) [1]}
+        [0, -1, 1, 1 / 2],
+        # 04 {E=3.29E-07 on F( 4,3) [1]}
+        [0, -1, 1, 1 / 2, -2],
+        # 05 {E=6.81E-07 on F( 5,3) [1]}
+        [0, -1, 1, 1 / 2, -2, -1 / 2],
+        # 06 {E=8.79E-07 on F( 6,3) [1]}
+        [0, -1, 1, 1 / 2, -1 / 2, 2, -2],
+        # 07 {E=3.71E-06 on F( 7,3) [1]}
+        [0, -1, 1, 1 / 2, -1 / 2, 2, -2, -1 / 4],
+        # 08 {E=7.35E-06 on F( 8,3) [1]}
+        [0, -1, 1, 1 / 2, -1 / 2, 2, -2, -1 / 4, 4],
+        # 09 {E=2.20E-05 on F( 9,3) [1]}
+        [0, -1, 1, 1 / 2, -1 / 2, 2, -2, -1 / 4, 3 / 4, -4 / 3],
+        # 10 {E=3.22E-05 on F(10,3) [1]}
+        [0, -1, 1, 1 / 2, -1 / 2, 2, -2, -1 / 4, 4, 3 / 4, -4 / 3],
+        # 11 {E=1.09E-04 on F(11,3) [1]}
+        [0, -1, 1, 1 / 2, -1 / 2, 2, -2, -1 / 4, 4, 3 / 4, -4 / 3, 1 / 4],
+        # 12 {E=1.99E-04 on F(12,3) [1]}
+        [0, -1, 1, 1 / 2, -1 / 2, 2, -2, -1 / 4, 4, 1 / 4, -3 / 4, 4 / 3, -4],
+        # 13 {E=5.54E-04 on F(13,3) [1]}
+        [0, -1, 1, 1 / 2, -1 / 2, 2, -2, -1 / 4, 4, 1 / 4, -3 / 4, 4 / 3, 3 / 4, -4 / 3],
+        # 14 {E=8.80E-04 on F(14,3) [1]}
+        [0, -1, 1, 1 / 2, -1 / 2, 2, -2, -1 / 4, 4, 1 / 4, -3 / 4, 4 / 3, -4, 3 / 4, -4 / 3],
+        # 15 {E=1.07E-02 on F(15,3) [1]}
+        [0, -1, 1, 1 / 2, -1 / 2, 2, -2, -1 / 4, 4, 1 / 4, -3 / 4, 4 / 3, -4, 2 / 3, -3 / 2, 3 / 2],
+        # 16 {E=1.93E-02 on F(16,3) [1]}
+        [
+            0,
+            -1,
+            1,
+            1 / 2,
+            -1 / 2,
+            2,
+            -2,
+            -1 / 4,
+            4,
+            1 / 4,
+            -3 / 4,
+            4 / 3,
+            -4,
+            2 / 3,
+            -3 / 2,
+            -2 / 3,
+            3 / 2,
+        ],
+    ]  # pylint: enable=bad-whitespace,line-too-long
+
+    return np.array(in_pts[degree - 1], dtype=np.float64)
+
+
+@memoize("topi.nn.winograd_matrices", save_at_exit=False)
+def winograd_transform_matrices(tile_size, kernel_size, out_dtype):
+    """Compute the A, B, and G transform matrices for `tile_size` as a `tvm.Expr`."""
+    if not 1 < tile_size < 9:
+        raise ValueError(f"Unsupported tile size for Winograd: {tile_size}")
+    if not 2 < kernel_size < 8:
+        raise ValueError(f"Unsupported kernel size for Winograd: {kernel_size}")
+
+    degree = tile_size + kernel_size - 2
+
+    intp_pts = _interpolation_points(degree)
+    A_data, B_data, G_data = _cook_toom_convolution(intp_pts, tile_size, kernel_size)
+
+    out_dtype = "uint16" if out_dtype == "bfloat16" else out_dtype
+    return (
+        const_matrix(A_data.astype(out_dtype), "A"),
+        const_matrix(B_data.astype(out_dtype), "B"),
+        const_matrix(G_data.astype(out_dtype), "G"),
+    )

python/tvm/topi/reduction.py

@@ -0,0 +1,281 @@
+# 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,consider-using-enumerate,no-member
+"""Reduce operators"""
+
+from . import cpp
+
+
+def _get_real_axis(ndim, axis):
+    if axis is None:
+        real_axis = list(range(ndim))
+    else:
+        if isinstance(axis, int):
+            axis = [axis]
+        else:
+            assert isinstance(axis, list | tuple)
+        real_axis = []
+        for ele in axis:
+            if ele < 0:
+                ele += ndim
+            if ele >= ndim:
+                raise ValueError(
+                    f"{ele} exceeds the maximum dimension {ndim}. Received axis={axis}"
+                )
+            real_axis.append(ele)
+        real_axis.sort()
+        real_axis = list(set(real_axis))  # Remove the duplicates
+    return real_axis
+
+
+def sum(data, axis=None, keepdims=False):
+    """Sum of array elements over a given axis or a list of axes
+
+    Parameters
+    ----------
+    data : tvm.te.Tensor
+        The input tvm tensor
+
+    axis : None or int or tuple of int
+        Axis or axes along which a sum is performed.
+        The default, axis=None, will sum all of the elements of the input array.
+        If axis is negative it counts from the last to the first axis.
+
+    keepdims : bool
+        If this is set to True, the axes which are reduced are left in the result as dimensions
+        with size one.
+        With this option, the result will broadcast correctly against the input array.
+
+    Returns
+    -------
+    ret : tvm.te.Tensor
+    """
+    return cpp.sum(data, axis, keepdims)
+
+
+def all(data, axis=None, keepdims=False):
+    """Logical AND of array elements over a given axis or a list of axes
+
+    Parameters
+    ----------
+    data : tvm.te.Tensor
+        The input tvm boolean tensor
+
+    axis : None or int or tuple of int
+        Axis or axes along which a logical AND is performed.
+        The default, axis=None, will perform logical AND over all elements of the input array.
+        If axis is negative it counts from the last to the first axis.
+
+    keepdims : bool
+        If this is set to True, the axes which are reduced are left in the result as dimensions
+        with size one.
+        With this option, the result will broadcast correctly against the input array.
+
+    Returns
+    -------
+    ret : tvm.te.Tensor
+    """
+    return cpp.all(data, axis, keepdims)
+
+
+def any(data, axis=None, keepdims=False):
+    """Logical OR of array elements over a given axis or a list of axes
+
+    Parameters
+    ----------
+    data : tvm.te.Tensor
+        The input tvm boolean tensor
+
+    axis : None or int or tuple of int
+        Axis or axes along which a logical OR is performed.
+        The default, axis=None, will perform logical OR over all elements of the input array.
+        If axis is negative it counts from the last to the first axis.
+
+    keepdims : bool
+        If this is set to True, the axes which are reduced are left in the result as dimensions
+        with size one.
+        With this option, the result will broadcast correctly against the input array.
+
+    Returns
+    -------
+    ret : tvm.te.Tensor
+    """
+    return cpp.any(data, axis, keepdims)
+
+
+def max(data, axis=None, keepdims=False):
+    """Maximum of array elements over a given axis or a list of axes
+
+    Parameters
+    ----------
+    data : tvm.te.Tensor
+        The input tvm tensor
+
+    axis : None or int or tuple of int
+        Axis or axes along which the max operation is performed.
+        The default, axis=None, will find the max element from all of the elements of the input
+        array. If axis is negative it counts from the last to the first axis.
+
+    keepdims : bool
+        If this is set to True, the axes which are reduced are left in the result as dimensions
+        with size one.
+        With this option, the result will broadcast correctly against the input array.
+
+    Returns
+    -------
+    ret : tvm.te.Tensor
+    """
+    return cpp.max(data, axis, keepdims)
+
+
+def min(data, axis=None, keepdims=False):
+    """Minimum of array elements over a given axis or a list of axes
+
+    Parameters
+    ----------
+    data : tvm.te.Tensor
+        The input tvm tensor
+
+    axis : None or int or tuple of int
+        Axis or axes along which a minimum operation is performed.
+        The default, axis=None, will find the minimum element from all of the elements of the
+        input array. If axis is negative it counts from the last to the first axis.
+
+    keepdims : bool
+        If this is set to True, the axes which are reduced are left in the result as dimensions
+        with size one.
+        With this option, the result will broadcast correctly against the input array.
+
+    Returns
+    -------
+    ret : tvm.te.Tensor
+    """
+    return cpp.min(data, axis, keepdims)
+
+
+def argmax(data, axis=None, keepdims=False, select_last_index=False):
+    """Returns the indices of the maximum values along an axis.
+
+    Parameters
+    ----------
+    data : tvm.te.Tensor
+        The input tvm tensor
+
+    axis : None or int or tuple of int
+        Axis or axes along which a argmax operation is performed.
+        The default, axis=None, will find the indices of the maximum element of the elements of
+        the input array. If axis is negative it counts from the last to the first axis.
+
+    keepdims : bool
+        If this is set to True, the axes which are reduced are left in the result as dimensions
+        with size one.
+        With this option, the result will broadcast correctly against the input array.
+
+    select_last_index: bool
+        Whether to select the last index if the maximum element appears multiple times, else
+        select the first index.
+
+    Returns
+    -------
+    ret : tvm.te.Tensor
+    """
+    return cpp.argmax(data, axis, keepdims, select_last_index)
+
+
+def argmin(data, axis=None, keepdims=False, select_last_index=False):
+    """Returns the indices of the minimum values along an axis.
+
+    Parameters
+    ----------
+    data : tvm.te.Tensor
+        The input tvm tensor
+
+    axis : None or int or tuple of int
+        Axis or axes along which a argmin operation is performed.
+        The default, axis=None, will find the indices of minimum element all of the elements of
+        the input array. If axis is negative it counts from the last to the first axis.
+
+    keepdims : bool
+        If this is set to True, the axes which are reduced are left in the result as dimensions
+        with size one.
+        With this option, the result will broadcast correctly against the input array.
+
+    select_last_index: bool
+        Whether to select the last index if the minimum element appears multiple times, else
+        select the first index.
+
+    Returns
+    -------
+    ret : tvm.te.Tensor
+    """
+    return cpp.argmin(data, axis, keepdims, select_last_index)
+
+
+def prod(data, axis=None, keepdims=False):
+    """Product of array elements over a given axis or a list of axes
+
+    Parameters
+    ----------
+    data : tvm.te.Tensor
+        The input tvm tensor
+
+    axis : None or int or tuple of int
+        Axis or axes along which a prod operation is performed.
+        The default, axis=None, will get the prod element over all of the elements of the
+        input array. If axis is negative it counts from the last to the first axis.
+
+    keepdims : bool
+        If this is set to True, the axes which are reduced are left in the result as dimensions
+        with size one.
+        With this option, the result will broadcast correctly against the input array.
+
+    Returns
+    -------
+    ret : tvm.te.Tensor
+    """
+    return cpp.prod(data, axis, keepdims)
+
+
+def collapse_sum(data, target_shape):
+    """Return a summation of data to the given shape.
+
+    collapse_sum is intended as the backward operator of topi broadcast operators in the automatic
+    differentiation process.
+
+    We expect that data is the result of broadcasting some tensor of target_shape in some
+    broadcast operation. Thus target_shape and data.shape must follow broadcast rules.
+
+    During computation, the axes of data.shape and target_shape are checked from right to left.
+    For every axis, if it either:
+    - exist in data but not in target_shape, or
+    - is larger than 1 in data and equals to 1 in target_shape,
+    data will be summed over this axis.
+
+    Parameters
+    ----------
+    data : tvm.te.Tensor
+        The input tensor.
+
+    shape : Tuple[int]
+        The shape to collapse to.
+
+    Returns
+    -------
+    ret : tvm.te.Tensor
+        The result tensor after summation.
+    """
+    return cpp.collapse_sum(data, target_shape)

python/tvm/topi/scan.py

@@ -0,0 +1,240 @@
+# 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
+"""Scan (cumulative binary) operators"""
+
+import operator
+from collections.abc import Callable
+
+import tvm
+from tvm.script.ir_builder import IRBuilder
+from tvm.script.ir_builder import tirx as T
+
+from ..te import extern
+from ..tirx import decl_buffer
+from . import utils
+from .math import cast
+
+
+def scanop(
+    data: tvm.te.Tensor,
+    binop: Callable[["tvm.Expr", "tvm.Expr"], "tvm.Expr"],
+    identity_value: "tvm.Expr",
+    op_name: str,
+    axis: int | None = None,
+    dtype: str | None = None,
+    exclusive: bool | 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: tvm.Expr
+        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. The cumulative operation of zero elements
+        is assumed to be the identity_value.
+
+    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 dtype is None or dtype == "":
+        dtype = data.dtype
+
+    if exclusive is None:
+        exclusive = False
+
+    def maybe_cast(x):
+        if dtype != data.dtype:
+            return cast(x, dtype)
+        return x
+
+    axis_mul_before = 1
+    axis_mul_after = 1
+
+    if axis is None:
+        axis = 0
+        cumsum_axis_len = utils.prod(data.shape)
+        shape = (cumsum_axis_len,)
+    else:
+        if not isinstance(axis, int):
+            axis = utils.get_const_int(axis)
+
+        shape = data.shape
+        cumsum_axis_len = shape[axis]
+
+        if axis < 0:
+            axis = len(shape) + axis
+
+        for i, value in enumerate(shape, 0):
+            if i < axis:
+                axis_mul_before *= value
+            elif i > axis:
+                axis_mul_after *= value
+
+    def gen_ir(data_buf, out_buf):
+        with IRBuilder() as ib:
+            data_buf = T.buffer_proxy(data_buf)
+            out_buf = T.buffer_proxy(out_buf)
+
+            with T.parallel(0, axis_mul_before * axis_mul_after) as fused:
+                i = fused // axis_mul_after
+                j = fused % axis_mul_after
+                base_idx = i * cumsum_axis_len * axis_mul_after + j
+                if exclusive:
+                    out_buf[base_idx] = cast(identity_value, dtype)
+                else:
+                    out_buf[base_idx] = maybe_cast(data_buf[base_idx])
+                with T.serial(0, cumsum_axis_len - 1) as _k:
+                    k = _k + 1
+                    cur_idx = base_idx + k * axis_mul_after
+                    prev_idx = base_idx + (k - 1) * axis_mul_after
+                    if exclusive:
+                        out_buf[cur_idx] = binop(out_buf[prev_idx], maybe_cast(data_buf[prev_idx]))
+                    else:
+                        out_buf[cur_idx] = binop(out_buf[prev_idx], maybe_cast(data_buf[cur_idx]))
+
+            return ib.get()
+
+    out_buf = decl_buffer(shape, dtype, "out_buf")
+
+    return extern(
+        [shape],
+        [data],
+        lambda ins, outs: gen_ir(ins[0], outs[0]),
+        dtype=dtype,
+        out_buffers=[out_buf],
+        name=op_name,
+        tag=op_name,
+    )
+
+
+def cumsum(
+    data: tvm.te.Tensor,
+    axis: int | None = None,
+    dtype: str | None = None,
+    exclusive: bool | 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.
+
+    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,
+        op_name="cumsum_generic",
+        axis=axis,
+        dtype=dtype,
+        exclusive=exclusive,
+    )
+
+
+def cumprod(
+    data: tvm.te.Tensor,
+    axis: int | None = None,
+    dtype: int | None = None,
+    exclusive: bool | None = None,
+) -> tvm.te.Tensor:
+    """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.
+
+    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,
+        op_name="cumprod_generic",
+        axis=axis,
+        dtype=dtype,
+        exclusive=exclusive,
+    )

python/tvm/topi/scatter.py

@@ -0,0 +1,165 @@
+# 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
+"""ScatterND operator"""
+
+from tvm import DataTypeCode, te, tirx  # hide redefinition of min and max
+from tvm.arith.analyzer import Analyzer
+from tvm.script.ir_builder import IRBuilder
+from tvm.script.ir_builder import tirx as T
+from tvm.tirx import expr
+
+
+def _verify_scatter_nd_inputs(data, indices, updates):
+    analyzer = Analyzer()
+    mdim = int(indices.shape[0])
+    assert mdim <= len(data.shape), (
+        f"The first dimension of the indices ({mdim}) must be less than or equal to "
+        f"the length of the shape of the output ({len(data.shape)})."
+    )
+    for i in range(len(indices.shape) - 1):
+        if isinstance(indices.shape[i + 1], expr.Var) or isinstance(updates.shape[i], expr.Var):
+            continue
+
+        assert analyzer.can_prove_equal(indices.shape[i + 1], updates.shape[i]), (
+            f"Dimension of indices[{i + 1}] ({indices.shape[i + 1]}) must equal dimension of "
+            f"updates[{i}] ({updates.shape[i]})."
+        )
+    for i in range(mdim, len(data.shape)):
+        data_ind = i - mdim + len(indices.shape) - 1
+        if isinstance(updates.shape[data_ind], expr.Var) or isinstance(data.shape[i], expr.Var):
+            continue
+        assert updates.shape[data_ind] == data.shape[i], (
+            f"Dimension of updates[{data_ind}] ({updates.shape[data_ind]}) must equal dimension "
+            f"of out_shape[{i}] ({data.shape[i]})."
+        )
+
+    assert indices.dtype.matches_code(DataTypeCode.INT, DataTypeCode.UINT), (
+        f"Indices must be a tensor of integers, but its elements are {indices.dtype}."
+    )
+
+
+def scatter_nd(data, indices, updates, mode):
+    """Scatter elements from a n-dimension array.
+
+    Given updates with shape (Y_0, ..., Y_{K-1}, X_M, ..., X_{N-1}), indices with shape
+    (M, Y_0, ..., Y_{K-1}), and output copied from data with shape (X_0, X_1, ..., X_{N-1}),
+    scatter_nd computes
+
+    .. code-block::
+
+        output[indices[0, y_0, ..., y_{K-1}],
+               ...,
+               indices[M-1, y_0, ..., y_{K-1}],
+               x_M,
+               ...,
+               x_{N-1}
+              ] = f(output[...], updates[y_0, ..., y_{K-1}, x_M, ..., x_{N-1}])
+
+    where the update function f is determinted by the mode.
+
+    Parameters
+    ----------
+    data : tvm.te.Tensor
+        The source array.
+
+    indices : tvm.te.Tensor
+        The indices of the values to extract.
+
+    updates : tvm.te.Tensor
+        The updates to apply at the Indices
+
+    mode : string
+        The update mode for the algorithm, either "update" or "add"
+        If update, the update values will replace the input data
+        If add, the update values will be added to the input data
+
+    Returns
+    -------
+    ret : tvm.te.Tensor
+    """
+    _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
+
+        with IRBuilder() as ib:
+            with T.seq_scope():
+                with T.serial(0, fused_shape) as i:
+                    out[i] = data[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[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.generic",
+        tag="scatter_nd.generic",
+    )

python/tvm/topi/scatter_elements.py

@@ -0,0 +1,174 @@
+# 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.
+"""ScatterElements operator"""
+
+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
+
+
+def scatter_elements(data, indices, updates, axis=0, reduction="update"):
+    """Scatter elements from updates to corresponding indices of copied data.
+
+    Data, indices, updates and output have the same shape.
+    Indices can not have duplicates (if idx1 != idx2, then indices[idx1] != indices[idx2])
+    if reduction == "update".
+
+    .. code-block::
+
+        output[indices[i][j]][j] = f(output[indices[i][j]][j], updates[i][j]) if axis = 0
+        output[i][indices[i][j]] = f(output[i][indices[i][j]], updates[i][j]) if axis = 1
+
+    where the update function f is determined by the reduction.
+    Five types of the function are supported: "update", "add", "mul", "min" and "max" (see below)
+
+    Parameters
+    ----------
+    data : tvm.te.Tensor
+        The source array.
+
+    indices : tvm.te.Tensor
+        The indices of the values to extract.
+
+    updates : tvm.te.Tensor
+        The updates to apply at the Indices
+
+    axis : optional, int
+        The axis to scatter on. It is zero by default.
+
+    reduction : optional, string
+        The update mode for the algorithm, either "update", "add", "mul", "min" or "max"
+        If update, the update values will replace the input data
+        If add, the update values will be added to the input data
+        If mul, the input data will be multiplied on the update values
+        If mean, the input data will be mean between the update values and the input data
+        If min, there is choice of minimal between the update values and the input data
+        If max, there is choice of maximal between the update values and the input data
+        It is "update" by default
+
+    Returns
+    -------
+    ret : tvm.te.Tensor
+    """
+    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
+
+    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)
+
+        # Copy initial input data to output
+        with IRBuilder() as ib:
+            with T.seq_scope():
+                with T.parallel(0, full_range) as i:
+                    out[i] = data[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[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.generic",
+        tag="scatter_elements.generic",
+    )

python/tvm/topi/searchsorted.py

@@ -0,0 +1,133 @@
+# 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
+"""searchsorted operator"""
+
+from tvm.script.ir_builder import IRBuilder
+from tvm.script.ir_builder import tirx as T
+
+from . import te, utils
+from .math import cast
+
+
+def binary_search(sequence_offset, search_range, sorted_sequence, value, right, out_dtype):
+    """Common IR generator for binary search used by CPU and GPU backends.
+
+    Must be called within an active IRBuilder context.
+
+    `sorted_sequence` is a N-D Buffer whose innermost dimension we want to search for `value`,
+    and `search_range` is the size of the innermost dimension. `sequence_offset` is
+    a 1-D linearlized offset specifying which of innermost sequences to search.
+
+    So the search for `value` is performed over
+    `sorted_sequence[sequence_offset:(sequence_offset + search_range)]`.
+    Note that we index N-D Buffer by 1-D linearlized indices.
+
+    """
+    lo_buf = T.decl_buffer([1], out_dtype, scope="local")
+    hi_buf = T.decl_buffer([1], out_dtype, scope="local")
+    lo = T.buffer_proxy(lo_buf)
+    hi = T.buffer_proxy(hi_buf)
+
+    lo[0] = cast(0, out_dtype)
+    hi[0] = cast(search_range, out_dtype)
+
+    # Reference: pytorch/aten/src/ATen/native/cuda/Bucketization.cu
+    def condition(current_val, target_val):
+        if right:
+            return current_val <= target_val
+        return current_val < target_val
+
+    with T.While(lo[0] < hi[0]):
+        mid = lo[0] + (hi[0] - lo[0] >> 1)
+        with T.If(condition(sorted_sequence[sequence_offset + mid], value)):
+            with T.Then():
+                lo[0] = mid + 1
+            with T.Else():
+                hi[0] = mid
+
+    return lo[0]
+
+
+def searchsorted(sorted_sequence, values, right=False, out_dtype="int64"):
+    """Find indices where elements should be inserted to maintain order.
+       If `sorted_sequence` is N-dimensional, the innermost dimension of
+       `values` are searched in the corresponding dimension of `sorted_sequence`.
+
+    Parameters
+    ----------
+    sorted_sequence : te.Tensor
+        N-D or 1-D Tensor, containing monotonically increasing sequence
+        on the innermost dimension.
+
+    values : te.Tensor
+        N-D Tensor containing the search values. When `sorted_sequence` is 1-D,
+        the shape of `values` can be arbitrary. Otherwise, ranks of `sorted_sequence`
+        and `values` must be the same, and outer N-1 axes must have the same size.
+
+    right : bool, optional
+        Controls which index is returned if a value lands exactly on one of sorted values. If
+        False, the index of the first suitable location found is given. If true, return the
+        last such index. If there is no suitable index, return either 0 or N (where N is the
+        size of the innermost dimension).
+
+    dtype : string, optional
+        The data type of the output indices.
+
+    Returns
+    -------
+    indices : te.Tensor
+        Tensor with same shape as values, representing the indices of
+        elements of `values` if they are inserted in `sorted_sequence`.
+    """
+
+    def ir(sorted_sequence, values, indices):
+        with IRBuilder() as ib:
+            sorted_sequence_shape = sorted_sequence.shape
+            values_shape = values.shape
+            num_search = utils.prod(values_shape)
+            search_range = sorted_sequence_shape[-1]
+
+            sorted_sequence = T.buffer_proxy(sorted_sequence)
+            values = T.buffer_proxy(values)
+            indices = T.buffer_proxy(indices)
+
+            with T.parallel(0, num_search) as i:
+                if len(sorted_sequence_shape) == 1:
+                    sequence_offset = 0
+                else:
+                    sequence_id = i // values_shape[-1]
+                    sequence_offset = sequence_id * search_range
+
+                indices[i] = binary_search(
+                    sequence_offset,
+                    search_range,
+                    sorted_sequence,
+                    values[i],
+                    right,
+                    out_dtype,
+                )
+
+            return ib.get()
+
+    return te.extern(
+        values.shape,
+        [sorted_sequence, values],
+        lambda ins, outs: ir(ins[0], ins[1], outs[0]),
+        name="searchsorted",
+        dtype=out_dtype,
+    )

python/tvm/topi/signal.py

@@ -0,0 +1,214 @@
+# 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-arguments, too-many-nested-blocks, unused-argument
+"""STFT operator"""
+
+from math import pi
+
+from tvm import te, tirx
+from tvm.script.ir_builder import IRBuilder
+from tvm.script.ir_builder import tirx as T
+
+
+def stft(
+    data,
+    n_fft,
+    hop_length,
+    win_length,
+    window,
+    normalized,
+    onesided,
+    output_shape,
+):
+    """
+    The STFT computes the Fourier transform of short overlapping windows of the input.
+    This gives frequency components of the signal as they change over time.
+    Parameters
+    ----------
+    data : te.Tensor
+        Either a 1-D tensor or a 2-D batch tensor.
+    n_fft : int
+        The size of Fourier transform
+    hop_length : int
+        The distance between neighboring sliding window frames
+    win_length : int
+        The size of window frame and STFT filter
+    window : te.Tensor
+        A 1-D tensor window frame
+    normalized : bool
+        Whether to return the normalized STFT results
+    onesided : bool
+        Whether to return onesided result or fill with conjugate symmetry
+    Returns
+    -------
+    output : te.Tensor
+        Tensor containing the STFT result
+    Examples
+    --------
+    .. code-block:: python
+
+        data = [1, 2, 3, 4, 5, 6]
+        window = [4, 3, 2]
+        [n_fft, hop_length, win_length, normalized, onesided] = [3, 3, 3, False, True]
+        topi.stft(data, n_fft, hop_length, win_length, window, normalized, onesided)
+        -> [[[15.0000,  0.0000], [34.0000,  0.0000]], [[ 4.5000,  0.8660], [ 1.0000, -1.7321]]]
+    """
+
+    def gen_ir(
+        data_ptr,
+        n_fft,
+        hop_length,
+        win_length,
+        window_ptr,
+        normalized,
+        onesided,
+        output_ptr,
+        loop_kind,
+    ):
+        col_loop = T.vectorized if loop_kind == "vectorize" else T.serial
+
+        with IRBuilder() as ib:
+            data = T.buffer_proxy(data_ptr)
+            window = T.buffer_proxy(window_ptr)
+            output = T.buffer_proxy(output_ptr)
+            # https://librosa.org/doc/0.7.2/_modules/librosa/core/spectrum.html#stft
+            with T.parallel(0, output_ptr.shape[0] * output_ptr.shape[1]) as batch_row:
+                with col_loop(0, output_ptr.shape[2]) as col:
+                    batch = tirx.floordiv(batch_row, output_ptr.shape[1])
+                    row = tirx.floormod(batch_row, output_ptr.shape[1])
+                    output[batch, row, col, 0] = tirx.Cast(data_ptr.dtype, 0)
+                    output[batch, row, col, 1] = tirx.Cast(data_ptr.dtype, 0)
+                    with T.serial(0, win_length) as wlen:
+                        output[batch, row, col, 0] += (
+                            window[wlen]
+                            * data[batch, col * hop_length + wlen]
+                            * tirx.cos(2 * pi * row * wlen / win_length)
+                        )
+                        output[batch, row, col, 1] -= (
+                            window[wlen]
+                            * data[batch, col * hop_length + wlen]
+                            * tirx.sin(2 * pi * row * wlen / win_length)
+                        )
+                    with T.If(normalized):
+                        with T.Then():
+                            output[batch, row, col, 0] /= tirx.sqrt(tirx.const(n_fft, "float32"))
+                            output[batch, row, col, 1] /= tirx.sqrt(tirx.const(n_fft, "float32"))
+
+            return ib.get()
+
+    output_buf = tirx.decl_buffer(output_shape, data.dtype, "output_buf", layout=None)
+    loop_kind = "vectorize"
+    if isinstance(output_shape[2], tirx.expr.Var):  # any_dim
+        loop_kind = "serial"
+
+    return te.extern(
+        output_shape,
+        [data, window],
+        lambda ins, outs: gen_ir(
+            ins[0], n_fft, hop_length, win_length, ins[1], normalized, onesided, outs[0], loop_kind
+        ),
+        dtype=[data.dtype],
+        out_buffers=[output_buf],
+        name="stft_cpu",
+        tag="stft_cpu",
+    )
+
+
+def dft(
+    re_data: te.Tensor,
+    im_data: te.Tensor,
+    inverse: tirx.IntImm,
+):
+    """
+    Computes the discrete Fourier transform of input (calculation along the last axis).
+    This gives frequency components of the signal as they change over time.
+
+    Parameters
+    ----------
+    re_data : te.Tensor
+        N-D tensor, real part of the input signal.
+
+    im_data : te.Tensor
+        N-D tensor, imaginary part of the input signal.
+        If the signal is real, then the values of this tensor are zeros.
+
+    inverse : bool
+        Whether to perform the inverse discrete fourier transform.
+
+    Returns
+    -------
+    re_output : te.Tensor
+        The Fourier Transform of the input (Real part).
+    im_output : te.Tensor
+        The Fourier Transform of the input (Imaginary part).
+    """
+
+    def gen_ir(
+        re_data_buf,
+        im_data_buf,
+        re_output_buf,
+        im_output_buf,
+    ):
+        with IRBuilder() as ib:
+            re_data_ptr = T.buffer_proxy(re_data_buf)
+            im_data_ptr = T.buffer_proxy(im_data_buf)
+            re_output_ptr = T.buffer_proxy(re_output_buf)
+            im_output_ptr = T.buffer_proxy(im_output_buf)
+
+            shape = re_data.shape
+            n_fft = shape[len(shape) - 1]
+            base_range = 1
+            for i in range(len(shape) - 1):
+                base_range *= shape[i]
+
+            sign = -1 if inverse else 1
+            factor = 1.0 / n_fft if inverse else 1.0
+
+            with T.parallel(0, base_range) as i:
+                base_idx = i * n_fft
+                with T.serial(0, n_fft) as n:
+                    n_idx = base_idx + n
+                    re_output_ptr[n_idx] = tirx.Cast(re_output_ptr.dtype, 0)
+                    im_output_ptr[n_idx] = tirx.Cast(im_output_ptr.dtype, 0)
+                    _w = sign * -2 * pi * n / n_fft
+                    with T.serial(0, n_fft) as k:
+                        k_idx = base_idx + k
+                        w = _w * k
+                        cos_w = tirx.Cast(re_output_ptr.dtype, tirx.cos(w))
+                        sin_w = tirx.Cast(re_output_ptr.dtype, tirx.sin(w))
+                        re_output_ptr[n_idx] += (
+                            re_data_ptr[k_idx] * cos_w - im_data_ptr[k_idx] * sin_w
+                        )
+                        im_output_ptr[n_idx] += (
+                            re_data_ptr[k_idx] * sin_w + im_data_ptr[k_idx] * cos_w
+                        )
+
+                    re_output_ptr[n_idx] *= tirx.Cast(re_output_ptr.dtype, factor)
+                    im_output_ptr[n_idx] *= tirx.Cast(im_output_ptr.dtype, factor)
+
+            return ib.get()
+
+    output_shape = [re_data.shape] * 2
+
+    return te.extern(
+        shape=output_shape,
+        inputs=[re_data, im_data],
+        fcompute=lambda ins, outs: gen_ir(ins[0], ins[1], outs[0], outs[1]),
+        dtype=[re_data.dtype, im_data.dtype],
+        name="dft_cpu",
+        tag="dft_cpu",
+    )

python/tvm/topi/slice_scatter.py

@@ -0,0 +1,76 @@
+# 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.
+"""SliceScatter operator"""
+
+from tvm import topi
+
+from . import utils
+
+
+def slice_scatter(input_tensor, src, start, end, step, axis):
+    """
+    Scatters a slice of src into input along the given axis (SSA form).
+
+    Args:
+        input_tensor (te.Tensor): The input tensor to scatter into.
+        src (te.Tensor): The source tensor to scatter from.
+        start (int): The starting index of the slice.
+        end (int): The ending index of the slice.
+        step (int): The step size of the slice.
+        axis (int): The axis to scatter along.
+
+    Returns:
+        list[te.Tensor]: A list containing the output tensor with the slice scattered.
+    """
+
+    dim_size_expr = input_tensor.shape[axis]  # Expression for dimension size
+    dim_size = utils.get_const_int(dim_size_expr)  # Dimension size (as constant int)
+
+    if start == 0 and end == dim_size and step == 1:
+        return topi.identity(src)
+
+    mask = topi.full((dim_size,), "bool", True)
+    idx = topi.arange(start=0, stop=dim_size, step=1, dtype="int64")
+
+    if start != 0:
+        mask = topi.logical_and(mask, topi.greater_equal(idx, start))
+
+    if end != dim_size:
+        mask = topi.logical_and(mask, topi.less(idx, end))
+
+    if step != 1:
+        step_mask = topi.equal(topi.floor_mod(idx - start, step), 0)
+        mask = topi.logical_and(mask, step_mask)
+
+    mask_shape_base = [1] * len(input_tensor.shape)
+    mask_shape_base[axis] = dim_size
+    mask_shape = tuple(mask_shape_base)
+
+    mask_reshaped = topi.reshape(mask, mask_shape)
+
+    idx_new_pre = idx - start + (step - 1)
+    idx_new_div = topi.floor_divide(idx_new_pre, step)
+    idx_new = topi.clip(idx_new_div, 0, dim_size - 1)
+
+    temp = topi.take(src, idx_new, axis=axis)
+
+    mask_shape_expanded_base = list(input_tensor.shape)
+    mask_shape_expanded = tuple(mask_shape_expanded_base)
+
+    mask_expanded = topi.broadcast_to(mask_reshaped, mask_shape_expanded)
+
+    output = topi.where(mask_expanded, temp, input_tensor)
+
+    return [output]

python/tvm/topi/sort.py

@@ -0,0 +1,219 @@
+# 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=too-many-arguments
+"""Argsort operator"""
+
+import tvm
+from tvm import te
+
+from .utils import get_const_tuple
+
+
+def sort(data, axis=-1, is_ascend=1):
+    """Performs sorting along the given axis and returns an array
+    in sorted order.
+
+    Parameters
+    ----------
+    data : tvm.te.Tensor
+        The input tensor.
+
+    axis : int, optional
+        Axis along which to sort the input tensor.
+        By default the flattened array is used.
+
+    is_ascend : boolean, optional
+        Whether to sort in ascending or descending order.
+
+    dtype : string, optional
+        DType of the output indices.
+
+    Returns
+    -------
+    out : tvm.te.Tensor
+        Sorted index tensor.
+
+    """
+    data_buf = tvm.tirx.decl_buffer(
+        data.shape, data.dtype, "data_buf", data_alignment=8, layout=None
+    )
+    out_buf = tvm.tirx.decl_buffer(data.shape, data.dtype, "out_buf", data_alignment=8, layout=None)
+    out = te.extern(
+        data.shape,
+        [data],
+        lambda ins, outs: tvm.tirx.call_packed(
+            "tvm.contrib.sort.sort", ins[0], outs[0], axis, is_ascend
+        ),
+        dtype=data.dtype,
+        in_buffers=[data_buf],
+        out_buffers=out_buf,
+        name="sort_cpu",
+        tag="sort_cpu",
+    )
+    return out
+
+
+def argsort(data, valid_count=None, axis=-1, is_ascend=1, dtype="float32"):
+    """Performs sorting along the given axis and returns an array
+    of indices having the same shape as an input array that index
+    data in sorted order.
+
+    Parameters
+    ----------
+    data : tvm.te.Tensor
+        The input tensor.
+
+    valid_count : tvm.te.Tensor, optional
+        1-D tensor for valid number of boxes.
+
+    axis : int, optional
+        Axis along which to sort the input tensor.
+        By default the flattened array is used.
+
+    is_ascend : boolean, optional
+        Whether to sort in ascending or descending order.
+
+    dtype : string, optional
+        DType of the output indices.
+
+    Returns
+    -------
+    out : tvm.te.Tensor
+        Sorted index tensor.
+
+    Example
+    --------
+    .. code-block:: python
+
+        # An example to use argsort
+        dshape = (1, 5, 6)
+        data = te.placeholder(dshape, name="data")
+        axis = 0
+        is_ascend = False
+        out = argsort(data, axis=axis, is_ascend=is_ascend)
+        np_data = np.random.uniform(dshape)
+        s = topi.generic.schedule_argsort(out)
+        f = tvm.compile(s, [data, out], "llvm")
+        dev = tvm.cpu()
+        tvm_data = tvm.runtime.tensor(np_data, dev)
+        tvm_out = tvm.runtime.tensor(np.zeros(dshape, dtype=data.dtype.dtype), dev)
+        f(tvm_data, tvm_out)
+    """
+    data_buf = tvm.tirx.decl_buffer(
+        data.shape, data.dtype, "data_buf", data_alignment=8, layout=None
+    )
+    if valid_count is not None:
+        valid_count_buf = tvm.tirx.decl_buffer(
+            valid_count.shape, valid_count.dtype, "valid_count_buf", data_alignment=4, layout=None
+        )
+        out_buf = tvm.tirx.decl_buffer(
+            data.shape, "int32", "out_buf", data_alignment=8, layout=None
+        )
+        out = te.extern(
+            data.shape,
+            [data, valid_count],
+            lambda ins, outs: tvm.tirx.call_packed(
+                "tvm.contrib.sort.argsort_nms", ins[0], ins[1], outs[0], axis, is_ascend
+            ),
+            dtype="int32",
+            in_buffers=[data_buf, valid_count_buf],
+            out_buffers=out_buf,
+            name="argsort_nms_cpu",
+            tag="argsort_nms_cpu",
+        )
+    else:
+        out_buf = tvm.tirx.decl_buffer(data.shape, dtype, "out_buf", data_alignment=8, layout=None)
+        out = te.extern(
+            data.shape,
+            [data],
+            lambda ins, outs: tvm.tirx.call_packed(
+                "tvm.contrib.sort.argsort", ins[0], outs[0], axis, is_ascend
+            ),
+            dtype=dtype,
+            in_buffers=[data_buf],
+            out_buffers=out_buf,
+            name="argsort_cpu",
+            tag="argsort_cpu",
+        )
+    return out
+
+
+def topk(data, k=1, axis=-1, ret_type="both", is_ascend=False, dtype="int64"):
+    """Get the top k elements in an input tensor along the given axis.
+
+    Parameters
+    ----------
+    data : tvm.te.Tensor
+        The input tensor.
+
+    k : int or tvm.te.Tensor, optional
+        Number of top elements to select. Return all elements if k < 1.
+
+    axis : int, optional
+        Axis long which to sort the input tensor.
+
+    ret_type: str, optional
+        The return type [both, values, indices].
+        "both": return both top k data and indices.
+        "values": return top k data only.
+        "indices": return top k indices only.
+
+    is_ascend : boolean, optional
+        Whether to sort in ascending or descending order.
+
+    dtype : string, optional
+        The data type of the indices output.
+
+    Returns
+    -------
+    out : tvm.te.Tensor or List[tvm.te.Tensor]
+        The computed result.
+    """
+    assert ret_type in ["both", "values", "indices"]
+    data_buf = tvm.tirx.decl_buffer(
+        data.shape, data.dtype, "data_buf", data_alignment=8, layout=None
+    )
+    out_shape = list(get_const_tuple(data.shape))
+    kvar = tvm.te.var("k")
+    if not isinstance(k, int):
+        out_shape[axis] = kvar
+    elif k >= 1:
+        out_shape[axis] = k
+    out_bufs = []
+    if ret_type in ["both", "values"]:
+        out_bufs.append(
+            tvm.tirx.decl_buffer(out_shape, data.dtype, "value_buf", data_alignment=8, layout=None)
+        )
+    if ret_type in ["both", "indices"]:
+        out_bufs.append(
+            tvm.tirx.decl_buffer(out_shape, dtype, "indices_buf", data_alignment=8, layout=None)
+        )
+    out_shapes = [out_shape] * len(out_bufs)
+
+    kv = kvar if not isinstance(k, int) else k
+    out = te.extern(
+        out_shapes,
+        [data],
+        lambda ins, outs: tvm.tirx.call_packed(
+            "tvm.contrib.sort.topk", ins[0], *outs, kv, axis, ret_type, is_ascend
+        ),
+        in_buffers=[data_buf],
+        out_buffers=out_bufs,
+        name="topk_cpu",
+        tag="topk_cpu",
+    )
+    return out

python/tvm/topi/sparse_reshape.py

@@ -0,0 +1,194 @@
+# 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-arguments, too-many-nested-blocks
+"""Sparse_Reshape operator"""
+
+from tvm.script.ir_builder import IRBuilder
+from tvm.script.ir_builder import tirx as T
+from tvm.te import div, extern, floordiv, floormod
+from tvm.tirx import Cast, decl_buffer
+
+
+def sparse_reshape(
+    sparse_indices,
+    prev_shape,
+    new_shape,
+    new_sparse_indices_shape,
+    new_shape_shape,
+):
+    """
+    Reshape a Sparse Tensor
+
+    Parameters
+    ----------
+    sparse_indices : te.Expr
+        A 2-D tensor[N, n_dim] of integers containing location of sparse values, where N is the
+        number of sparse values and n_dim is the number of dimensions of the dense_shape
+
+    prev_shape : te.Expr
+        A 1-D tensor containing the previous shape of the dense tensor
+
+    new_shape : te.Expr
+        A 1-D tensor containing the new shape of the dense tensor
+
+    Returns
+    -------
+    result: te.Expr
+        Output tensor.
+
+    Examples
+    --------
+    .. code-block:: python
+
+        sparse_indices = [[0, 0, 0],
+                            [0, 0, 1],
+                            [0, 1, 0],
+                            [1, 0, 0],
+                            [1, 2, 3]]
+        prev_shape = [2, 3, 4]
+        new_shape = [9, -1]
+        new_sparse_indices, new_shape = topi.sparse_reshape(
+            sparse_indices, prev_shape, new_shape)
+        new_sparse_indices = [[0, 0],
+                              [0, 1],
+                              [1, 2],
+                              [4, 2],
+                              [8, 1]]
+        new_shape = [9, 4]
+    """
+
+    def gen_ir(
+        sparse_indices_ptr,
+        prev_shape_ptr,
+        new_shape_ptr,
+        new_sparse_indices_ptr,
+        out_new_shape_ptr,
+    ):
+        with IRBuilder() as ib:
+            sparse_indices = T.buffer_proxy(sparse_indices_ptr)
+            prev_shape = T.buffer_proxy(prev_shape_ptr)
+
+            new_shape = T.buffer_proxy(new_shape_ptr)
+            out_new_shape = T.buffer_proxy(out_new_shape_ptr)
+            new_sparse_indices = T.buffer_proxy(new_sparse_indices_ptr)
+
+            prev_shape_size = prev_shape_ptr.shape[0]
+            new_shape_size = new_shape_ptr.shape[0]
+
+            multipliers_buf = T.alloc_buffer([prev_shape_size], new_shape_ptr.dtype, scope="local")
+            multipliers = T.buffer_proxy(multipliers_buf)
+            dividers_buf = T.alloc_buffer([new_shape_size], new_shape_ptr.dtype, scope="local")
+            dividers = T.buffer_proxy(dividers_buf)
+            flattened_indices_buf = T.alloc_buffer(
+                [sparse_indices_ptr.shape[0]], new_shape_ptr.dtype, scope="local"
+            )
+            flattened_indices = T.buffer_proxy(flattened_indices_buf)
+            total_ele_buf = T.alloc_buffer([1], new_shape_ptr.dtype, scope="local")
+            total_ele = T.buffer_proxy(total_ele_buf)
+            division_total_ele_buf = T.alloc_buffer([1], new_shape_ptr.dtype, scope="local")
+            division_total_ele = T.buffer_proxy(division_total_ele_buf)
+            equal_shape_buf = T.alloc_buffer([1], "bool", scope="local")
+            equal_shape = T.buffer_proxy(equal_shape_buf)
+
+            total_ele[0] = prev_shape[0]
+
+            # Cumulative Reverse Exclusive Multiply
+            multipliers[prev_shape_size - 1] = Cast(new_shape_ptr.dtype, 1)
+            with T.serial(0, prev_shape_size - 1) as i_:
+                i = i_ + 1
+                multipliers[prev_shape_size - 1 - i] = (
+                    prev_shape[prev_shape_size - i] * multipliers[prev_shape_size - i]
+                )
+                total_ele[0] *= prev_shape[prev_shape_size - i]
+
+            division_total_ele[0] = Cast(new_shape_ptr.dtype, 1)
+            with T.serial(0, new_shape_size) as i:
+                with T.If(new_shape[i] != -1):
+                    with T.Then():
+                        division_total_ele[0] *= new_shape[i]
+
+            # Compute true output shape (replace negative ones)
+            with T.serial(0, new_shape_size) as i:
+                with T.If(new_shape[i] == -1):
+                    with T.Then():
+                        out_new_shape[i] = Cast(
+                            new_shape_ptr.dtype, div(total_ele[0], division_total_ele[0])
+                        )
+                    with T.Else():
+                        out_new_shape[i] = new_shape[i]
+
+            # Check if prev_shape and new_shape are equal
+            equal_shape[0] = True
+            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[i] != out_new_shape[i]):
+                            with T.Then():
+                                equal_shape[0] = False
+                with T.Else():
+                    equal_shape[0] = False
+
+            # Return same inputs if shapes are equal
+            with T.If(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:
+                            new_sparse_indices[i, j] = sparse_indices[i, j]
+
+                # Else compute new_sparse_indices
+                with T.Else():
+                    dividers[new_shape_size - 1] = Cast(new_shape_ptr.dtype, 1)
+                    with T.serial(0, new_shape_size - 1) as i_:
+                        i = i_ + 1
+                        dividers[new_shape_size - 1 - i] = (
+                            dividers[new_shape_size - i] * out_new_shape[new_shape_size - i]
+                        )
+
+                    with T.parallel(0, sparse_indices_ptr.shape[0]) as i:
+                        flattened_indices[i] = Cast(new_shape_ptr.dtype, 0)
+                        with T.serial(0, sparse_indices_ptr.shape[1]) as j:
+                            flattened_indices[i] += sparse_indices[i, j] * multipliers[j]
+
+                    with T.parallel(0, new_sparse_indices_ptr.shape[0]) as i:
+                        current_element_buf = T.alloc_buffer(
+                            [1], new_shape_ptr.dtype, scope="local"
+                        )
+                        current_element = T.buffer_proxy(current_element_buf)
+                        current_element[0] = flattened_indices[i]
+
+                        with T.serial(0, new_sparse_indices_ptr.shape[1]) as j:
+                            new_sparse_indices[i, j] = Cast(
+                                sparse_indices_ptr.dtype,
+                                floordiv(current_element[0], dividers[j]),
+                            )
+                            current_element[0] = floormod(current_element[0], dividers[j])
+
+            return ib.get()
+
+    new_sparse_indices_buf = decl_buffer(
+        new_sparse_indices_shape, sparse_indices.dtype, "new_sparse_indices_buf"
+    )
+    new_shape_buf = decl_buffer(new_shape_shape, prev_shape.dtype, "new_shape_buf")
+
+    return extern(
+        [new_sparse_indices_shape, new_shape_shape],
+        [sparse_indices, prev_shape, new_shape],
+        lambda ins, outs: gen_ir(ins[0], ins[1], ins[2], outs[0], outs[1]),
+        out_buffers=[new_sparse_indices_buf, new_shape_buf],
+        name="sparse_reshape_cpu",
+        tag="sparse_reshape_cpu",
+    )

python/tvm/topi/tag.py

@@ -0,0 +1,86 @@
+# 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.
+"""Namespace of all tag system in tvm
+
+Each operator can be tagged by a tag, which indicate its type.
+
+Generic categories
+
+- tag.ELEMWISE="elemwise":
+   Elementwise operator, for example :code:`out[i, j] = input[i, j]`
+- tag.BROADCAST="broadcast":
+    Broadcasting operator, can always map output axis to the input in order.
+    for example :code:`out[i, ax1, j, ax2] = input[i, j]`.
+    Note that the axis need to be in order so transpose is not a bcast operator.
+    If an input of broadcast operator has same shape as output,
+    we can ensure that it is elementwise relation.
+- tag.INJECTIVE="injective":
+    Injective operator, can always injectively map output axis to a single input axis.
+    All injective operator can still be safely fused similar to ewise to reduction.
+
+- tag.COMM_REDUCE="comm_reduce":
+    Communicative reduction operator
+- If an op does not belong to these generic categories, it should have a special tag.
+
+Note
+----
+When we add a new topi operator, the op need to be tagged as generic as possible.
+We can also compose tags like "injective,pad" to give generic and specific information.
+When we use composed tags, we must always put generic tag in the first location.
+"""
+
+ELEMWISE = "elemwise"
+BROADCAST = "broadcast"
+INJECTIVE = "injective"
+COMM_REDUCE = "comm_reduce"
+COMM_REDUCE_IDX = "comm_reduce_idx"
+
+
+def is_broadcast(tag):
+    """Check if a tag is bcast
+
+    Parameters
+    ----------
+    tag : str
+        The input tag
+
+    Returns
+    -------
+    ret : bool
+        Whether a tag is broadcast
+    """
+    if tag in (ELEMWISE, BROADCAST):
+        return True
+    return tag.startswith(ELEMWISE) or tag.startswith(BROADCAST)
+
+
+def is_injective(tag):
+    """Check if a tag is injective
+
+    Parameters
+    ----------
+    tag : str
+        The input tag
+
+    Returns
+    -------
+    ret : bool
+        Whether a tag is injective
+    """
+    if tag in (ELEMWISE, BROADCAST, INJECTIVE):
+        return True
+    return tag.startswith(ELEMWISE) or tag.startswith(BROADCAST) or tag.startswith(INJECTIVE)

python/tvm/topi/tensor.py

@@ -0,0 +1,113 @@
+# 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,consider-using-enumerate,unused-argument,len-as-condition
+"""Elementwise operators"""
+
+import math as _math
+
+from tvm import te
+
+from . import cpp
+
+
+def elemwise_sum(xs):
+    """Perform element-wise sum on inputs
+
+    Parameters
+    ----------
+    xs : list of tvm.te.Tensor
+        Input arguments.
+
+    Returns
+    -------
+    y : tvm.te.Tensor
+        The result.
+    """
+    return cpp.elemwise_sum(xs)
+
+
+def full(shape, dtype, fill_value):
+    """Fill tensor with fill_value
+
+    Parameters
+    ----------
+    shape : tuple
+        Input tensor shape.
+    dtype : str
+        Data type
+    fill_value : float
+        Value to be filled
+
+    Returns
+    -------
+    y : tvm.te.Tensor
+        The result.
+    """
+
+    if isinstance(fill_value, int | float) and (_math.isinf(fill_value) or _math.isnan(fill_value)):
+        if not ("float" in dtype or "bfloat16" in dtype):
+            raise ValueError("Infinite and NaN require a floating-point dtype.")
+
+    return cpp.full(shape, dtype, fill_value)
+
+
+def full_like(x, fill_value):
+    """Construct a tensor with same shape as input tensor,
+       then fill tensor with fill_value.
+
+    Parameters
+    ----------
+    x : tvm.te.Tensor
+        Input argument.
+    fill_value : float
+        Value to be filled
+
+    Returns
+    -------
+    y : tvm.te.Tensor
+        The result.
+    """
+    return cpp.full_like(x, fill_value)
+
+
+def eye(n: int, m: int | None = None, k: int = 0, dtype: str = "float32") -> te.Tensor:
+    """Generate an identity matrix or a matrix with ones on the k-th diagonal.
+
+    Parameters
+    ----------
+    n : int
+        Number of rows
+    m : int, optional
+        Number of columns. If None, defaults to n.
+    k : int, optional
+        Index of the diagonal. 0 (default) refers to the main diagonal.
+        A positive value refers to an upper diagonal, and a negative value
+        to a lower diagonal.
+    dtype : str, optional
+        Data type of the returned array.
+
+    Returns
+    -------
+    y : tvm.te.Tensor
+        The result.
+    """
+    m = m if m is not None else n
+    return te.compute(
+        (n, m),
+        lambda i, j: te.if_then_else(i == j - k, te.const(1, dtype), te.const(0, dtype)),
+        name="eye",
+    )

python/tvm/topi/testing/__init__.py

@@ -0,0 +1,80 @@
+# 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.
+
+"""TOPI Testing Util functions.
+
+Used to verify the correctness of operators in TOPI .
+"""
+
+from .conv1d_ncw_python import conv1d_ncw_python, group_conv1d_ncw_python
+from .conv2d_hwcn_python import conv2d_hwcn_python
+from .conv2d_nchw_python import conv2d_nchw_python
+from .conv2d_nhwc_python import conv2d_nhwc_python
+from .conv3d_ncdhw_python import conv3d_ncdhw_python
+from .conv3d_ndhwc_python import conv3d_ndhwc_python
+from .conv3d_transpose_ncdhw_python import conv3d_transpose_ncdhw_python
+from .conv2d_transpose_python import conv2d_transpose_nchw_python, conv2d_transpose_nhwc_python
+from .conv1d_transpose_ncw_python import (
+    conv1d_transpose_ncw_python,
+    group_conv1d_transpose_ncw_python,
+)
+from .correlation_nchw_python import correlation_nchw_python
+from .deformable_conv2d_python import deformable_conv2d_nchw_python, deformable_conv2d_nhwc_python
+from .depthwise_conv2d_python import (
+    depthwise_conv2d_python_nchw,
+    depthwise_conv2d_python_nhwc,
+    depthwise_conv2d_python_nchwc,
+)
+from .dilate_python import dilate_python
+from .softmax_python import softmax_python, log_softmax_python
+from .resize_python import resize1d_python, resize2d_python, resize3d_python
+from .reorg_python import reorg_python
+from .roi_align_python import roi_align_nchw_python, roi_align_nhwc_python
+from .roi_pool_python import roi_pool_nchw_python
+from .instance_norm_python import instance_norm_python
+from .layer_norm_python import layer_norm_python
+from .group_norm_python import group_norm_python
+from .rms_norm_python import rms_norm_python
+from .lrn_python import lrn_python
+from .l2_normalize_python import l2_normalize_python
+from .gather_python import gather_python
+from .gather_nd_python import gather_nd_python
+from .get_valid_counts_python import get_valid_counts_python
+from .strided_slice_python import strided_slice_python, strided_set_python
+from .batch_matmul import batch_matmul
+from .batch_norm import batch_norm
+from .nms_python import non_max_suppression_python
+from .slice_axis_python import slice_axis_python
+from .sequence_mask_python import sequence_mask
+from .poolnd_python import poolnd_python
+from .pool_grad_python import pool_grad_nchw
+from .one_hot import one_hot
+from .depth_to_space import depth_to_space_python
+from .space_to_depth import space_to_depth_python
+from .crop_and_resize_python import crop_and_resize_python
+from .adaptive_pool_python import adaptive_pool
+from .grid_sample_python import affine_grid_python, grid_sample_python
+from .matrix_set_diag import matrix_set_diag
+from .space_to_batch_nd import space_to_batch_nd_python
+from .batch_to_space_nd import batch_to_space_nd_python
+from .nll_loss import nll_loss
+from .dense import dense
+from .searchsorted import searchsorted_ref
+from .conv2d_backcward_weight_python import conv2d_backward_weight_python
+from .lstm_python import lstm_python
+from .attention_python import attention_python

python/tvm/topi/testing/adaptive_pool_python.py

@@ -0,0 +1,130 @@
+# 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, unused-argument, unused-variable
+# ruff: noqa: E741, RUF005
+"""adaptive pool in python"""
+
+import numpy as np
+
+
+def _start_index(index, odim, idim):
+    return int(np.floor(index * idim / odim))
+
+
+def _end_index(index, odim, idim):
+    return int(np.ceil((index + 1) * idim / odim))
+
+
+def _pool1d(in_size, out_size, np_data, np_op):
+    out = np.zeros(out_size).astype(np_data.dtype)
+    ow = out_size[0]
+    for l in range(ow):
+        l_start = _start_index(l, ow, in_size[0])
+        l_end = _end_index(l, ow, in_size[0])
+        l_sl = slice(l_start, l_end)
+        out[l] = np_op(np_data[l_sl])
+    return out
+
+
+def _pool2d(in_size, out_size, np_data, np_op):
+    out = np.zeros(out_size).astype(np_data.dtype)
+    oh, ow = out_size
+    for k in range(oh):
+        k_start = _start_index(k, oh, in_size[0])
+        k_end = _end_index(k, oh, in_size[0])
+        k_sl = slice(k_start, k_end)
+        for l in range(ow):
+            l_start = _start_index(l, ow, in_size[1])
+            l_end = _end_index(l, ow, in_size[1])
+            l_sl = slice(l_start, l_end)
+            out[k, l] = np_op(np_data[k_sl, l_sl])
+    return out
+
+
+def _pool3d(in_size, out_size, np_data, np_op):
+    out = np.zeros(out_size).astype(np_data.dtype)
+    od, oh, ow = out_size
+    for m in range(od):
+        m_start = _start_index(m, od, in_size[0])
+        m_end = _end_index(m, od, in_size[0])
+        m_sl = slice(m_start, m_end)
+        for k in range(oh):
+            k_start = _start_index(k, oh, in_size[1])
+            k_end = _end_index(k, oh, in_size[1])
+            k_sl = slice(k_start, k_end)
+            for l in range(ow):
+                l_start = _start_index(l, ow, in_size[2])
+                l_end = _end_index(l, ow, in_size[2])
+                l_sl = slice(l_start, l_end)
+                out[m, k, l] = np_op(np_data[m_sl, k_sl, l_sl])
+    return out
+
+
+def adaptive_pool_channel_first(np_data, out_size, pool_op, np_op):
+    """The reference function for adaptive pool, channel first layout"""
+    ishape = np_data.shape
+    n, c = ishape[:2]
+    oshape = (n, c) + out_size
+    np_out = np.zeros(oshape).astype(np_data.dtype)
+
+    for i in range(n):
+        for j in range(c):
+            np_out[i, j] = pool_op(ishape[2:], out_size, np_data[i, j], np_op)
+
+    return np_out
+
+
+def adaptive_pool_channel_last(np_data, out_size, pool_op, np_op):
+    """The reference function for adaptive pool, channel last layout"""
+    ishape = np_data.shape
+    n, c = ishape[0], ishape[-1]
+    oshape = (n,) + out_size + (c,)
+    np_out = np.zeros(oshape).astype(np_data.dtype)
+
+    for i in range(n):
+        for j in range(c):
+            if len(out_size) == 1:
+                np_out[i, :, j] = pool_op(ishape[1:-1], out_size, np_data[i, :, j], np_op)
+            elif len(out_size) == 2:
+                np_out[i, :, :, j] = pool_op(ishape[1:-1], out_size, np_data[i, :, :, j], np_op)
+            else:
+                np_out[i, :, :, :, j] = pool_op(
+                    ishape[1:-1], out_size, np_data[i, :, :, :, j], np_op
+                )
+
+    return np_out
+
+
+def adaptive_pool(np_data, out_size, pool_type, layout):
+    """The reference function for adaptive pool, for 2d and 3d"""
+    if isinstance(out_size, int):
+        out_size = (out_size,)
+    if len(out_size) == 1:
+        pool_op = _pool1d
+    elif len(out_size) == 2:
+        pool_op = _pool2d
+    else:
+        assert len(out_size) == 3
+        pool_op = _pool3d
+
+    np_op = np.mean if pool_type == "avg" else np.max
+
+    if layout in ["NCW", "NCHW", "NCDHW"]:
+        return adaptive_pool_channel_first(np_data, out_size, pool_op, np_op)
+
+    assert layout in ["NWC", "NHWC", "NDHWC"]
+    return adaptive_pool_channel_last(np_data, out_size, pool_op, np_op)

python/tvm/topi/testing/attention_python.py

@@ -0,0 +1,125 @@
+# 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.
+
+"""Attention operator in python"""
+
+import numpy as np
+
+from .softmax_python import softmax_python
+
+
+def attention_python(
+    q: np.ndarray,
+    k: np.ndarray,
+    v: np.ndarray,
+    bias: np.ndarray | None,
+    qk_scale: float,
+    causal: str,
+    window_size: int | None = None,
+    layout: str = "BSNH",
+):  # pylint: disable=too-many-arguments, too-many-locals, invalid-name
+    """Attention operator in python
+
+    Parameters
+    ----------
+    q : np.ndarray
+        Query tensor with shape [batch, seq_length, num_heads, head_dim] in the layout specified by
+        `layout`.
+    k : np.ndarray
+        Key tensor with shape [batch, seq_length_kv, num_kv_heads, head_dim] in the layout specified
+        by `layout`.
+    v : np.ndarray
+        Value tensor with shape [batch, seq_length_kv, num_kv_heads, head_dim_v] in the layout
+        specified by `layout`.
+    bias : np.ndarray
+        Bias tensor with shape [batch, num_heads, seq_length, seq_length]
+    qk_scale : float
+        Scale factor for the query-key product.
+    causal : str
+        The type of causal mask to apply. Can be "none", "TopLeft", or "BottomRight".
+    window_size : Optional[int]
+        The window size for the causal mask.
+    layout : str
+        The layout of the input tensors, e.g. "BSNH" or "BNSH".
+
+    Returns
+    -------
+    np.ndarray
+        The output tensor with shape [batch, seq_length, num_heads, head_dim_v] in the layout
+        specified by `layout`.
+    """
+    assert layout in ["BSNH", "BNSH", "SBNH"]
+
+    dim_b = layout.find("B")
+    dim_s = layout.find("S")
+    dim_n = layout.find("N")
+    dim_h = layout.find("H")
+
+    q = q.transpose(dim_b, dim_n, dim_s, dim_h)  # b, n, s, h
+    k = k.transpose(dim_b, dim_n, dim_s, dim_h)  # b, n, s_kv, h
+    kt = k.transpose(0, 1, 3, 2)  # b, n, h, s_kv
+    v = v.transpose(dim_b, dim_n, dim_s, dim_h)
+
+    num_heads = q.shape[1]
+    num_kv_heads = k.shape[1]
+    s = q.shape[2]
+    s_kv = k.shape[2]
+
+    if num_heads != num_kv_heads:
+        assert num_heads % num_kv_heads == 0
+        factor = num_heads // num_kv_heads
+        kt = np.repeat(kt, factor, axis=1)
+        v = np.repeat(v, factor, axis=1)
+
+    if not qk_scale == "none":
+        score = q @ kt * qk_scale  # b, n, s, s_kv
+    else:
+        score = q @ kt / np.sqrt(q.shape[-1])  # b, n, s, s_kv
+    if bias is not None:
+        score = score + bias  # b, n, s, s_kv
+    if causal == "none":
+        attn = softmax_python(score, -1)
+    else:
+        if causal == "TopLeft":
+            offset = 0
+        elif causal == "BottomRight":
+            offset = abs(s - s_kv)
+        else:
+            raise ValueError(f"Unsupported causal type: {causal}")
+        score_masked = np.tril(score, k=offset)
+
+        if window_size:
+            score_masked = np.triu(
+                score_masked,
+                -window_size + 1,  # pylint: disable=invalid-unary-operand-type
+            )
+
+        score_masked_exp = np.tril(
+            np.exp(score_masked - np.max(score_masked, axis=-1, keepdims=True)), k=offset
+        )
+
+        if window_size:
+            score_masked_exp = np.triu(
+                score_masked_exp,
+                -window_size + 1,  # pylint: disable=invalid-unary-operand-type
+            )
+
+        score_masked_sum = np.sum(score_masked_exp, axis=-1, keepdims=True)
+        attn = np.divide(score_masked_exp, score_masked_sum)
+
+    out = attn @ v  # b, n, s, h_v
+    return out.transpose(*np.argsort([dim_b, dim_n, dim_s, dim_h]).tolist())

python/tvm/topi/testing/batch_matmul.py

@@ -0,0 +1,60 @@
+# 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
+"""Batch matmul in python"""
+
+import numpy as np
+
+
+def batch_matmul(x, y, out_dtype=None, trans_x=False, trans_y=True):
+    """batch_matmul operator implemented in numpy.
+
+    Parameters
+    ----------
+    x : numpy.ndarray
+        3-D with shape [batch, M, K]
+
+    y : numpy.ndarray
+        3-D with shape [batch, N, K]
+
+    out_dtype: string, optional
+        Specify the dtype of output
+
+    Returns
+    -------
+    out : numpy.ndarray
+        3-D with shape [batch, M, N]
+    """
+    if trans_x:
+        XB, _, M = x.shape
+    else:
+        XB, M, _ = x.shape
+    if trans_y:
+        YB, N, _ = y.shape
+    else:
+        YB, _, N = y.shape
+    batch = max(XB, YB)
+    dtype = x.dtype if out_dtype is None else out_dtype
+    out = np.zeros((batch, M, N)).astype(dtype)
+    for i in range(batch):
+        xx = x[i if XB != 1 else 0].astype(dtype)
+        yy = y[i if YB != 1 else 0].astype(dtype)
+        out[i] = np.dot(
+            xx.T if trans_x else xx,
+            yy.T if trans_y else yy,
+        )
+    return out

python/tvm/topi/testing/batch_norm.py

@@ -0,0 +1,115 @@
+# 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.
+"""Batch Normalization implemented in Numpy."""
+
+import numpy as np
+
+
+def batch_norm(
+    x: np.ndarray,
+    gamma: np.ndarray,
+    beta: np.ndarray,
+    moving_mean: np.ndarray,
+    moving_var: np.ndarray,
+    axis: int,
+    epsilon: float,
+    center: bool,
+    scale: bool,
+    training: bool,
+    momentum: float,
+):
+    """Batch Normalization operator implemented in Numpy.
+
+    Parameters
+    ----------
+    data : np.ndarray
+        Input to be batch-normalized.
+
+    gamma : np.ndarray
+        Scale factor to be applied to the normalized tensor.
+
+    beta : np.ndarray
+        Offset to be applied to the normalized tensor.
+
+    moving_mean : np.ndarray
+        Running mean of input.
+
+    moving_var : np.ndarray
+        Running variance of input.
+
+    axis : int
+        Specify along which shape axis the normalization should occur.
+
+    epsilon : float
+        Small float added to variance to avoid dividing by zero.
+
+    center : bool
+        If True, add offset of beta to normalized tensor, If False,
+        beta is ignored.
+
+    scale : bool
+        If True, scale normalized tensor by gamma. If False, gamma
+        is ignored.
+
+    training : bool
+        Indicating whether it is in training mode. If True, update
+        moving_mean and moving_var.
+
+    momentum : float
+        The value used for the moving_mean and moving_var update
+
+    Returns
+    -------
+    output : np.ndarray
+        Normalized data with same shape as input
+
+    moving_mean : np.ndarray
+        Running mean of input.
+
+    moving_var : np.ndarray
+        Running variance of input.
+    """
+    shape = [1] * len(x.shape)
+    shape[axis] = x.shape[axis]
+
+    if training:
+        reduce_axes = list(range(len(x.shape)))
+        reduce_axes.remove(axis)
+        reduce_axes = tuple(reduce_axes)
+        data_mean = np.mean(x, axis=reduce_axes)
+        data_var = np.var(x, axis=reduce_axes)
+        data_mean_rs = np.reshape(data_mean, shape)
+        data_var_rs = np.reshape(data_var, shape)
+        out = (x - data_mean_rs) / np.sqrt(data_var_rs + epsilon)
+    else:
+        moving_mean_rs = moving_mean.reshape(shape)
+        moving_var_rs = moving_var.reshape(shape)
+        out = (x - moving_mean_rs) / np.sqrt(moving_var_rs + epsilon)
+
+    if scale:
+        out = out * gamma.reshape(shape)
+    if center:
+        out = out + beta.reshape(shape)
+
+    if training:
+        return [
+            out,
+            (1 - momentum) * moving_mean + momentum * data_mean,
+            (1 - momentum) * moving_var + momentum * data_var,
+        ]
+
+    return [out, moving_mean, moving_var]

python/tvm/topi/testing/batch_to_space_nd.py

@@ -0,0 +1,99 @@
+# 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, line-too-long, unused-variable, too-many-locals
+"""Batch to space ND in python"""
+
+import numpy as np
+
+from . import strided_slice_python
+
+
+def batch_to_space_nd_python(data, block_shape, crop_begin_list, crop_end_list):
+    """Batch to Space operator in python for NHWC layout.
+
+    Parameters
+    ----------
+    data : np.ndarray
+        N-D with shape [batch, spatial_shape, remaining_shapes],
+        where spatial_shape has M dimensions.
+
+    block_shape : list of ints
+        1-D array of size [M] where M is number of spatial dims, specifies block
+        size for each spatial dimension.
+
+    crop_begin_list : list of ints
+        list of shape [M] where M is number of spatial dims, specifies
+        begin crop size for each spatial dimension.
+
+    crop_end_list : list of ints
+        list of shape [M] where M is number of spatial dims, specifies
+        end crop size for each spatial dimension.
+
+    Returns
+    -------
+    b2s_out : np.ndarray
+        N-D with shape
+        [batch / prod(block_shape),
+        in_shape[1] * block_shape[0] - crop_begin_list[0] - crop_end_list[0], ...,
+        in_shape[M] * block_shape[M-1] - crop_begin_list[M-1] - crop_end_list[M-1],
+        remaining_shape]
+    """
+    in_shape = data.shape
+    N = len(in_shape)
+    M = len(block_shape)
+    block_shape_prod = np.prod(block_shape)
+    in_batch = data.shape[0]
+    axis = []
+    r_p_shape = []
+
+    r_shape = [block_shape[i] for i in range(0, M)]
+    axis.append(len(r_shape))
+    r_shape.append(in_batch // block_shape_prod)
+
+    for i in range(1, N):
+        axis.append(len(r_shape))
+        if len(axis) < (M + N):
+            axis.append(len(r_shape) - (M + 1))
+        r_shape.append(in_shape[i])
+
+    r_p_shape.append(int(in_batch / block_shape_prod))
+    for i in range(1, M + 1):
+        r_p_shape.append(in_shape[i] * block_shape[i - 1])
+    for i in range(M + 1, N):
+        r_p_shape.append(in_shape[i])
+
+    b2s_out = np.reshape(data, newshape=r_shape)
+    b2s_out = np.transpose(b2s_out, axes=axis)
+    b2s_out = np.reshape(b2s_out, newshape=r_p_shape)
+
+    # Crop the start and end of dimensions of b2s_out
+    begin_idx = []
+    end_idx = []
+    strides = []
+
+    for i, _ in enumerate(r_p_shape):
+        strides.append(1)
+        if 0 < i <= M:
+            # begin and end index for spatial dimensions
+            begin_idx.append(crop_begin_list[i - 1])
+            end_idx.append(r_p_shape[i] - crop_end_list[i - 1])
+        else:
+            begin_idx.append(0)
+            end_idx.append(r_p_shape[i])
+
+    b2s_out = strided_slice_python(b2s_out, begin_idx, end_idx, strides)
+    return b2s_out

python/tvm/topi/testing/common.py

@@ -0,0 +1,70 @@
+# 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
+"""Common utility for topi test"""
+
+import numpy as np
+import scipy.signal
+
+
+def _convolve2d(data, weights):
+    """2d convolution operator in HW layout.
+
+    This is intended to be used as a replacement for
+    scipy.signals.convolve2d, with wider support for different dtypes.
+    scipy.signal.convolve2d does not support all TVM-supported
+    dtypes (e.g. float16).  Where possible, this function uses
+    scipy.signal.convolve2d to take advantage of compiled scipy
+    routines, falling back to an explicit loop only where needed.
+
+    Parameters
+    ----------
+    data : numpy.ndarray
+        2-D with shape [in_height, in_width]
+
+    weights : numpy.ndarray
+        2-D with shape [filter_height, filter_width].
+
+    Returns
+    -------
+    b_np : np.ndarray
+        2-D with shape [out_height, out_width]
+
+        Return value and layout conventions are matched to
+        ``scipy.signal.convolve2d(data, weights, mode="valid")``
+    """
+
+    try:
+        return scipy.signal.convolve2d(data, weights, mode="valid")
+    except ValueError:
+        pass
+
+    weights = np.rot90(weights, k=2)
+
+    assert len(data.shape) == len(weights.shape) == 2
+
+    dtype = data.dtype
+    kernel_h, kernel_w = weights.shape
+
+    output_shape = [a_dim - w_dim + 1 for a_dim, w_dim in zip(data.shape, weights.shape)]
+    output = np.zeros(output_shape, dtype=dtype)
+
+    for y in range(output_shape[0]):
+        for x in range(output_shape[1]):
+            output[y][x] = np.sum(data[y : y + kernel_h, x : x + kernel_w] * weights)
+
+    return output

python/tvm/topi/testing/conv1d_ncw_python.py

@@ -0,0 +1,110 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+# pylint: disable=unused-variable, invalid-name
+"""1D convolution in python"""
+
+import numpy as np
+
+from tvm.topi.nn.utils import get_pad_tuple1d
+
+
+def dilate_np(x, dilation):
+    """1D dilation using numpy
+
+    Parameters
+    ----------
+    x : numpy.ndarray
+        Array to dilate with shape [batch, in_channel, in_width]
+
+    dilation : int
+        dilation rate of output
+
+    Returns
+    -------
+    out : numpy.ndarray
+        Dilated output with shape [batch, in_channel, (in_width - 1) * dilation + 1]
+    """
+    irange = range(len(x) - 1)
+    for d in range(dilation - 1):
+        indices = [(d + 1) * (i + 1) for i in irange]
+        x = np.insert(x, indices, 0)
+    return x
+
+
+def group_conv1d_ncw_python(a_np, w_np, stride, padding, dilation, groups):
+    "Grouped version of `conv1d_ncw_python`, see that for documentation"
+    a_slices = np.array_split(a_np, groups, axis=1)
+    w_slices = np.array_split(w_np, groups, axis=0)
+    b_slices = [
+        conv1d_ncw_python(a_slice, w_slice, stride, padding, dilation)
+        for a_slice, w_slice in zip(a_slices, w_slices)
+    ]
+    return np.concatenate(b_slices, axis=1)
+
+
+def conv1d_ncw_python(a_np, w_np, stride, padding, dilation):
+    """1D convolution operator in NCW layout
+
+    Parameters
+    ----------
+    a_np : numpy.ndarray
+        3-D with shape [batch, in_channel, in_width]
+
+    w_np : numpy.ndarray
+        3-D with shape [num_filter, in_channel, filter_width]
+
+    stride : int
+        Stride size
+
+    padding : int, tuple, or str
+        Single int for padding size or tuple of (left, right) padding
+        or a string in ['VALID', 'SAME']
+
+    dilation : int
+        Dilation rate of the kernel
+
+    groups : int
+        Number of groups in the convolution
+
+    Returns
+    -------
+    b_np : numpy.ndarray
+        3-D with shape [batch, out_channel, out_width]
+    """
+    batch, in_c, in_w = a_np.shape
+    out_c, _, filter_w = w_np.shape
+    if isinstance(stride, tuple | list):
+        stride = stride[0]
+    if isinstance(dilation, tuple | list):
+        dilation = dilation[0]
+
+    dilated_filter_w = (filter_w - 1) * dilation + 1
+    pad_left, pad_right = get_pad_tuple1d(padding, (dilated_filter_w,))
+    out_w = ((in_w - dilated_filter_w + pad_left + pad_right) // stride) + 1
+
+    padded_a_np = np.zeros((batch, in_c, in_w + pad_left + pad_right))
+    padded_a_np[:, :, pad_left : (in_w + pad_left)] = a_np
+
+    b_np = np.zeros((batch, out_c, out_w))
+    for n in range(batch):
+        for f in range(out_c):
+            for c in range(in_c):
+                out = np.convolve(
+                    padded_a_np[n, c], np.flip(dilate_np(w_np[f, c], dilation)), mode="valid"
+                )
+                b_np[n, f] += out[::stride]
+    return b_np

python/tvm/topi/testing/conv1d_transpose_ncw_python.py

@@ -0,0 +1,92 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+# pylint: disable=unused-variable
+"""Transposed 1D convolution in python"""
+
+import numpy as np
+import scipy
+
+import tvm.topi.testing
+from tvm.topi.nn.utils import get_pad_tuple1d
+
+
+def group_conv1d_transpose_ncw_python(a_np, w_np, stride, padding, output_padding, groups=1):
+    "Grouped version of `conv1d_transpose_ncw_python`, see that for documentation"
+    a_slices = np.array_split(a_np, groups, axis=1)
+    w_slices = np.array_split(w_np, groups, axis=0)
+    b_slices = [
+        conv1d_transpose_ncw_python(a_slice, w_slice, stride, padding, output_padding)
+        for a_slice, w_slice in zip(a_slices, w_slices)
+    ]
+    b_np = np.concatenate(b_slices, axis=1)
+    return b_np
+
+
+def conv1d_transpose_ncw_python(a_np, w_np, stride, padding, output_padding):
+    """Transposed 1D convolution operator in NCW layout.
+
+    Parameters
+    ----------
+    a_np : numpy.ndarray
+        3-D with shape [batch, in_channel, in_width]
+
+    w_np : numpy.ndarray
+        3-D with shape [in_channel, num_filter, filter_width]
+
+    stride : int or a list/tuple of one int
+        Stride size, or [stride_width]
+
+    padding : int, tuple, or str
+        Single int for padding size, or
+        tuple of 2 ints for left and right padding, or
+        ['VALID', 'SAME']
+
+    output_padding : tuple
+        Used to recover the actual output shape in case more than one
+        is possible
+
+    Returns
+    -------
+    b_np : np.ndarray
+        3-D with shape [batch, out_channel, out_width]
+
+    """
+    batch, in_c, in_w = a_np.shape
+    _, out_c, filter_w = w_np.shape
+    opad = output_padding[0]
+    if isinstance(stride, int):
+        stride_w = stride
+    else:
+        stride_w = stride[0]
+    assert opad < stride_w
+    fpad_left, fpad_right = get_pad_tuple1d(padding, filter_w)
+    # dilate stage
+    dilated_a_np = tvm.topi.testing.dilate_python(a_np, [1, 1, stride_w])
+    # padding stage
+    bpad_left = filter_w - 1 - fpad_left
+    bpad_right = filter_w - 1 - fpad_right + opad
+    padded_a_np = np.zeros((batch, in_c, dilated_a_np.shape[2] + bpad_left + bpad_right))
+    padded_a_np[:, :, bpad_left : dilated_a_np.shape[2] + bpad_left] = dilated_a_np
+    # convolution stage
+    out_w = (in_w - 1) * stride_w - fpad_left - fpad_right + filter_w + opad
+    b_np = np.zeros((batch, out_c, out_w))
+    for n in range(batch):
+        for f in range(out_c):
+            for c in range(in_c):
+                out = scipy.signal.convolve(padded_a_np[n, c], w_np[c, f], mode="valid")
+                b_np[n, f] += out
+    return b_np

python/tvm/topi/testing/conv2d_backcward_weight_python.py

@@ -0,0 +1,150 @@
+# 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-nested-blocks
+"""Gradient of conv2d with respect to weight in python"""
+
+import numpy as np
+
+
+# Reference: cutlass/tools/util/include/cutlass/util/reference/host/convolution.h
+def conv2d_backward_weight_nchw_python(
+    dy_np, x_np, kernel_size, stride, padding, groups=1, channels=None
+):
+    """Gradient of the conv2d op with respect to weight, in NCHW layout.
+
+    Parameters
+    ----------
+    dy_np : numpy.ndarray
+        4-D with shape [batch, in_channel, out_height, out_width]
+
+    x_np : numpy.ndarray
+        4-D with shape [batch, in_channel, in_height, in_width]
+
+    kernel_size : tuple of two ints
+        Height and width of the weight
+
+    stride : tuple of two ints
+        Stride size, or [stride_height, stride_width]
+
+    padding : tuple of two ints
+        Spatial padding, or [pad_h, pad_w]
+
+    Returns
+    -------
+    dw_np : np.ndarray
+        4-D with shape [num_filter, in_channel, filter_height, filter_width]
+
+    """
+    N, C, H, W = x_np.shape
+    _, K, P, Q = dy_np.shape
+    R, S = kernel_size
+    pad_h, pad_w = padding
+    stride_h, stride_w = stride
+    is_depth_wise = C == K and C == groups
+
+    if is_depth_wise:
+        assert channels == groups, "Only channel_mult == 1 supported for now."
+        dw = np.zeros((K, 1, R, S)).astype(dy_np.dtype)
+    else:
+        assert groups == 1, "General grouped conv2d not supported for now."
+        dw = np.zeros((K, C, R, S)).astype(dy_np.dtype)
+
+    for k in range(K):
+        for r in range(R):
+            for s in range(S):
+                for c in range(dw.shape[1]):
+                    acc = 0
+                    for n in range(N):
+                        for p in range(P):
+                            for q in range(Q):
+                                if not is_depth_wise:
+                                    in_c = c
+                                else:
+                                    in_c = k
+
+                                coord = (
+                                    n,
+                                    in_c,
+                                    p * stride_h - pad_h + r,
+                                    q * stride_w - pad_w + s,
+                                )
+
+                                if (
+                                    coord[2] < H
+                                    and coord[2] >= 0
+                                    and coord[3] < W
+                                    and coord[3] >= 0
+                                ):
+                                    acc += dy_np[n, k, p, q] * x_np[coord]
+
+                    dw[k, c, r, s] = acc
+
+    return dw
+
+
+def conv2d_backward_weight_python(
+    dy_np, x_np, kernel_size, stride, padding, layout="NCHW", groups=1, channels=None
+):
+    """Gradient of the conv2d op with respect to weight, in NCHW or NHWC layout.
+
+    Parameters
+    ----------
+    dy_np : numpy.ndarray
+        4-D with shape [batch, in_channel, out_height, out_width] for NCHW layout
+
+    x_np : numpy.ndarray
+        4-D with shape [batch, in_channel, in_height, in_width] for NCHW layout
+
+    kernel_size : tuple of two ints
+        Height and width of the weight
+
+    stride : tuple of two ints
+        Stride size, or [stride_height, stride_width]
+
+    padding : tuple of two ints
+        Spatial padding, or [pad_h, pad_w]
+
+    layout: string
+        Layout of dy_np and x_np
+
+    groups: int
+        Number of groups for grouped convolution.
+
+    channels : int
+        Number of output channels of this convolution.
+
+    Returns
+    -------
+    dw_np : np.ndarray
+        Tensor of shape [num_filter, in_channel, filter_height, filter_width] for NCHW layout,
+        [num_filter, filter_height, filter_width, in_channel] for NHWC layout.
+    """
+    if layout == "NCHW":
+        return conv2d_backward_weight_nchw_python(
+            dy_np, x_np, kernel_size, stride, padding, groups, channels
+        )
+
+    dw_np_oihw = conv2d_backward_weight_nchw_python(
+        np.transpose(dy_np, [0, 3, 1, 2]),
+        np.transpose(x_np, [0, 3, 1, 2]),
+        kernel_size,
+        stride,
+        padding,
+        groups,
+        channels,
+    )
+    return np.transpose(dw_np_oihw, [0, 2, 3, 1])

python/tvm/topi/testing/conv2d_hwcn_python.py

@@ -0,0 +1,79 @@
+# 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, line-too-long, unused-variable, too-many-locals
+"""Convolution in python"""
+
+import numpy as np
+import scipy.signal
+
+from tvm.topi.nn.utils import get_pad_tuple
+
+
+def conv2d_hwcn_python(a_np, w_np, stride, padding):
+    """Convolution operator in HWCN layout.
+
+    Parameters
+    ----------
+    a_np : numpy.ndarray
+        4-D with shape [in_height, in_width, in_channel, batch]
+
+    w_np : numpy.ndarray
+        4-D with shape [filter_height, filter_width, in_channel, num_filter]
+
+    stride : int or a list/tuple of two ints
+        Stride size, or [stride_height, stride_width]
+
+    padding : int or str or a list/tuple of 2 or 4 ints
+        Padding size, or ['VALID', 'SAME'], or
+        [pad_height, pad_width] for 2 ints, or
+        [pad_top, pad_left, pad_bottom, pad_right] for 2 ints
+
+    Returns
+    -------
+    b_np : np.ndarray
+        4-D with shape [out_height, out_width, out_channel, batch]
+    """
+    in_height, in_width, in_channel, batch = a_np.shape
+    kernel_h, kernel_w, _, num_filter = w_np.shape
+    if isinstance(stride, int):
+        stride_h = stride_w = stride
+    else:
+        stride_h, stride_w = stride
+
+    pad_top, pad_left, pad_bottom, pad_right = get_pad_tuple(padding, (kernel_h, kernel_w))
+    pad_h = pad_top + pad_bottom
+    pad_w = pad_left + pad_right
+    # compute the output shape
+    out_channel = num_filter
+    out_height = (in_height - kernel_h + pad_h) // stride_h + 1
+    out_width = (in_width - kernel_w + pad_w) // stride_w + 1
+    # change the layout from HWCN to NCHW
+    at = a_np.transpose((3, 2, 0, 1))
+    wt = w_np.transpose((3, 2, 0, 1))
+    bt = np.zeros((batch, out_channel, out_height, out_width))
+    # computation
+    for n in range(batch):
+        for f in range(out_channel):
+            for c in range(in_channel):
+                if pad_h > 0 or pad_w > 0:
+                    apad = np.zeros((in_height + pad_h, in_width + pad_w))
+                    apad[pad_top : pad_top + in_height, pad_left : pad_left + in_width] = at[n, c]
+                else:
+                    apad = at[n, c]
+                out = scipy.signal.convolve2d(apad, np.rot90(np.rot90(wt[f, c])), mode="valid")
+                bt[n, f] += out[::stride_h, ::stride_w]
+    return bt.transpose((2, 3, 1, 0))

python/tvm/topi/testing/conv2d_nchw_python.py

@@ -0,0 +1,159 @@
+# 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, line-too-long, unused-variable, too-many-locals, too-many-branches
+# ruff: noqa: F841
+"""Convolution in python"""
+
+import numpy as np
+import scipy
+
+from tvm.topi.nn.utils import get_pad_tuple
+
+
+def _conv2d_nchw_python(a_np, w_np, stride, padding):
+    """Convolution operator in NCHW layout.
+
+    Parameters
+    ----------
+    a_np : numpy.ndarray
+        4-D with shape [batch, in_channel, in_height, in_width]
+
+    w_np : numpy.ndarray
+        4-D with shape [num_filter, in_channel, filter_height, filter_width]
+
+    stride : int or a list/tuple of two ints
+        Stride size, or [stride_height, stride_width]
+
+    padding : int or str or a list/tuple of 2 or 4 ints
+        Padding size, or ['VALID', 'SAME'], or
+        [pad_height, pad_width] for 2 ints, or
+        [pad_top, pad_left, pad_bottom, pad_right] for 2 ints
+
+    Returns
+    -------
+    b_np : np.ndarray
+        4-D with shape [batch, out_channel, out_height, out_width]
+    """
+    batch, in_channel, in_height, in_width = a_np.shape
+    num_filter, _, kernel_h, kernel_w = w_np.shape
+    if isinstance(stride, int):
+        stride_h = stride_w = stride
+    else:
+        stride_h, stride_w = stride
+    pad_top, pad_left, pad_bottom, pad_right = get_pad_tuple(padding, (kernel_h, kernel_w))
+    pad_h = pad_top + pad_bottom
+    pad_w = pad_left + pad_right
+    # compute the output shape
+    out_channel = num_filter
+    out_height = (in_height - kernel_h + pad_h) // stride_h + 1
+    out_width = (in_width - kernel_w + pad_w) // stride_w + 1
+    b_np = np.zeros((batch, out_channel, out_height, out_width), dtype=a_np.dtype)
+    # computation
+    for n in range(batch):
+        for f in range(out_channel):
+            for c in range(in_channel):
+                if pad_h > 0 or pad_w > 0:
+                    apad = np.zeros((in_height + pad_h, in_width + pad_w), dtype=a_np.dtype)
+                    apad[pad_top : pad_top + in_height, pad_left : pad_left + in_width] = a_np[n, c]
+                else:
+                    apad = a_np[n, c]
+
+                out = _conv2d_hw(apad, w_np[f, c])
+                b_np[n, f] += out[::stride_h, ::stride_w]
+    return b_np
+
+
+def _conv2d_hw(apad, w_np_fc):
+    """2d convolution operator in HW layout.
+
+    This is intended to be used as a subroutine from
+    _conv2d_nchw_python.  Using scipy.signal.convolve2d directly does
+    not work for all dtypes (e.g. float16).  Where possible, this
+    function uses scipy.signal.convolve2d to take advantage of
+    compiled scipy routines, falling back to an explicit loop only
+    where needed
+
+    Parameters
+    ----------
+    a_np : numpy.ndarray
+        2-D with shape [in_height, in_width]
+
+    w_np : numpy.ndarray
+        2-D with shape [filter_height, filter_width].
+
+    Returns
+    -------
+    b_np : np.ndarray
+        2-D with shape [out_height, out_width]
+    """
+
+    try:
+        return scipy.signal.convolve2d(apad, np.rot90(np.rot90(w_np_fc)), mode="valid")
+    except ValueError:
+        pass
+
+    assert len(apad.shape) == len(w_np_fc.shape) == 2
+
+    dtype = apad.dtype
+    in_height, in_width = apad.shape
+    kernel_h, kernel_w = w_np_fc.shape
+
+    output_shape = [a_dim - w_dim + 1 for a_dim, w_dim in zip(apad.shape, w_np_fc.shape)]
+    output = np.zeros(output_shape, dtype=apad.dtype)
+
+    for y in range(output_shape[0]):
+        for x in range(output_shape[1]):
+            output[y][x] = np.sum(apad[y : y + kernel_h, x : x + kernel_w] * w_np_fc)
+
+    return output
+
+
+def conv2d_nchw_python(a_np, w_np, stride, padding, groups=1):
+    """Convolution operator in NCHW layout.
+
+    Parameters
+    ----------
+    a_np : numpy.ndarray
+        4-D with shape [batch, in_channel, in_height, in_width]
+
+    w_np : numpy.ndarray
+        4-D with shape [num_filter, in_channel // groups, filter_height, filter_width]
+
+    stride : int or a list/tuple of two ints
+        Stride size, or [stride_height, stride_width]
+
+    padding : int or str or a list/tuple of 2 or 4 ints
+        Padding size, or ['VALID', 'SAME'], or
+        [pad_height, pad_width] for 2 ints, or
+        [pad_top, pad_left, pad_bottom, pad_right] for 2 ints
+
+    groups : int
+        Number of groups
+
+    Returns
+    -------
+    b_np : np.ndarray
+        4-D with shape [batch, out_channel, out_height, out_width]
+    """
+    a_slices = np.array_split(a_np, groups, axis=1)
+    w_slices = np.array_split(w_np, groups, axis=0)
+    b_slices = [
+        _conv2d_nchw_python(a_slice, w_slice, stride, padding)
+        for a_slice, w_slice in zip(a_slices, w_slices)
+    ]
+    b_np = np.concatenate(b_slices, axis=1)
+    return b_np

python/tvm/topi/testing/conv2d_nhwc_python.py

@@ -0,0 +1,116 @@
+# 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, line-too-long, unused-variable, too-many-locals
+"""Convolution in python"""
+
+import numpy as np
+import scipy.signal
+
+from tvm.topi.nn.utils import get_pad_tuple
+
+
+def _conv2d_nhwc_python(a_np, w_np, stride, padding):
+    """Convolution operator in NHWC layout.
+
+    Parameters
+    ----------
+    a_np : numpy.ndarray
+        4-D with shape [batch, in_height, in_width, in_channel]
+
+    w_np : numpy.ndarray
+        4-D with shape [filter_height, filter_width, in_channel, num_filter]
+
+    stride : int or a list/tuple of two ints
+        Stride size, or [stride_height, stride_width]
+
+    padding : int or str or a list/tuple of two ints
+        Padding size, or ['VALID', 'SAME'], or [pad_height, pad_width]
+
+    Returns
+    -------
+    b_np : np.ndarray
+        4-D with shape [batch, out_height, out_width, out_channel]
+    """
+    batch, in_height, in_width, in_channel = a_np.shape
+    kernel_h, kernel_w, _, num_filter = w_np.shape
+    if isinstance(stride, int):
+        stride_h = stride_w = stride
+    else:
+        stride_h, stride_w = stride
+
+    pad_top, pad_left, pad_bottom, pad_right = get_pad_tuple(padding, (kernel_h, kernel_w))
+    pad_h = pad_top + pad_bottom
+    pad_w = pad_left + pad_right
+
+    # compute the output shape
+    out_channel = num_filter
+    out_height = (in_height - kernel_h + pad_h) // stride_h + 1
+    out_width = (in_width - kernel_w + pad_w) // stride_w + 1
+    # change the layout from NHWC to NCHW
+    at = a_np.transpose((0, 3, 1, 2))
+    wt = w_np.transpose((3, 2, 0, 1))
+    bt = np.zeros((batch, out_channel, out_height, out_width))
+    # computation
+    for n in range(batch):
+        for f in range(out_channel):
+            for c in range(in_channel):
+                if pad_h > 0 or pad_w > 0:
+                    apad = np.zeros((in_height + pad_h, in_width + pad_w))
+                    apad[pad_top : pad_top + in_height, pad_left : pad_left + in_width] = at[n, c]
+                else:
+                    apad = at[n, c]
+                out = scipy.signal.convolve2d(apad, np.rot90(np.rot90(wt[f, c])), mode="valid")
+                bt[n, f] += out[::stride_h, ::stride_w]
+    return bt.transpose((0, 2, 3, 1))
+
+
+def conv2d_nhwc_python(a_np, w_np, stride, padding, groups=1):
+    """Convolution operator in NHWC layout.
+
+    Parameters
+    ----------
+    a_np : numpy.ndarray
+        4-D with shape [batch, in_height, in_width, in_channel]
+
+    w_np : numpy.ndarray
+        4-D with shape [filter_height, filter_width, in_channel // groups, num_filter]
+
+    stride : int or a list/tuple of two ints
+        Stride size, or [stride_height, stride_width]
+
+    padding : int or str or a list/tuple of 2 or 4 ints
+        Padding size, or ['VALID', 'SAME'], or
+        [pad_height, pad_width] for 2 ints, or
+        [pad_top, pad_left, pad_bottom, pad_right] for 2 ints
+
+    groups : int
+        Number of groups
+
+    Returns
+    -------
+    b_np : np.ndarray
+        4-D with shape [batch, out_height, out_width, out_channel]
+    """
+
+    a_slices = np.array_split(a_np, groups, axis=3)
+    w_slices = np.array_split(w_np, groups, axis=3)
+    b_slices = [
+        _conv2d_nhwc_python(a_slice, w_slice, stride, padding)
+        for a_slice, w_slice in zip(a_slices, w_slices)
+    ]
+    b_np = np.concatenate(b_slices, axis=3)
+    return b_np

python/tvm/topi/testing/conv2d_transpose_python.py

@@ -0,0 +1,183 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+# pylint: disable=unused-variable
+"""Transposed convolution in python"""
+
+import numpy as np
+import scipy
+
+import tvm.topi.testing
+from tvm.topi.nn.utils import get_pad_tuple
+
+
+def _conv2d_transpose_nchw_python(a_np, w_np, stride, padding, output_padding):
+    """Transposed convolution operator in NCHW layout.
+
+    Parameters
+    ----------
+    a_np : numpy.ndarray
+        4-D with shape [batch, in_channel, in_height, in_width]
+
+    w_np : numpy.ndarray
+        4-D with shape [in_channel, num_filter, filter_height, filter_width]
+
+    stride : int or a list/tuple of two ints
+        Stride size, or [stride_height, stride_width]
+
+    padding : int or str
+        Padding size, or ['VALID', 'SAME']
+
+    output_padding : int or a list/tuple of two ints
+        Use to disambiguate the output shape.
+
+    Returns
+    -------
+    b_np : np.ndarray
+        4-D with shape [batch, out_channel, out_height, out_width]
+    """
+    batch, in_c, in_h, in_w = a_np.shape
+    _, out_c, filter_h, filter_w = w_np.shape
+    if isinstance(stride, int):
+        stride_h = stride_w = stride
+    else:
+        stride_h, stride_w = stride
+    if isinstance(output_padding, int):
+        opad_h = opad_w = output_padding
+    else:
+        opad_h, opad_w = output_padding
+    assert opad_h < stride_h and opad_w < stride_w
+    # dilate stage
+    dilated_a_np = tvm.topi.testing.dilate_python(a_np, [1, 1, stride_h, stride_w])
+    # padding stage
+    fpad_top, fpad_left, fpad_bottom, fpad_right = get_pad_tuple(padding, (filter_h, filter_w))
+    bpad_top = filter_h - 1 - fpad_top
+    bpad_bottom = filter_h - 1 - fpad_bottom + opad_h
+    bpad_left = filter_w - 1 - fpad_left
+    bpad_right = filter_w - 1 - fpad_right + opad_w
+    padded_a_np = np.zeros(
+        (
+            batch,
+            in_c,
+            dilated_a_np.shape[2] + bpad_top + bpad_bottom,
+            dilated_a_np.shape[3] + bpad_left + bpad_right,
+        )
+    ).astype(a_np.dtype)
+    padded_a_np[
+        :,
+        :,
+        bpad_top : dilated_a_np.shape[2] + bpad_top,
+        bpad_left : dilated_a_np.shape[3] + bpad_left,
+    ] = dilated_a_np
+    # convolution stage
+    out_h = (in_h - 1) * stride_h - fpad_top - fpad_bottom + filter_h + opad_h
+    out_w = (in_w - 1) * stride_w - fpad_left - fpad_right + filter_w + opad_w
+    b_np = np.zeros((batch, out_c, out_h, out_w)).astype(a_np.dtype)
+    for n in range(batch):
+        for f in range(out_c):
+            for c in range(in_c):
+                out = scipy.signal.convolve2d(padded_a_np[n, c], w_np[c, f], mode="valid")
+                b_np[n, f] += out
+    return b_np
+
+
+def conv2d_transpose_nhwc_python(
+    a_nhwc, weight, weight_format, stride, padding, output_padding=(0, 0)
+):
+    """Transposed convolution operator in NHWC layout.
+
+    Parameters
+    ----------
+    a_nhwc : numpy.ndarray
+        4-D with shape [batch, in_height, in_width, in_channel]
+
+    weight : numpy.ndarray
+        4-D in formats HWIO, HWOI, OIHW or IOHW
+
+    weight_format : str
+        ['HWIO', 'HWOI', 'OIHW', 'IOHW']
+
+    stride : int or a list/tuple of two ints
+        Stride size, or [stride_height, stride_width]
+
+    padding : int or str
+        Padding size, or ['VALID', 'SAME']
+
+    Returns
+    -------
+    b_np : np.ndarray
+        4-D with shape [batch, out_channel, out_height, out_width]
+    """
+    assert a_nhwc.ndim == 4, "a_nhwc number of dimensions should be 4"
+    assert weight.ndim == 4, "weight number of dimensions should be 4"
+
+    a_nchw = np.transpose(a_nhwc, (0, 3, 1, 2))
+
+    # conv2d_transpose_nchw_python needs kernel layout to be IOHW
+    if weight_format == "HWIO":
+        w_iohw = np.transpose(weight, (2, 3, 0, 1))
+    elif weight_format == "HWOI":
+        w_iohw = np.transpose(weight, (3, 2, 0, 1))
+    elif weight_format == "OIHW":
+        w_iohw = np.transpose(weight, (1, 0, 2, 3))
+    elif weight_format == "IOHW":
+        w_iohw = weight
+    else:
+        raise ValueError("Valid weight_formats are HWIO, HWOI, OIHW or IOHW")
+
+    res_nchw = conv2d_transpose_nchw_python(
+        a_nchw, w_iohw, stride, padding, output_padding=output_padding
+    )
+    res_nhwc = np.transpose(res_nchw, (0, 2, 3, 1))
+    return res_nhwc
+
+
+def conv2d_transpose_nchw_python(a_np, w_np, stride, padding, output_padding, groups=1):
+    """Convolution operator in NCHW layout.
+
+    Parameters
+    ----------
+    a_np : numpy.ndarray
+        4-D with shape [batch, in_channel, in_height, in_width]
+
+    w_np : numpy.ndarray
+        4-D with shape [in_channel, num_filter // groups, filter_height, filter_width]
+
+    stride : int or a list/tuple of two ints
+        Stride size, or [stride_height, stride_width]
+
+    padding : int or str
+        Padding size, or ['VALID', 'SAME']
+
+    output_padding : int or a list/tuple of two ints
+        Use to disambiguate the output shape.
+
+    groups : int
+        Number of groups
+
+    Returns
+    -------
+    b_np : np.ndarray
+        4-D with shape [batch, out_channel, out_height, out_width]
+    """
+    a_slices = np.array_split(a_np, groups, axis=1)
+    w_slices = np.array_split(w_np, groups, axis=0)
+    b_slices = [
+        _conv2d_transpose_nchw_python(a_slice, w_slice, stride, padding, output_padding)
+        for a_slice, w_slice in zip(a_slices, w_slices)
+    ]
+    b_np = np.concatenate(b_slices, axis=1)
+    return b_np

python/tvm/topi/testing/conv3d_ncdhw_python.py

@@ -0,0 +1,97 @@
+# 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, line-too-long, unused-variable, too-many-locals, too-many-branches
+"""Convolution 3D in python"""
+
+import numpy as np
+import scipy.signal
+
+from tvm.topi.nn.utils import get_pad_tuple3d
+
+
+def _conv3d_ncdhw_python(a_np, w_np, stride, padding):
+    batch, in_channel, in_depth, in_height, in_width = a_np.shape
+    num_filter, _, kernel_d, kernel_h, kernel_w = w_np.shape
+    if isinstance(stride, int):
+        stride_d = stride_h = stride_w = stride
+    else:
+        stride_d, stride_h, stride_w = stride
+
+    pad_front, pad_top, pad_left, pad_back, pad_bottom, pad_right = get_pad_tuple3d(
+        padding, (kernel_d, kernel_h, kernel_w)
+    )
+    pad_d = pad_front + pad_back
+    pad_h = pad_top + pad_bottom
+    pad_w = pad_left + pad_right
+
+    # compute the output shape
+    out_channel = num_filter
+    out_depth = (in_depth - kernel_d + pad_d) // stride_d + 1
+    out_height = (in_height - kernel_h + pad_h) // stride_h + 1
+    out_width = (in_width - kernel_w + pad_w) // stride_w + 1
+    b_np = np.zeros((batch, out_channel, out_depth, out_height, out_width))
+    # computation
+    for n in range(batch):
+        for f in range(out_channel):
+            for c in range(in_channel):
+                if pad_d > 0 or pad_h > 0 or pad_w > 0:
+                    apad = np.zeros((in_depth + pad_d, in_height + pad_h, in_width + pad_w))
+                    apad[
+                        pad_front : pad_front + in_depth,
+                        pad_top : pad_top + in_height,
+                        pad_left : pad_left + in_width,
+                    ] = a_np[n, c]
+                else:
+                    apad = a_np[n, c]
+                out = scipy.signal.convolve(apad, np.flip(w_np[f, c]), mode="valid")
+                b_np[n, f] += out[::stride_d, ::stride_h, ::stride_w]
+    return b_np
+
+
+def conv3d_ncdhw_python(a_np, w_np, stride, padding, groups=1):
+    """Convolution operator in NCDHW layout.
+
+    Parameters
+    ----------
+    a_np : numpy.ndarray
+        5-D with shape [batch, in_channel, in_depth, in_height, in_width]
+
+    w_np : numpy.ndarray
+        5-D with shape [num_filter, in_channel, filter_depth, filter_height, filter_width]
+
+    stride : int or a list/tuple of three ints
+        Stride size, or [stride_depth, stride_height, stride_width]
+
+    padding : int or str or a list/tuple of three ints
+        Padding size, or ['VALID', 'SAME'], or [pad_depth, pad_height, pad_width]
+
+    groups : int
+        Number of groups
+
+    Returns
+    -------
+    b_np : np.ndarray
+        5-D with shape [batch, out_channel, out_depth, out_height, out_width]
+    """
+    a_slices = np.array_split(a_np, groups, axis=1)
+    w_slices = np.array_split(w_np, groups, axis=0)
+    b_slices = [
+        _conv3d_ncdhw_python(a_slice, w_slice, stride, padding)
+        for a_slice, w_slice in zip(a_slices, w_slices)
+    ]
+    b_np = np.concatenate(b_slices, axis=1)
+    return b_np

python/tvm/topi/testing/conv3d_ndhwc_python.py

@@ -0,0 +1,122 @@
+# 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, line-too-long, unused-variable, too-many-locals
+"""Convolution 3D in python"""
+
+import numpy as np
+import scipy.signal
+
+from tvm.topi.nn.utils import get_pad_tuple3d
+
+
+def _conv3d_ndhwc_python(a_np, w_np, stride, padding):
+    """Convolution 3D operator in NDHWC layout.
+
+    Parameters
+    ----------
+    a_np : numpy.ndarray
+        5-D with shape [batch, in_channel, in_depth, in_height, in_width]
+
+    w_np : numpy.ndarray
+        5-D with shape [num_filter, in_channel, filter_depth, filter_height, filter_width]
+
+    stride : int or a list/tuple of three ints
+        Stride size, or [stride_depth, stride_height, stride_width]
+
+    padding : int or str or a list/tuple of three ints
+        Padding size, or ['VALID', 'SAME'], or [pad_depth, pad_height, pad_width]
+
+    Returns
+    -------
+    b_np : np.ndarray
+        5-D with shape [batch, out_channel, out_depth, out_height, out_width]
+    """
+    batch, in_depth, in_height, in_width, in_channel = a_np.shape
+    kernel_d, kernel_h, kernel_w, _, num_filter = w_np.shape
+    if isinstance(stride, int):
+        stride_d = stride_h = stride_w = stride
+    else:
+        stride_d, stride_h, stride_w = stride
+
+    pad_front, pad_top, pad_left, pad_back, pad_bottom, pad_right = get_pad_tuple3d(
+        padding, (kernel_d, kernel_h, kernel_w)
+    )
+    pad_d = pad_front + pad_back
+    pad_h = pad_top + pad_bottom
+    pad_w = pad_left + pad_right
+    # compute the output shape
+    out_channel = num_filter
+    out_depth = (in_depth - kernel_d + pad_d) // stride_d + 1
+    out_height = (in_height - kernel_h + pad_h) // stride_h + 1
+    out_width = (in_width - kernel_w + pad_w) // stride_w + 1
+    # change the layout from NHWC to NCHW
+    at = a_np.transpose((0, 4, 1, 2, 3))
+    wt = w_np.transpose((4, 3, 0, 1, 2))
+    bt = np.zeros((batch, out_channel, out_depth, out_height, out_width), dtype=a_np.dtype)
+    # computation
+    for n in range(batch):
+        for f in range(out_channel):
+            for c in range(in_channel):
+                if pad_d > 0 or pad_h > 0 or pad_w > 0:
+                    apad = np.zeros(
+                        (in_depth + pad_d, in_height + pad_h, in_width + pad_w), dtype=a_np.dtype
+                    )
+                    apad[
+                        pad_front : pad_front + in_depth,
+                        pad_top : pad_top + in_height,
+                        pad_left : pad_left + in_width,
+                    ] = at[n, c]
+                else:
+                    apad = at[n, c]
+                out = scipy.signal.convolve(apad, np.flip(wt[f, c]), mode="valid")
+                bt[n, f] += out[::stride_d, ::stride_h, ::stride_w]
+    return bt.transpose((0, 2, 3, 4, 1))
+
+
+def conv3d_ndhwc_python(a_np, w_np, stride, padding, groups=1):
+    """Convolution 3D operator in NDHWC layout.
+
+    Parameters
+    ----------
+    a_np : numpy.ndarray
+        5-D with shape [batch, in_channel, in_depth, in_height, in_width]
+
+    w_np : numpy.ndarray
+        5-D with shape [num_filter, in_channel, filter_depth, filter_height, filter_width]
+
+    stride : int or a list/tuple of three ints
+        Stride size, or [stride_depth, stride_height, stride_width]
+
+    padding : int or str or a list/tuple of three ints
+        Padding size, or ['VALID', 'SAME'], or [pad_depth, pad_height, pad_width]
+
+    groups : int
+        Number of groups
+
+    Returns
+    -------
+    b_np : np.ndarray
+        5-D with shape [batch, out_channel, out_depth, out_height, out_width]
+    """
+    a_slices = np.array_split(a_np, groups, axis=4)
+    w_slices = np.array_split(w_np, groups, axis=4)
+    b_slices = [
+        _conv3d_ndhwc_python(a_slice, w_slice, stride, padding)
+        for a_slice, w_slice in zip(a_slices, w_slices)
+    ]
+    b_np = np.concatenate(b_slices, axis=4)
+    return b_np

python/tvm/topi/testing/conv3d_transpose_ncdhw_python.py

@@ -0,0 +1,145 @@
+# 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, line-too-long, unused-variable, too-many-locals, too-many-branches
+# ruff: noqa: F841
+"""Convolution 3D transpose in python"""
+
+import numpy as np
+
+import tvm.topi.testing
+from tvm.topi.nn.utils import get_pad_tuple3d
+
+
+def _conv3d_transpose_ncdhw_python(a_np, w_np, stride, padding, output_padding):
+    """Transposed 3d convolution operator in NCDHW layout.
+
+    Parameters
+    ----------
+    a_np : numpy.ndarray
+        5-D with shape [batch, in_channel, in_depth, in_height, in_width]
+
+    w_np : numpy.ndarray
+        5-D with shape [in_channel, num_filter, filter_depth, filter_height, filter_width]
+
+    stride : int or a list/tuple of two ints
+        Stride size, or [stride_depth, stride_height, stride_width]
+
+    padding : int or str
+        Padding size
+
+    output_padding : int or list/tuple of three ints
+        Used to disambiguate output shape.
+
+    Returns
+    -------
+    b_np : np.ndarray
+        5-D with shape [batch, out_channel, out_depth, out_height, out_width]
+    """
+    batch, in_c, in_d, in_h, in_w = a_np.shape
+    _, out_c, filter_d, filter_h, filter_w = w_np.shape
+    if isinstance(stride, int):
+        stride_d = stride_h = stride_w = stride
+    else:
+        stride_d, stride_h, stride_w = stride
+    if isinstance(output_padding, int):
+        opad_d = opad_h = opad_w = output_padding
+    else:
+        opad_d, opad_h, opad_w = output_padding
+    assert opad_d < stride_d and opad_h < stride_h and opad_w < stride_w
+
+    # dilate stage
+    dilated_a_np = tvm.topi.testing.dilate_python(a_np, [1, 1, stride_d, stride_h, stride_w])
+
+    # padding stage
+    fpad_front, fpad_top, fpad_left, fpad_back, fpad_bottom, fpad_right = get_pad_tuple3d(
+        padding, (filter_d, filter_h, filter_w)
+    )
+
+    bpad_front = filter_d - 1 - fpad_front
+    bpad_back = filter_d - 1 - fpad_back + opad_d
+    bpad_top = filter_h - 1 - fpad_top
+    bpad_bottom = filter_h - 1 - fpad_bottom + opad_h
+    bpad_left = filter_w - 1 - fpad_left
+    bpad_right = filter_w - 1 - fpad_right + opad_w
+
+    padded_a_np = np.zeros(
+        (
+            batch,
+            in_c,
+            dilated_a_np.shape[2] + bpad_front + bpad_back,
+            dilated_a_np.shape[3] + bpad_top + bpad_bottom,
+            dilated_a_np.shape[4] + bpad_left + bpad_right,
+        )
+    )
+
+    padded_a_np[
+        :,
+        :,
+        bpad_front : dilated_a_np.shape[2] + bpad_front,
+        bpad_top : dilated_a_np.shape[3] + bpad_top,
+        bpad_left : dilated_a_np.shape[4] + bpad_left,
+    ] = dilated_a_np
+
+    # convolution stage
+    out_d = (in_d - 1) * stride_d - bpad_front - bpad_back + filter_d
+    out_h = (in_h - 1) * stride_h - fpad_top - fpad_bottom + filter_h
+    out_w = (in_w - 1) * stride_w - fpad_left - fpad_right + filter_w
+
+    w_np = np.flip(w_np, axis=[2, 3, 4]).transpose((1, 0, 2, 3, 4))
+    b_np = tvm.topi.testing.conv3d_ncdhw_python(
+        padded_a_np, w_np, stride=(1, 1, 1), padding=(0, 0, 0)
+    )
+
+    return b_np
+
+
+def conv3d_transpose_ncdhw_python(a_np, w_np, stride, padding, output_padding, groups=1):
+    """Transposed 3d convolution operator in NCDHW layout.
+
+    Parameters
+    ----------
+    a_np : numpy.ndarray
+        5-D with shape [batch, in_channel, in_depth, in_height, in_width]
+
+    w_np : numpy.ndarray
+        5-D with shape [in_channel, num_filter, filter_depth, filter_height, filter_width]
+
+    stride : int or a list/tuple of two ints
+        Stride size, or [stride_depth, stride_height, stride_width]
+
+    padding : int or str
+        Padding size
+
+    output_padding : int or list/tuple of three ints
+        Used to disambiguate output shape.
+
+    groups : int
+        Number of groups
+
+    Returns
+    -------
+    b_np : np.ndarray
+        5-D with shape [batch, out_channel, out_depth, out_height, out_width]
+    """
+    a_slices = np.array_split(a_np, groups, axis=1)
+    w_slices = np.array_split(w_np, groups, axis=0)
+    b_slices = [
+        _conv3d_transpose_ncdhw_python(a_slice, w_slice, stride, padding, output_padding)
+        for a_slice, w_slice in zip(a_slices, w_slices)
+    ]
+    b_np = np.concatenate(b_slices, axis=1)
+    return b_np

python/tvm/topi/testing/correlation_nchw_python.py

@@ -0,0 +1,109 @@
+# 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, line-too-long, unused-variable, too-many-locals
+# ruff: noqa: E731
+"""Convolution 3D in python"""
+
+import numpy as np
+
+
+def correlation_nchw_python(
+    data1, data2, kernel_size, max_displacement, stride1, stride2, padding, is_multiply
+):
+    """Correlationn operator in NCHW layout.
+
+    Parameters
+    ----------
+    data1_np : numpy.ndarray
+        4-D with shape [batch, in_channel, in_height, in_width]
+
+    data2_np : numpy.ndarray
+        4-D with shape [batch, in_channel, in_height, in_width]
+
+    kernel_size: int
+        Kernel size for correlation, must be an odd number
+
+    max_displacement: int
+        Max displacement of Correlation
+
+    stride1: int
+        Stride for data1
+
+    stride2: int
+        Stride for data2 within the neightborhood centered around data1
+
+    padding: int
+        Padding for correlation
+
+    is_multiply: bool
+        operation type is either multiplication or substraction
+
+    Returns
+    -------
+    c_np : np.ndarray
+        4-D with shape [batch, out_channel, out_height, out_width]
+    """
+    # compute output's dimension
+    pad_data_height = data1.shape[2] + 2 * padding
+    pad_data_width = data1.shape[3] + 2 * padding
+    kernel_radius = (kernel_size - 1) // 2
+    border_size = max_displacement + kernel_radius
+    out_width = (pad_data_width - border_size * 2) // stride1
+    out_height = (pad_data_height - border_size * 2) // stride1
+    neighborhood_grid_radius = max_displacement // stride2
+    neighborhood_grid_width = neighborhood_grid_radius * 2 + 1
+    out_channel = neighborhood_grid_width * neighborhood_grid_width
+
+    out = np.zeros((data1.shape[0], out_channel, out_height, out_width))
+    pad_data1 = np.zeros((data1.shape[0], data1.shape[1], pad_data_height, pad_data_width))
+    pad_data2 = np.zeros((data1.shape[0], data1.shape[1], pad_data_height, pad_data_width))
+
+    pad_data1[:, :, padding : padding + data1.shape[2], padding : padding + data1.shape[3]] = data1[
+        :, :, :, :
+    ]
+    pad_data2[:, :, padding : padding + data2.shape[2], padding : padding + data2.shape[3]] = data2[
+        :, :, :, :
+    ]
+
+    if is_multiply:
+        corr_func = lambda x, y: x * y
+    else:
+        corr_func = lambda x, y: abs(x - y)
+
+    # pylint: disable=too-many-nested-blocks
+    for i in range(out_height):
+        for j in range(out_width):
+            for nbatch in range(data1.shape[0]):
+                # x1,y1 is the location in data1 , i,j is the location in output
+                x1 = j * stride1 + max_displacement
+                y1 = i * stride1 + max_displacement
+
+                for q in range(out_channel):
+                    # location in data2
+                    x2 = x1 + (q % neighborhood_grid_width - neighborhood_grid_radius) * stride2
+                    y2 = y1 + (q // neighborhood_grid_width - neighborhood_grid_radius) * stride2
+
+                    for h in range(kernel_size):
+                        for w in range(kernel_size):
+                            for channel in range(data1.shape[1]):
+                                out[nbatch, q, i, j] += corr_func(
+                                    pad_data1[nbatch, channel, y1 + h, x1 + w],
+                                    pad_data2[nbatch, channel, y2 + h, x2 + w],
+                                )
+
+    out /= float(kernel_size**2 * data1.shape[1])
+    return out

python/tvm/topi/testing/crop_and_resize_python.py

@@ -0,0 +1,121 @@
+# 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, line-too-long, unused-variable, too-many-locals, too-many-nested-blocks
+"""crop and resize in python"""
+
+import math
+
+import numpy as np
+
+
+def crop_and_resize_python(
+    image, boxes, box_indices, crop_size, layout, method="bilinear", extrapolation_value=0
+):
+    """Crop and resize using python"""
+    (target_h, target_w) = crop_size
+
+    if layout == "NHWC":
+        batch = boxes.shape[0]
+        image_height, image_width, channel = image.shape[1], image.shape[2], image.shape[3]
+        scaled_image = np.ones((batch, target_h, target_w, channel))
+    else:
+        batch = boxes.shape[0]
+        channel, image_height, image_width = image.shape[1], image.shape[2], image.shape[3]
+        scaled_image = np.ones((batch, channel, target_h, target_w))
+
+    for n, box in enumerate(boxes):
+        b_in = box_indices[n]
+        y1, x1 = boxes[n][0], boxes[n][1]
+        y2, x2 = boxes[n][2], boxes[n][3]
+
+        in_h = (image_height - 1) * (y2 - y1)
+        in_w = (image_width - 1) * (x2 - x1)
+        h_scale = np.float32(in_h) / np.float32(target_h - 1)
+        w_scale = np.float32(in_w) / np.float32(target_w - 1)
+
+        for y in range(target_h):
+            in_y = y1 * (image_height - 1) + h_scale * y
+
+            if in_y < 0 or in_y > image_height - 1:
+                for x in range(target_w):
+                    for d in range(channel):
+                        if layout == "NHWC":
+                            scaled_image[n][y][x][d] = extrapolation_value
+                        else:
+                            scaled_image[n][d][y][x] = extrapolation_value
+                continue
+
+            if method == "bilinear":
+                top_y_index = math.floor(in_y)
+                bottom_y_index = math.ceil(in_y)
+                y_lerp = in_y - top_y_index
+
+                for x in range(target_w):
+                    in_x = x1 * (image_width - 1) + x * w_scale
+                    if in_x < 0 or in_x > image_width - 1:
+                        for d in range(channel):
+                            if layout == "NHWC":
+                                scaled_image[n][y][x][d] = extrapolation_value
+                            else:
+                                scaled_image[n][d][y][x] = extrapolation_value
+                        continue
+
+                    left_x_index = math.floor(in_x)
+                    right_x_index = math.ceil(in_x)
+                    x_lerp = in_x - left_x_index
+
+                    for d in range(channel):
+                        if layout == "NHWC":
+                            top_left = image[b_in][top_y_index][left_x_index][d]
+                            top_right = image[b_in][top_y_index][right_x_index][d]
+                            bottom_left = image[b_in][bottom_y_index][left_x_index][d]
+                            bottom_right = image[b_in][bottom_y_index][right_x_index][d]
+                            top = top_left + (top_right - top_left) * x_lerp
+                            bottom = bottom_left + (bottom_right - bottom_left) * x_lerp
+                            scaled_image[n][y][x][d] = top + (bottom - top) * y_lerp
+                        else:
+                            top_left = image[b_in][d][top_y_index][left_x_index]
+                            top_right = image[b_in][d][top_y_index][right_x_index]
+                            bottom_left = image[b_in][d][bottom_y_index][left_x_index]
+                            bottom_right = image[b_in][d][bottom_y_index][right_x_index]
+                            top = top_left + (top_right - top_left) * x_lerp
+                            bottom = bottom_left + (bottom_right - bottom_left) * x_lerp
+                            scaled_image[n][d][y][x] = top + (bottom - top) * y_lerp
+
+            elif method == "nearest_neighbor":
+                for x in range(target_w):
+                    in_x = x1 * (image_width - 1) + x * w_scale
+                    if in_x < 0 or in_x > image_width - 1:
+                        for d in range(channel):
+                            if layout == "NHWC":
+                                scaled_image[n][y][x][d] = extrapolation_value
+                            else:
+                                scaled_image[n][d][y][x] = extrapolation_value
+                        continue
+                    closest_x_index = np.round(in_x).astype("int32")
+                    closest_y_index = np.round(in_y).astype("int32")
+                    for d in range(channel):
+                        if layout == "NHWC":
+                            scaled_image[n][y][x][d] = image[b_in][closest_y_index][
+                                closest_x_index
+                            ][d]
+                        else:
+                            scaled_image[n][d][y][x] = image[b_in][d][closest_y_index][
+                                closest_x_index
+                            ]
+
+    return scaled_image

python/tvm/topi/testing/deformable_conv2d_python.py

@@ -0,0 +1,181 @@
+# 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-arguments
+"""Deformable convolution in python"""
+
+import itertools
+import math
+
+import numpy as np
+
+from tvm.topi.nn.utils import get_pad_tuple
+
+
+def deformable_conv2d_nchw_python(
+    a_np, offset_np, w_np, stride, padding, dilation, deformable_groups, groups
+):
+    """Deformable convolution operator in NCHW layout.
+
+    Parameters
+    ----------
+    a_np : numpy.ndarray
+        4-D with shape [batch, in_channel, in_height, in_width]
+
+    offset_np : numpy.ndarray
+        4-D with shape [batch, deformable_groups * filter_height * filter_width * 2,
+                        out_height, out_width]
+
+    w_np : numpy.ndarray
+        4-D with shape [num_filter, in_channel, filter_height, filter_width]
+
+    stride : int or a list/tuple of two ints
+        Stride size, or [stride_height, stride_width]
+
+    padding : int or str or a list/tuple of 2 or 4 ints
+        Padding size, or ['VALID', 'SAME'], or
+        [pad_height, pad_width] for 2 ints, or
+        [pad_top, pad_left, pad_bottom, pad_right] for 2 ints
+
+    dilation : int or a list/tuple of two ints
+        Dilation size, or [dilate_height, dilate_width]
+
+    deformable_groups : int
+        Number of deformable groups
+
+    groups : int
+        Number of groups
+
+    Returns
+    -------
+    b_np : np.ndarray
+        4-D with shape [batch, out_channel, out_height, out_width]
+    """
+    batch, in_channel, in_height, in_width = a_np.shape
+    out_channel, _, kernel_h, kernel_w = w_np.shape
+    out_height, out_width = offset_np.shape[-2:]
+    dtype = a_np.dtype
+    ic_per_dgroup = in_channel // deformable_groups
+    assert groups == 1, "deformable_conv2d_nchw_python does not support groups > 1"
+
+    if isinstance(stride, int):
+        stride_h = stride_w = stride
+    else:
+        stride_h, stride_w = stride
+
+    pad_top, pad_left, _, _ = get_pad_tuple(padding, (kernel_h, kernel_w))
+
+    if isinstance(dilation, int):
+        dilation_h = dilation_w = dilation
+    else:
+        dilation_h, dilation_w = dilation
+
+    def _bilinear(n, c, h, w):
+        y_low = math.floor(h)
+        x_low = math.floor(w)
+        y_high = y_low + 1
+        x_high = x_low + 1
+
+        wy_h = h - y_low
+        wx_h = w - x_low
+        wy_l = 1 - wy_h
+        wx_l = 1 - wx_h
+
+        val = 0
+        for wx, xp in zip((wx_l, wx_h), (x_low, x_high)):
+            for wy, yp in zip((wy_l, wy_h), (y_low, y_high)):
+                if 0 <= yp < in_height and 0 <= xp < in_width:
+                    val += wx * wy * a_np[n, c, yp, xp]
+        return val
+
+    a_deform = np.zeros((batch, in_channel, out_height, out_width, kernel_h, kernel_w), dtype=dtype)
+    for n, h, w in itertools.product(range(batch), range(out_height), range(out_width)):
+        offset = offset_np[n, :, h, w].reshape(deformable_groups, kernel_h, kernel_w, 2)
+        in_h = h * stride_h - pad_top
+        in_w = w * stride_w - pad_left
+
+        index_h_base, index_w_base = np.meshgrid(
+            np.arange(in_h, in_h + kernel_h * dilation_h, dilation_h, dtype=offset_np.dtype),
+            np.arange(in_w, in_w + kernel_w * dilation_w, dilation_w, dtype=offset_np.dtype),
+            indexing="ij",
+        )
+
+        for c, kh, kw in itertools.product(range(in_channel), range(kernel_h), range(kernel_w)):
+            dg = c // ic_per_dgroup
+            index_h = index_h_base + offset[dg, ..., 0]
+            index_w = index_w_base + offset[dg, ..., 1]
+
+            y, x = index_h[kh, kw], index_w[kh, kw]
+            if y < 0 or y >= in_height or x < 0 or x >= in_width:
+                continue
+            a_deform[n, c, h, w, kh, kw] = _bilinear(n, c, y, x)
+
+    b_np = np.zeros((batch, out_channel, out_height, out_width), dtype=dtype)
+    for n, c, f, h, w in itertools.product(
+        range(batch), range(in_channel), range(out_channel), range(out_height), range(out_width)
+    ):
+        b_np[n, f, h, w] += np.tensordot(a_deform[n, c, h, w], w_np[f, c])
+
+    return b_np
+
+
+def deformable_conv2d_nhwc_python(
+    a_np, offset_np, w_np, stride, padding, dilation, deformable_groups, groups
+):
+    """Deformable convolution operator in NHWC layout.
+
+    Parameters
+    ----------
+    a_np : numpy.ndarray
+        4-D with shape [batch, in_height, in_width, in_channel]
+
+    offset_np : numpy.ndarray
+        4-D with shape [batch, out_height, out_width,
+                        deformable_groups * filter_height * filter_width * 2]
+
+    w_np : numpy.ndarray
+        4-D with shape [filter_height, filter_width, in_channel, num_filter]
+
+    stride : int or a list/tuple of two ints
+        Stride size, or [stride_height, stride_width]
+
+    padding : int or str or a list/tuple of 2 or 4 ints
+        Padding size, or ['VALID', 'SAME'], or
+        [pad_height, pad_width] for 2 ints, or
+        [pad_top, pad_left, pad_bottom, pad_right] for 2 ints
+
+    dilation : int or a list/tuple of two ints
+        Dilation size, or [dilate_height, dilate_width]
+
+    deformable_groups : int
+        Number of deformable groups
+
+    groups : int
+        Number of groups
+
+    Returns
+    -------
+    b_np : np.ndarray
+        4-D with shape [batch, out_channel, out_height, out_width]
+    """
+    a_np = np.transpose(a_np, [0, 3, 1, 2])  # NHWC -> NCHW
+    offset_np = np.transpose(offset_np, [0, 3, 1, 2])  # NHWC -> NCHW
+    w_np = np.transpose(w_np, [3, 2, 0, 1])  # HWIO -> OIHW
+    b_np = deformable_conv2d_nchw_python(
+        a_np, offset_np, w_np, stride, padding, dilation, deformable_groups, groups
+    )
+    b_np = np.transpose(b_np, [0, 2, 3, 1])  # NCHW -> NHWC
+    return b_np

python/tvm/topi/testing/dense.py

@@ -0,0 +1,54 @@
+# 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
+"""Dense in python"""
+
+import numpy as np
+
+
+def dense(x, y, bias, use_bias=False, use_relu=False, out_dtype=None):
+    """dense operator implemented in numpy.
+
+    Parameters
+    ----------
+    x : numpy.ndarray
+        2-D with shape [M, K]
+
+    y : numpy.ndarray
+        2-D with shape [N, K]
+
+    bias: numpy.ndarray
+        1-D with shape [M,]
+
+    out_dtype: string, optional
+        Specify the dtype of output
+
+    Returns
+    -------
+    out : numpy.ndarray
+        2-D with shape [M, N]
+    """
+    dtype = x.dtype if out_dtype is None else out_dtype
+    if use_bias:
+        out = np.dot(x.astype(dtype), y.T.astype(dtype)) + bias
+    else:
+        out = np.dot(x.astype(dtype), y.T.astype(dtype))
+
+    if use_relu:
+        out = np.maximum(out, 0)
+
+    return out

python/tvm/topi/testing/depth_to_space.py

@@ -0,0 +1,53 @@
+# 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, line-too-long, unused-variable, too-many-locals
+"""Depth to space in python"""
+
+import numpy as np
+
+
+def depth_to_space_python(data, block_size, mode="DCR"):
+    """Depth to Space operator in python for NCHW layout.
+
+    Parameters
+    ----------
+    data : np.ndarray
+        4-D with shape [batch, in_channel, in_height, in_width]
+
+    block_size : int
+        Size of blocks to convert channel pixels into.
+
+    Returns
+    -------
+    d2s_out : np.ndarray
+        4-D with shape [batch, in_channel / (block_size * block_size),
+                        out_height * block_size, out_width * block_size]
+    """
+    in_n, in_c, in_h, in_w = data.shape
+    new_h = int(in_h * block_size)
+    new_w = int(in_h * block_size)
+    new_c = int(in_c / (block_size * block_size))
+
+    if mode == "DCR":
+        expanded = np.reshape(data, newshape=[in_n, block_size, block_size, new_c, in_h, in_w])
+        transposed = np.transpose(expanded, axes=[0, 3, 4, 1, 5, 2])
+    else:
+        expanded = np.reshape(data, newshape=(in_n, new_c, block_size, block_size, in_h, in_w))
+        transposed = np.transpose(expanded, axes=(0, 1, 4, 2, 5, 3))
+    newshape = [in_n, new_c, new_h, new_w]
+    d2s_out = np.reshape(transposed, newshape=newshape)
+    return d2s_out

python/tvm/topi/testing/depthwise_conv2d_python.py

@@ -0,0 +1,167 @@
+# 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, unused-variable, line-too-long
+"""Depthwise convolution in python"""
+
+import numpy as np
+
+from tvm.topi.nn.utils import get_pad_tuple
+
+from .common import _convolve2d
+
+
+def depthwise_conv2d_python_nchw(input_np, filter_np, stride, padding):
+    """Depthwise convolution operator in NCHW layout.
+
+    Parameters
+    ----------
+    input_np : numpy.ndarray
+        4-D with shape [batch, in_channel, in_height, in_width]
+
+    filter_np : numpy.ndarray
+        4-D with shape [in_channel, channel_multiplier, filter_height, filter_width]
+
+    stride : list / tuple of 2 ints
+        [stride_height, stride_width]
+
+    padding : str
+        'VALID' or 'SAME'
+
+    Returns
+    -------
+    output_np : np.ndarray
+        4-D with shape [batch, out_channel, out_height, out_width]
+    """
+    batch, in_channel, in_height, in_width = input_np.shape
+    _, channel_multiplier, filter_height, filter_width = filter_np.shape
+    if isinstance(stride, int):
+        stride_h = stride_w = stride
+    else:
+        stride_h, stride_w = stride
+
+    pad_top, pad_left, pad_bottom, pad_right = get_pad_tuple(padding, (filter_height, filter_width))
+    pad_h = pad_top + pad_bottom
+    pad_w = pad_left + pad_right
+
+    out_channel = in_channel * channel_multiplier
+    out_height = (in_height - filter_height + pad_h) // stride_h + 1
+    out_width = (in_width - filter_width + pad_w) // stride_w + 1
+    output_np = np.zeros((batch, out_channel, out_height, out_width))
+
+    for i in range(batch):
+        for j in range(out_channel):
+            apad = input_np[i, j // channel_multiplier, :, :]
+            if pad_h or pad_w:
+                apad = np.pad(apad, [(pad_top, pad_bottom), (pad_left, pad_right)], "constant")
+
+            conv = _convolve2d(
+                apad,
+                np.rot90(filter_np[j // channel_multiplier, j % channel_multiplier, :, :], k=2),
+            )
+            output_np[i, j, :, :] = conv[
+                ::stride_h,
+                ::stride_w,
+            ]
+
+    return output_np
+
+
+def depthwise_conv2d_python_nchwc(input_np, filter_np, stride, padding):
+    """Depthwise convolution operator in NCHWc layout.
+
+    Parameters
+    ----------
+    input_np : numpy.ndarray
+        5-D with shape [batch, in_channel_chunk, in_height, in_width, in_channel_block]
+
+    filter_np : numpy.ndarray
+        6-D with shape [out_channel_chunk, channel_multiplier_chunk,
+                        filter_height, filter_width,
+                        channel_multiplier_block, out_channel_block]
+
+    stride : list / tuple of 2 ints
+        [stride_height, stride_width]
+
+    padding : str
+        'VALID' or 'SAME'
+
+    Returns
+    -------
+    output_np : np.ndarray
+        5-D with shape [batch, out_channel_chunk, out_height, out_width, out_channel_block]
+    """
+    # Transform to NCHW
+    batch_size, in_channel_chunk, in_height, in_width, in_channel_block = input_np.shape
+    input_nchw = input_np.transpose(0, 1, 4, 2, 3).reshape(
+        (batch_size, in_channel_chunk * in_channel_block, in_height, in_width)
+    )
+
+    (
+        out_channel_chunk,
+        channel_multiplier_chunk,
+        filter_height,
+        filter_width,
+        channel_multiplier_block,
+        out_channel_block,
+    ) = filter_np.shape
+    filter_nchw = filter_np.transpose(0, 5, 1, 4, 2, 3).reshape(
+        (
+            out_channel_chunk * out_channel_block,
+            channel_multiplier_chunk * channel_multiplier_block,
+            filter_height,
+            filter_width,
+        )
+    )
+
+    # Perform conv2d
+    output_np = depthwise_conv2d_python_nchw(input_nchw, filter_nchw, stride, padding)
+
+    # Transform back to NCHWc
+
+    # pylint: disable=unpacking-non-sequence
+    batch_size, out_channel, out_height, out_width = output_np.shape
+    return output_np.reshape(
+        (batch_size, out_channel_chunk, out_channel_block, out_height, out_width)
+    ).transpose(0, 1, 3, 4, 2)
+
+
+def depthwise_conv2d_python_nhwc(input_np, filter_np, stride, padding):
+    """Depthwise convolution operator in nhwc layout.
+
+    Parameters
+    ----------
+    input_np : numpy.ndarray
+        4-D with shape [batch, in_height, in_width, in_channel]
+
+    filter_np : numpy.ndarray
+        4-D with shape [filter_height, filter_width, in_channel, channel_multiplier]
+
+    stride : list / tuple of 2 ints
+        [stride_height, stride_width]
+
+    padding : str
+        'VALID' or 'SAME'
+
+    Returns
+    -------
+    output_np : np.ndarray
+        4-D with shape [batch, out_height, out_width, out_channel]
+    """
+    input_nchw = input_np.transpose(0, 3, 1, 2)
+    filter_nchw = filter_np.transpose(2, 3, 0, 1)
+    output_nchw = depthwise_conv2d_python_nchw(input_nchw, filter_nchw, stride, padding)
+    return output_nchw.transpose(0, 2, 3, 1)