megenginelite-sys 1.8.2

/**
 * \file dnn/src/naive/warp_perspective/opr_impl.h
 * MegEngine is Licensed under the Apache License, Version 2.0 (the "License")
 *
 * Copyright (c) 2014-2021 Megvii Inc. All rights reserved.
 *
 * Unless required by applicable law or agreed to in writing,
 * software distributed under the License is distributed on an
 * "AS IS" BASIS, WITHOUT ARRANTIES OR CONDITIONS OF ANY KIND, either express or
 * implied.
 */
#pragma once
#include "megdnn/oprs.h"
#include "src/common/utils.h"

namespace megdnn {
namespace naive {

class WarpPerspectiveForwardImpl : public WarpPerspectiveForward {
protected:
    using Format = Param::Format;
    template <typename ctype, typename mtype>
    struct KernParam {
        Format format;
        BorderMode bmode;
        float border_val;
        size_t n_src, n_mat, c, ih, iw, oh, ow;
        DType src_dtype, dst_dtype;
        RefPtr src_ptr, mat_ptr, dst_ptr;
        RefPtr midx_ptr;  //!< can be null
        Workspace workspace;

        static KernParam from_tensors(
                Format format, BorderMode bmode, float border_val,
                _megdnn_tensor_in src, _megdnn_tensor_in mat, _megdnn_tensor_in mat_idx,
                _megdnn_tensor_out dst, _megdnn_workspace workspace) {
            KernParam ret;
            ret.format = format;
            ret.bmode = bmode;
            ret.border_val = border_val;
            ret.n_src = src.layout.shape[0];
            ret.src_dtype = src.layout.dtype;
            ret.dst_dtype = dst.layout.dtype;

            if (mat_idx.raw_ptr()) {
                megdnn_assert(mat_idx.layout.ndim == 1);
                ret.n_mat = mat_idx.layout.shape[0];
                ret.midx_ptr = mat_idx.get_ref_ptr();
            } else {
                megdnn_assert(mat_idx.layout.ndim == 0);
                ret.n_mat = ret.n_src;
                ret.midx_ptr = nullptr;
            }

            if (format == Format::NCHW || format == Format::NCHW_NCHW4_IC_SMALL) {
                ret.c = src.layout.shape[1];
                ret.ih = src.layout.shape[2];
                ret.iw = src.layout.shape[3];
                ret.oh = dst.layout.shape[2];
                ret.ow = dst.layout.shape[3];
            } else if (format == Format::NHWC) {
                ret.c = src.layout.shape[3];
                ret.ih = src.layout.shape[1];
                ret.iw = src.layout.shape[2];
                ret.oh = dst.layout.shape[1];
                ret.ow = dst.layout.shape[2];
            } else if (
                    format == Format::NHWC_NCHW ||
                    format == Format::NHWC_NCHW4_IC_SMALL) {
                ret.c = src.layout.shape[3];
                ret.ih = src.layout.shape[1];
                ret.iw = src.layout.shape[2];
                ret.oh = dst.layout.shape[2];
                ret.ow = dst.layout.shape[3];
            } else if (format == Format::NCHW4) {
                ret.c = src.layout.shape[1] * 4;
                ret.ih = src.layout.shape[2];
                ret.iw = src.layout.shape[3];
                ret.oh = dst.layout.shape[2];
                ret.ow = dst.layout.shape[3];
            } else if (format == Format::NCHW64) {
                ret.c = src.layout.shape[1] * 64;
                ret.ih = src.layout.shape[2];
                ret.iw = src.layout.shape[3];
                ret.oh = dst.layout.shape[2];
                ret.ow = dst.layout.shape[3];
            } else {
                megdnn_assert(format == Format::NHWCD4);
                ret.c = src.layout.shape[2] * 4;
                ret.ih = src.layout.shape[1];
                ret.iw = src.layout.shape[3];
                ret.oh = dst.layout.shape[1];
                ret.ow = dst.layout.shape[3];
            }

            if ((src.layout.dtype.enumv() == DTypeEnum::Float32 ||
                 DNN_FLOAT16_SELECT(
                         (src.layout.dtype.enumv() == DTypeEnum::Float16 ||
                          src.layout.dtype.enumv() == DTypeEnum::BFloat16),
                         false) ||
                 src.layout.dtype.enumv() == DTypeEnum::Int8 ||
                 src.layout.dtype.enumv() == DTypeEnum::Uint8 ||
                 src.layout.dtype.enumv() == DTypeEnum::QuantizedS8 ||
                 src.layout.dtype.enumv() == DTypeEnum::Quantized8Asymm) &&
                (src.layout.dtype == dst.layout.dtype)) {
                ret.src_ptr = src.get_ref_ptr();
                ret.mat_ptr = mat.get_ref_ptr();
                ret.dst_ptr = dst.get_ref_ptr();
            } else if (
                    src.layout.dtype.enumv() == DTypeEnum::QuantizedS8 ||
                    src.layout.dtype.enumv() == DTypeEnum::QuantizedS4 ||
                    src.layout.dtype.enumv() == DTypeEnum::Quantized4Asymm) {
                ret.src_ptr = src.get_ref_ptr();
                ret.mat_ptr = mat.get_ref_ptr();
                ret.dst_ptr = dst.get_ref_ptr();
            } else if (
                    (src.layout.dtype.enumv() == DTypeEnum::Uint8 ||
                     src.layout.dtype.enumv() == DTypeEnum::Quantized8Asymm) &&
                    src.layout.dtype.enumv() != dst.layout.dtype.enumv()) {
                ret.src_ptr = src.get_ref_ptr();
                ret.mat_ptr = mat.get_ref_ptr();
                ret.dst_ptr = dst.get_ref_ptr();
            } else {
                ret.src_ptr = nullptr;
                ret.mat_ptr = nullptr;
                ret.dst_ptr = nullptr;
            }
            ret.workspace = workspace;
            return ret;
        }
    };

