[Relax][Op] Add missing TOPI-backed operators: correlation, lrn, dilate, dilation2d

include/tvm/relax/op/image.h

@@ -80,6 +80,18 @@ struct AffineGridAttrs : public AttrsNode {
   TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.AffineGridAttrs", AffineGridAttrs, AttrsNode);
 };  // struct AffineGridAttrs
 
+/*! \brief Attributes used in image dilation2d operator */
+struct Dilation2DAttrs : public AttrsNode {
+  ffi::Array<int64_t> strides;
+  ffi::Array<int64_t> padding;
+  ffi::Array<int64_t> dilations;
+  ffi::String data_layout;
+  ffi::Optional<DLDataType> out_dtype;
+
+  static void RegisterReflection();
+  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.Dilation2DAttrs", Dilation2DAttrs, AttrsNode);
+};  // struct Dilation2DAttrs
+
 }  // namespace relax
 }  // namespace tvm
 

include/tvm/relax/op/nn.h

@@ -339,6 +339,41 @@ struct PixelShuffleAttrs : public AttrsNode {
   TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.PixelShuffleAttrs", PixelShuffleAttrs, AttrsNode);
 };
 
+/*! \brief Attributes used for the correlation operator */
+struct CorrelationAttrs : public AttrsNode {
+  int kernel_size;
+  int max_displacement;
+  int stride1;
+  int stride2;
+  int padding;
+  bool is_multiply;
+  ffi::String data_layout;
+
+  static void RegisterReflection();
+  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.CorrelationAttrs", CorrelationAttrs, AttrsNode);
+};  // struct CorrelationAttrs
+
+/*! \brief Attributes used for the local response normalization (LRN) operator */
+struct LRNAattrs : public AttrsNode {
+  int size;
+  int axis;
+  double alpha;
+  double beta;
+  double bias;
+
+  static void RegisterReflection();
+  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.LRNAattrs", LRNAattrs, AttrsNode);
+};  // struct LRNAattrs
+
+/*! \brief Attributes used for the dilate (upsample with holes) operator */
+struct DilateAttrs : public AttrsNode {
+  ffi::Array<int64_t> strides;
+  double dilation_value;
+
+  static void RegisterReflection();
+  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.DilateAttrs", DilateAttrs, AttrsNode);
+};  // struct DilateAttrs
+
 }  // namespace relax
 }  // namespace tvm
 

python/tvm/relax/op/image/__init__.py

@@ -19,7 +19,9 @@
 
 from .image import (
     Resize2DAttrs,
+    Dilation2DAttrs,
     affine_grid,
+    dilation2d,
     grid_sample,
     resize2d,
     resize3d,

python/tvm/relax/op/image/image.py

@@ -317,3 +317,99 @@ def affine_grid(
             ty=ty,
         ),
     )
+
+
+@tvm_ffi.register_object("relax.attrs.Dilation2DAttrs")
+class Dilation2DAttrs(Attrs):
+    """Attributes for image.dilation2d"""
+
+
+def dilation2d(
+    input: Expr,
+    filter: Expr,
+    strides: int | tuple[int, int] = (1, 1),
+    padding: int | tuple[int, int, int, int] | str = 0,
+    dilations: int | tuple[int, int] = (1, 1),
+    data_layout: str = "NCHW",
+    out_dtype: str | tvm_ffi.dtype | None = None,
+    *,
+    ty=None,
+    loc: Location = UNKNOWN_LOC,
+) -> Expr:
+    r"""Morphological Dilation 2D operator.
+
+    This operator computes the morphological dilation of the input image
+    with the given filter (structuring element). For each position, the
+    output is the maximum of `input + filter` over the filter window.
+
+    .. math::
+
+        output[n, c, h, w] = \max_{f_h, f_w}
+            \left( input[n, c, h' + d_h \cdot f_h, w' + d_w \cdot f_w]
+                 + filter[c, f_h, f_w] \right)
+
+    where :math:`h', w'` are the top-left positions of the filter window
+    and :math:`d_h, d_w` are the dilation factors.
+
+    Parameters
+    ----------
+    input : relax.Expr
+        4-D input tensor. With NCHW layout: [batch, in_channel, in_height, in_width].
+        With NHWC layout: [batch, in_height, in_width, in_channel].
+
+    filter : relax.Expr
+        3-D filter tensor with shape [in_channel, filter_height, filter_width].
+
+    strides : Union[int, Tuple[int, int]]
+        Stride size for the filter window, or [stride_h, stride_w].
+
+    padding : Union[int, Tuple[int, ...]]
+        Padding size. Can be an int for symmetric padding,
+        [pad_top, pad_left, pad_bottom, pad_right], or 2-tuple
+        [pad_h, pad_w] which is expanded to 4-tuple.
+
+    dilations : Union[int, Tuple[int, int]]
+        Dilation size for the filter, or [dilation_h, dilation_w].
+
+    data_layout : str
+        Layout of the input data. Supported: "NCHW", "NHWC".
+
+    out_dtype : Optional[str | tvm_ffi.dtype]
+        Specifies the output data type. If None, uses the input dtype.
+
+    Returns
+    -------
+    result : relax.Expr
+        The dilated output tensor.
+    """
+    if isinstance(strides, int):
+        strides = (strides, strides)
+    if isinstance(dilations, int):
+        dilations = (dilations, dilations)
+
+    if isinstance(padding, int):
+        padding = (padding, padding, padding, padding)
+    elif isinstance(padding, tuple | list):
+        if len(padding) == 2:
+            padding = (padding[0], padding[1], padding[0], padding[1])
+        elif len(padding) == 4:
+            padding = tuple(padding)
+        else:
+            raise ValueError(f"padding must be int, 2-tuple, or 4-tuple, got {len(padding)}-tuple")
+    else:
+        raise TypeError(f"padding must be int, tuple, or list, got {type(padding)}")
+
+    return _Call(
+        "relax.image.dilation2d",
+        [input, filter],
+        attrs=_make_attrs(
+            "relax.attrs.Dilation2DAttrs",
+            strides=strides,
+            padding=padding,
+            dilations=dilations,
+            data_layout=data_layout,
+            out_dtype=(out_dtype.dtype if isinstance(out_dtype, _PrimType) else out_dtype),
+        ),
+        ty=ty,
+        loc=loc,
+    )  # type: ignore

python/tvm/relax/op/nn/__init__.py

@@ -44,6 +44,9 @@
     GroupNormAttrs,
     RMSNormAttrs,
     NLLLossAttrs,
+    CorrelationAttrs,
+    LRNAattrs,
+    DilateAttrs,
     adaptive_avg_pool1d,
     adaptive_avg_pool2d,
     adaptive_avg_pool3d,
@@ -61,7 +64,9 @@
     conv2d_transpose,
     conv3d,
     conv3d_transpose,
+    correlation,
     cross_entropy_with_logits,
+    dilate,
     dropout,
     gelu,
     gelu_tanh,
@@ -70,6 +75,7 @@
     layer_norm,
     leakyrelu,
     log_softmax,
+    lrn,
     max_pool1d,
     max_pool2d,
     max_pool3d,