# 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.
"""External function interface to MIOpen library."""
# pylint: disable-msg=C0103
import ctypes
import numpy as np
import tvm
import tvm._ffi

from tvm import te


def _get_np_int32_array_handle(arr):
    """Return a void_p handle for a numpy array

    Parameters
    ----------
    arr: numpy.NDArray
        source numpy array

    Returns
    -------
    ptr:  ctypes.c_void_p
        pointer to the data
    """
    assert arr.dtype == np.int32
    ptr = arr.ctypes.data_as(ctypes.POINTER(ctypes.c_int32))
    return ctypes.cast(ptr, ctypes.c_void_p)


def conv2d_forward(
    x,
    w,
    stride_h=1,
    stride_w=1,
    pad_h=0,
    pad_w=0,
    dilation_h=1,
    dilation_w=1,
    conv_mode=0,
    data_type=1,
    group_count=1,
):
    """Create an extern op that compute 2D convolution with MIOpen

    Parameters
    ----------
    x: Tensor
        input feature map
    w: Tensor
        convolution weight
    stride_h: int
        height stride
    stride_w: int
        width stride
    pad_h: int
        height pad
    pad_w: int
        weight pad
    dilation_h: int
        height dilation
    dilation_w: int
        width dilation
    conv_mode: int
        0: miopenConvolution
        1: miopenTranspose
    data_type: int
        0: miopenHalf (fp16)
        1: miopenFloat (fp32)
    group_count: int
        number of groups
    Returns
    -------
    y: Tensor
        The result tensor
    """
    assert 0 <= conv_mode <= 2, "0: miopenConvolution / 1: miopenTranspose / 2: miopenGroupConv"
    if group_count > 1:
        conv_mode = 2
    oshape = np.zeros((len(x.shape)), dtype=np.int32)
    xshape = x.shape
    wshape = w.shape
    setup_func = tvm._ffi.get_global_func("tvm.contrib.miopen.conv2d.setup")
    algo = setup_func(
        conv_mode,
        data_type,
        pad_h,
        pad_w,
        stride_h,
        stride_w,
        dilation_h,
        dilation_w,
        xshape[0].value,
        xshape[1].value,
        xshape[2].value,
        xshape[3].value,
        wshape[0].value,
        wshape[1].value,
        wshape[2].value,
        wshape[3].value,
        group_count,
        _get_np_int32_array_handle(oshape),
    )

    return te.extern(
        list(oshape),
        [x, w],
        lambda ins, outs: tvm.tir.call_packed(
            "tvm.contrib.miopen.conv2d.forward",
            conv_mode,
            data_type,
            pad_h,
            pad_w,
            stride_h,
            stride_w,
            dilation_h,
            dilation_w,
            algo,
            ins[0],
            ins[1],
            outs[0],
        ),
        name="y",
    )