    // ctype: C type of input data type.
    // mtype: C type of transformation matrix data type.
    template <typename ctype, typename mtype>
    void kern_naive(const KernParam<ctype, mtype>& kern_param, size_t task_id);

public:
    using WarpPerspectiveForward::WarpPerspectiveForward;
    void exec(
            _megdnn_tensor_in src, _megdnn_tensor_in mat, _megdnn_tensor_in mat_idx,
            _megdnn_tensor_out dst, _megdnn_workspace workspace) override;
    size_t get_workspace_in_bytes(
            const TensorLayout&, const TensorLayout&, const TensorLayout&,
            const TensorLayout&) override {
        return 0;
    }

private:
    template <typename ctype, typename mtype>
    void kern_naive_nhwcd4(const KernParam<ctype, mtype>& kern_param, size_t task_id);
    template <typename ctype, typename mtype>
    void kern_naive_int4(const KernParam<ctype, mtype>& kern_param, size_t task_id);
    template <typename ctype, typename dst_ctype, typename mtype>
    void kern_naive_dimshuffle_typecvt(
            const KernParam<ctype, mtype>& kern_param, size_t task_id);
};

class WarpPerspectiveBackwardDataImpl : public WarpPerspectiveBackwardData {
protected:
    template <typename ctype, typename mtype>
    struct KernParam {
        size_t n_src, n_mat, c, ih, iw, oh, ow;
        RefPtr grad_ptr, diff_ptr;
        RefPtr mat_ptr;
        RefPtr midx_ptr;

        static KernParam from_tensors(
                _megdnn_tensor_in mat, _megdnn_tensor_in mat_idx,
                _megdnn_tensor_in diff, _megdnn_tensor_out grad) {
            KernParam ret;
            ret.n_src = grad.layout.shape[0], ret.c = grad.layout.shape[1];
            ret.ih = grad.layout.shape[2], ret.iw = grad.layout.shape[3];
            ret.oh = diff.layout.shape[2], ret.ow = diff.layout.shape[3];
            ret.diff_ptr = diff.get_ref_ptr();
            ret.mat_ptr = mat.get_ref_ptr();
            ret.grad_ptr = grad.get_ref_ptr();
            if (mat_idx.raw_ptr()) {
                megdnn_assert(mat_idx.layout.ndim == 1);
                ret.n_mat = mat_idx.layout.shape[0];
                ret.midx_ptr = mat_idx.get_ref_ptr();
            } else {
                megdnn_assert(mat_idx.layout.ndim == 0);
                ret.n_mat = ret.n_src;
                ret.midx_ptr = nullptr;
            }
            return ret;
        }
    };

public:
    using WarpPerspectiveBackwardData::WarpPerspectiveBackwardData;
    void exec(
            _megdnn_tensor_in mat, _megdnn_tensor_in mat_idx, _megdnn_tensor_in diff,
            _megdnn_tensor_out grad, _megdnn_workspace workspace) override;
    size_t get_workspace_in_bytes(
            const TensorLayout&, const TensorLayout&, const TensorLayout&,
            const TensorLayout&) override {
        return 0;
    }

private:
    template <typename ctype, typename mtype>
    void kern_naive(const KernParam<ctype, mtype>& kern_param);
};

class WarpPerspectiveBackwardMatImpl : public WarpPerspectiveBackwardMat {
protected:
    template <typename ctype, typename mtype>
    struct KernParam {
        size_t n_src, n_mat, c, ih, iw, oh, ow;
        RefPtr src_ptr, grad_ptr, diff_ptr;
        RefPtr mat_ptr;
        RefPtr midx_ptr;
        float border_val;

        static KernParam from_tensors(
                float border_val_, _megdnn_tensor_in src, _megdnn_tensor_in mat,
                _megdnn_tensor_in mat_idx, _megdnn_tensor_in diff,
                _megdnn_tensor_out grad) {
            KernParam ret;
            ret.border_val = border_val_;
            ret.n_src = src.layout.shape[0], ret.c = src.layout.shape[1];
            ret.ih = src.layout.shape[2], ret.iw = src.layout.shape[3];
            ret.oh = diff.layout.shape[2], ret.ow = diff.layout.shape[3];
            ret.src_ptr = src.get_ref_ptr();
            ret.diff_ptr = diff.get_ref_ptr();
            ret.mat_ptr = mat.get_ref_ptr();
            ret.grad_ptr = grad.get_ref_ptr();
            if (mat_idx.raw_ptr()) {
                megdnn_assert(mat_idx.layout.ndim == 1);
                ret.n_mat = mat_idx.layout.shape[0];
                ret.midx_ptr = mat_idx.get_ref_ptr();
            } else {
                megdnn_assert(mat_idx.layout.ndim == 0);
                ret.n_mat = ret.n_src;
                ret.midx_ptr = nullptr;
            }
            return ret;
        }
    };

public:
    using WarpPerspectiveBackwardMat::WarpPerspectiveBackwardMat;
    void exec(
            _megdnn_tensor_in src, _megdnn_tensor_in mat, _megdnn_tensor_in mat_idx,
            _megdnn_tensor_in diff, _megdnn_tensor_out grad,
            _megdnn_workspace workspace) override;
    size_t get_workspace_in_bytes(
            const TensorLayout&, const TensorLayout&, const TensorLayout&,
            const TensorLayout&, const TensorLayout&) override {
        return 0;
    }

private:
    template <typename ctype, typename mtype>
    void kern_naive(const KernParam<ctype, mtype>& kern_param);
};

#define UNPACK_WARP_PERSPECTIVE_FWD_KERN_PARAM(p)                                    \
    auto N_SRC = p.n_src, N_MAT = p.n_mat, C = p.c, IH = p.ih, IW = p.iw, OH = p.oh, \
         OW = p.ow;                                                                  \
    auto sptr = static_cast<const ctype*>(p.src_ptr.get_ptr());                      \
    auto mptr = static_cast<const mtype*>(p.mat_ptr.get_ptr());                      \
    auto dptr = static_cast<ctype*>(p.dst_ptr.get_ptr());                            \
    auto midx_ptr = static_cast<int*>(p.midx_ptr.get_ptr());                         \
    auto bmode = p.bmode;                                                            \
    float border_val = p.border_val

}  // namespace naive
}  // namespace megdnn

// vim: syntax=cpp.doxygen