# 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")