<gh_stars>1-10
# specific language governing permissions and limitations
# under the License.
import numpy as np
import tvm
from tvm import te
from tvm import autotvm
from tvm import topi
import tvm.topi.testing
from tvm.contrib.pickle_memoize import memoize
from tvm.topi.utils import get_const_tuple

import tvm.testing

_modulated_deformable_conv2d_nchw_implement = {
    "generic": (topi.nn.modulated_deformable_conv2d_nchw, topi.generic.schedule_modulated_deformable_conv2d_nchw),
    "cuda": (topi.cuda.modulated_deformable_conv2d_nchw, topi.cuda.schedule_modulated_deformable_conv2d_nchw),
}


def check_numpy_array(a, b):
    assert a.shape == b.shape
    return np.max(a - b)

def verify_modulated_deformable_conv2d_nchw(
    batch,
    in_channel,
    in_size,
    num_filter,
    kernel,
    stride,
    padding,
    dilation=1,
    deformable_groups=1,
    groups=1,
):
    # verify_modulated_deformable_conv2d_nchw(1, 3, 16, 2, 3, 1, 1, deformable_groups=1)
    print(
        "Workload: (%d, %d, %d, %d, %d, %d, %d, %d, %d, %d)"
        % (
            batch, # 1
            in_channel, # 16
            in_size, # 7
            num_filter, # 16
            kernel, # 1
            stride, # 1
            padding, # 
            dilation,
            deformable_groups,
            groups,
        )
    )
    A = te.placeholder((batch, in_channel, in_size, in_size), name="A")
    out_size = (in_size - (kernel - 1) * dilation - 1 + 2 * padding) // stride + 1
    Offset = te.placeholder(
        (batch, deformable_groups * kernel * kernel * 2, out_size, out_size), name="offset"
    )
    W = te.placeholder((num_filter, in_channel, kernel, kernel), name="W")

    bias = te.placeholder((num_filter, 1, 1), name="bias")

    a_shape = get_const_tuple(A.shape)
    offset_shape = get_const_tuple(Offset.shape)
    w_shape = get_const_tuple(W.shape)
    bias_shape = get_const_tuple(bias.shape)
    dtype = A.dtype

    mask_shape = (batch, kernel * kernel, out_size, out_size)
    Mask = te.placeholder(mask_shape, name = "mask")


    a_np = np.random.uniform(size=a_shape).astype(dtype)
    offset_np = np.random.uniform(size = offset_shape).astype(dtype)
    mask_np = np.random.uniform(size = mask_shape).astype(dtype)
    w_np = np.random.uniform(size=w_shape).astype(dtype)

    output_shape = (batch, num_filter, out_size, out_size)
    print("output_shape: ")
    print(output_shape)


    def check_device(device):
        dev = tvm.device(device, 0)
        if not tvm.testing.device_enabled(device):
            print("Skip because %s is not enabled" % device)
            return
        print("Running on target: %s" % device)
        fcompute, fschedule = tvm.topi.testing.dispatch(device, _modulated_deformable_conv2d_nchw_implement)
        with tvm.target.Target(device):
            C = fcompute(A, Offset, Mask, W, stride, padding, dilation, deformable_groups, groups, dtype)
            s = fschedule([C])

            a = tvm.nd.array(a_np, dev)
            offset = tvm.nd.array(offset_np, dev)
            mask = tvm.nd.array(mask_np, dev)
            w = tvm.nd.array(w_np, dev)

            func = tvm.build(s, [A, Offset, Mask, W, C], device)
            c = tvm.nd.empty(output_shape, dtype="float32", device=dev)
            func(a, offset, mask, w, c)

            final_output = c.asnumpy()

            # check numpy implementaion result:
            c_np = tvm.topi.testing.modulated_deformable_conv2d_nchw_python(
                a_np, offset_np, mask_np, w_np, stride, padding, dilation, deformable_groups, groups
            )
            
            print("check the validness of tvm topi : ")
            assert check_numpy_array(final_output, c_np) <= 0.0001

    for device in ["llvm", "cuda"]:
        check_device(device)



if __name__ == "__main__":
    print("testing modulated deformable op")
    verify_modulated_deformable_conv2d_nchw(1, 3, 16, 2, 3, 1, 1, deformable_groups=1)
# specific language governing permissions and limitations
# under the License.
import numpy as np
import tvm
from tvm import te
from tvm import autotvm
from tvm import topi
import tvm.topi.testing
from tvm.contrib.pickle_memoize import memoize
from tvm.topi.utils import get_const_tuple

import tvm.testing

_modulated_deformable_conv2d_nchw_implement = {
    "generic": (topi.nn.modulated_deformable_conv2d_nchw, topi.generic.schedule_modulated_deformable_conv2d_nchw),
    "cuda": (topi.cuda.modulated_deformable_conv2d_nchw, topi.cuda.schedule_modulated_deformable_conv2d_nchw),
}


def check_numpy_array(a, b):
    assert a.shape == b.shape
    return np.max(a - b)

def verify_modulated_deformable_conv2d_nchw(
    batch,
    in_channel,
    in_size,
    num_filter,
    kernel,
    stride,
    padding,
    dilation=1,
    deformable_groups=1,
    groups=1,
):
    # verify_modulated_deformable_conv2d_nchw(1, 3, 16, 2, 3, 1, 1, deformable_groups=1)
    print(
        "Workload: (%d, %d, %d, %d, %d, %d, %d, %d, %d, %d)"
        % (
            batch, # 1
            in_channel, # 16
            in_size, # 7
            num_filter, # 16
            kernel, # 1
            stride, # 1
            padding, # 
            dilation,
            deformable_groups,
            groups,
        )
    )
    A = te.placeholder((batch, in_channel, in_size, in_size), name="A")
    out_size = (in_size - (kernel - 1) * dilation - 1 + 2 * padding) // stride + 1
    Offset = te.placeholder(
        (batch, deformable_groups * kernel * kernel * 2, out_size, out_size), name="offset"
    )
    W = te.placeholder((num_filter, in_channel, kernel, kernel), name="W")

    bias = te.placeholder((num_filter, 1, 1), name="bias")

    a_shape = get_const_tuple(A.shape)
    offset_shape = get_const_tuple(Offset.shape)
    w_shape = get_const_tuple(W.shape)
    bias_shape = get_const_tuple(bias.shape)
    dtype = A.dtype

    mask_shape = (batch, kernel * kernel, out_size, out_size)
    Mask = te.placeholder(mask_shape, name = "mask")


    a_np = np.random.uniform(size=a_shape).astype(dtype)
    offset_np = np.random.uniform(size = offset_shape).astype(dtype)
    mask_np = np.random.uniform(size = mask_shape).astype(dtype)
    w_np = np.random.uniform(size=w_shape).astype(dtype)

    output_shape = (batch, num_filter, out_size, out_size)
    print("output_shape: ")
    print(output_shape)


    def check_device(device):
        dev = tvm.device(device, 0)
        if not tvm.testing.device_enabled(device):
            print("Skip because %s is not enabled" % device)
            return
        print("Running on target: %s" % device)
        fcompute, fschedule = tvm.topi.testing.dispatch(device, _modulated_deformable_conv2d_nchw_implement)
        with tvm.target.Target(device):
            C = fcompute(A, Offset, Mask, W, stride, padding, dilation, deformable_groups, groups, dtype)
            s = fschedule([C])

            a = tvm.nd.array(a_np, dev)
            offset = tvm.nd.array(offset_np, dev)
            mask = tvm.nd.array(mask_np, dev)
            w = tvm.nd.array(w_np, dev)

            func = tvm.build(s, [A, Offset, Mask, W, C], device)
            c = tvm.nd.empty(output_shape, dtype="float32", device=dev)
            func(a, offset, mask, w, c)

            final_output = c.asnumpy()

            # check numpy implementaion result:
            c_np = tvm.topi.testing.modulated_deformable_conv2d_nchw_python(
                a_np, offset_np, mask_np, w_np, stride, padding, dilation, deformable_groups, groups
            )
            
            print("check the validness of tvm topi : ")
            assert check_numpy_array(final_output, c_np) <= 0.0001

    for device in ["llvm", "cuda"]:
        check_device(device)



if __name__ == "__main__":
    print("testing modulated deformable op")
    verify_modulated_deformable_conv2d_nchw(1, 3, 16, 2, 3, 1, 1, deformable_groups=1)
# specific language governing permissions and limitations # under the License. # verify_modulated_deformable_conv2d_nchw(1, 3, 16, 2, 3, 1, 1, deformable_groups=1) # 1 # 16 # 7 # 16 # 1 # 1 # # check numpy implementaion result: