elemwise_helper.cuh

/**
 * \file dnn/src/cuda/elemwise_helper.cuh
 * MegEngine is Licensed under the Apache License, Version 2.0 (the "License")
 *
 * Copyright (c) 2014-2020 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 "src/common/elemwise_helper.cuh"
#include "src/cuda/utils.cuh"
#include "src/cuda/int_fastdiv.cuh"
#include "src/cuda/query_blocksize.cuh"

/*
 * please note that all arithmetics on GPU are 32-bit for best performance; this
 * limits max possible size
 */

namespace megdnn {
namespace cuda {

//! internals for element-wise
namespace elemwise_intl {
#define devfunc __device__ __forceinline__

    /*!
     * \brief get cuda launch specs for element-wise kernel
     * \param kern kernel function address
     * \param size total size of elements
     */
    void get_launch_spec(
            const void *kern, size_t size, int *grid_size, int *block_size);

    MEGDNN_NORETURN void on_bad_ndim(int ndim);

    /*!
     * \brief broadcast type
     * BCAST_x[0]x[1]...: x[i] == !stride[i]
     */
    enum BcastType {
        BCAST_OTHER,
        BCAST_1010,
        BCAST_101,
        BCAST_10,
        BCAST_01,
        BCAST_FULL
    };

    /*!
     * \brief read and write type trait for byte width integer type
     */
    template <typename ctype>
    class VectTypeTrait;

    struct __attribute__((aligned(8))) half4 {
        dt_float16 x, y, z, w;
    };

    __device__ __forceinline__ half4 make_half4(dt_float16 x, dt_float16 y,
                                                dt_float16 z, dt_float16 w) {
        half4 t;
        t.x = x, t.y = y, t.z = z, t.w = w;
        return t;
    }

    struct __attribute__((aligned(8))) bhalf4 {
        dt_bfloat16 x, y, z, w;
    };

    __device__ __forceinline__ bhalf4 make_bhalf4(dt_bfloat16 x, dt_bfloat16 y,
                                                dt_bfloat16 z, dt_bfloat16 w) {
        bhalf4 t;
        t.x = x, t.y = y, t.z = z, t.w = w;
        return t;
    }

#define INST(_ctype, _vect_type)                                               \
    template <>                                                                \
    class VectTypeTrait<_ctype> {                                              \
    public:                                                                    \
        using vect_type = _vect_type;                                          \
        static const size_t packed_size = sizeof(_vect_type) / sizeof(_ctype); \
        static __device__ __forceinline__ vect_type make_vector(_ctype x,      \
                                                                _ctype y,      \
                                                                _ctype z,      \
                                                                _ctype w) {    \
            return make_##_vect_type(as_raw(x), as_raw(y), as_raw(z),          \
                                     as_raw(w));                               \
        }                                                                      \
    }
#define as_raw(x) x
    INST(dt_int8, char4);
    INST(dt_uint8, uchar4);
    INST(dt_float32, float4);
    INST(dt_float16, half4);
    INST(dt_bfloat16, bhalf4);
    INST(dt_int32, int4);
    INST(dt_int16, short4);
#undef as_raw
#define as_raw(x) x.as_int8()
    INST(dt_qint8, char4);
#undef as_raw
#define as_raw(x) x.as_uint8()
    INST(dt_quint8, uchar4);
#undef as_raw
#define as_raw(x) x.as_int32()
    INST(dt_qint32, int4);
#undef as_raw
#undef INST

    /*!
     * \brief visitor to access an elemeent in a tensor at given logic index
     * \tparam ctype plain element ctype (i.e. ctype in DTypeTrait)
     * \tparam brdcast_mask bit mask for broadcast of params; (i.e. stride[i] is
     *      0 iff (brdcast_mask & (1<<(ndim-1-i))) is 1.
     *
     * host interface:
     *      void host_init(
     *              const TensorND &tensor, int grid_size, int block_size)
     *
     * device interface:
     *      void thread_init(uint32_t idx)
     *          called on thread entrance, with logical indexing; the index may
     *          go beyond buffer range
     *
     *      ctype* ptr()
     *          return buffer pointer; can be used by specialized OpCaller
     *
     *      void next()
     *          called before moving to next chunk on each thread
     *
     *      int offset(uint32_t idx)
     *          get physical offset from logical index
     *
     *      ctype& at(uint32_t idx)
     *          ptr()[offset(idx)]
     *
     */
    template <int ndim, typename ctype, BcastType brd_type>
    class ParamElemVisitor;

    /*!
     * \brief visitor to access vector element in a tensor at given logic index
     * \tparam ctype same as ParamElemVisitor, vect_type packed vector type of
     * element ctype (i.e. vect_type in VectTypeTrait) \tparam brdcast_mask same
     * as ParamElemVisitor
     *
     *
     * device interface:
     *      vect_type& at(uint32_t idx)
     *          ptr()[offset(idx)]
     *
     */
    template <int ndim, typename ctype, BcastType brd_type>
    class ParamVectVisitor;

    /* f{{{ ParamElemVisitor specializations */

#define PARAM_ELEM_VISITOR_COMMON_DEV      \
    devfunc ctype* ptr() { return m_ptr; } \
    devfunc ctype& at(uint32_t idx) { return m_ptr[offset(idx)]; }
#define PARAM_ELEM_VISITOR_COMMON_HOST static const int packed_size = 1;

    //! specialization for BCAST_OTHER
    template <int ndim, typename ctype>
    class ParamElemVisitor<ndim, ctype, BCAST_OTHER> {
    protected:
        ctype* __restrict m_ptr;

    private:
        int m_stride[ndim];

        //! m_shape_highdim[i] = original_shape[i + 1]
#ifdef _MSC_VER
        Uint32Fastdiv m_shape_highdim[ndim > 1 ? ndim - 1 : 1];
#else
        Uint32Fastdiv m_shape_highdim[ndim];
#endif

    public:
        static const int NDIM = ndim;
        PARAM_ELEM_VISITOR_COMMON_HOST

        void host_init(const TensorND& rv, int grid_size, int block_size);

#if MEGDNN_CC_CUDA
        devfunc void thread_init(uint32_t) {}

        devfunc void next() {}

        devfunc int offset(uint32_t idx) {
            int offset = 0;
#pragma unroll
            for (int i = ndim - 1; i >= 1; --i) {
                Uint32Fastdiv& shp = m_shape_highdim[i - 1];
                uint32_t idx_div = idx / shp;
                offset += (idx - idx_div * shp.divisor()) * m_stride[i];
                idx = idx_div;
            }
            offset += idx * m_stride[0];
            return offset;
        }

        PARAM_ELEM_VISITOR_COMMON_DEV
#endif
    };
    
    /*!
     * \brief specialization for ndim == 3 and BCAST_101
     * (for dimshuffle 'x', 0, 'x')
     *
     * visit: idx / m_shape2 % m_shape1
     */
    template <typename ctype>
    class ParamElemVisitor<3, ctype, BCAST_101> {
        StridedDivSeq2 m_shape12;
        int m_stride1;

    protected:
        ctype* __restrict m_ptr;

    public:
        static const int NDIM = 3;
        PARAM_ELEM_VISITOR_COMMON_HOST

        void host_init(const TensorND& rv, int grid_size, int block_size);

#if MEGDNN_CC_CUDA
        devfunc void thread_init(uint32_t idx) { m_shape12.device_init(idx); }

        devfunc void next() { m_shape12.next(); }

        devfunc int offset(uint32_t idx) { return m_shape12.get() * m_stride1; }

        PARAM_ELEM_VISITOR_COMMON_DEV
#endif
    };

    /*!
     * \brief specialization for ndim == 2 and BCAST_10
     *
     * visit: idx % m_shape1
     */
    template <typename ctype>
    class ParamElemVisitor<2, ctype, BCAST_10> {
        StridedDivSeq<false> m_shape1;
        int m_stride1;

    protected:
        ctype* __restrict m_ptr;

    public:
        static const int NDIM = 2;
        PARAM_ELEM_VISITOR_COMMON_HOST

        void host_init(const TensorND& rv, int grid_size, int block_size);

#if MEGDNN_CC_CUDA
        devfunc void thread_init(uint32_t idx) { m_shape1.device_init(idx); }

        devfunc void next() { m_shape1.next(); }

        devfunc int offset(uint32_t idx) { return m_shape1.r() * m_stride1; }

        PARAM_ELEM_VISITOR_COMMON_DEV
#endif
    };

    /*!
     * \brief specialization for ndim == 2 and BCAST_01
     *
     * visit: idx / shape1
     */
    template <typename ctype>
    class ParamElemVisitor<2, ctype, BCAST_01> {
        StridedDivSeq<true> m_shape1;
        int m_stride0;

    protected:
        ctype* __restrict m_ptr;

    public:
        static const int NDIM = 2;
        PARAM_ELEM_VISITOR_COMMON_HOST

        void host_init(const TensorND& rv, int grid_size, int block_size);

#if MEGDNN_CC_CUDA
        devfunc void thread_init(uint32_t idx) { m_shape1.device_init(idx); }

        devfunc void next() { m_shape1.next(); }

        devfunc int offset(uint32_t idx) { return m_shape1.q() * m_stride0; }

        PARAM_ELEM_VISITOR_COMMON_DEV
#endif
    };

    //! specialization for ndim == 1 and BCAST_FULL
    template <typename ctype>
    class ParamElemVisitor<1, ctype, BCAST_FULL> {
    protected:
        ctype* __restrict m_ptr;

    public:
        static const int NDIM = 1;
        PARAM_ELEM_VISITOR_COMMON_HOST

        void host_init(const TensorND& rv, int grid_size, int block_size);

#if MEGDNN_CC_CUDA
        devfunc void thread_init(uint32_t) {}

        devfunc void next() {}

        devfunc int offset(uint32_t idx) {
            MEGDNN_MARK_USED_VAR(idx);
            return 0;
        }

        PARAM_ELEM_VISITOR_COMMON_DEV
#endif
    };

#undef PARAM_ELEM_VISITOR_COMMON_DEV
#undef PARAM_ELEM_VISITOR_COMMON_HOST

    /* f}}} */

    /* f{{{ ParamVectVisitor specializations */

#if MEGDNN_CC_CUDA
#define DEVICE_WRAPPER(x) x
#else
#define DEVICE_WRAPPER(x)
#endif
#define INST_PARAM_VECT_VISITOR                                        \
    template <int ndim, typename ctype>                                \
    class ParamVectVisitor<ndim, ctype, _brdcast_mask>                 \
            : public ParamElemVisitor<ndim, ctype, _brdcast_mask> {    \
    public:                                                            \
        using Super = ParamElemVisitor<ndim, ctype, _brdcast_mask>;    \
        using rwtype = typename VectTypeTrait<ctype>::vect_type;       \
        static const int packed_size = sizeof(rwtype) / sizeof(ctype); \
        DEVICE_WRAPPER(devfunc rwtype& at(uint32_t idx) {              \
            return *(rwtype*)(&Super::m_ptr[Super::offset(idx)]);      \
        })                                                             \
    };
#define _brdcast_mask BCAST_OTHER
    INST_PARAM_VECT_VISITOR;
#undef _brdcast_mask
#define _brdcast_mask BCAST_01
    INST_PARAM_VECT_VISITOR;
#undef _brdcast_mask
#define _brdcast_mask BCAST_10
    INST_PARAM_VECT_VISITOR;
#undef _brdcast_mask
#define _brdcast_mask BCAST_101
    INST_PARAM_VECT_VISITOR;
#undef _brdcast_mask
#define INST_DT_IBYTE(ctype)                                                \
    template <int ndim>                                                     \
    class ParamVectVisitor<ndim, ctype, BCAST_FULL>                         \
            : public ParamElemVisitor<ndim, ctype, BCAST_FULL> {            \
    public:                                                                 \
        using Super = ParamElemVisitor<ndim, ctype, BCAST_FULL>;            \
        using rwtype = typename VectTypeTrait<ctype>::vect_type;            \
        static const int packed_size = sizeof(rwtype) / sizeof(ctype);      \
        DEVICE_WRAPPER(rwtype vect_scalar;                                  \
                       devfunc rwtype & at(uint32_t /* idx */) {            \
                           ctype v = Super::m_ptr[0];                       \
                           vect_scalar = VectTypeTrait<ctype>::make_vector( \
                                   v, v, v, v);                             \
                           return vect_scalar;                              \
                       })                                                   \
    }
    INST_DT_IBYTE(dt_int8);
    INST_DT_IBYTE(dt_uint8);
    INST_DT_IBYTE(dt_qint8);
    INST_DT_IBYTE(dt_quint8);
#undef INST_DT_IBYTE
#undef DEVICE_WRAPPER
#undef INST_PARAM_VECT_VISITOR

    /*!
     * \brief specialization for ndim == 4 and BCAST_1010
     *
     * visit: (idx % m_shape3) * m_stride3 + (idx / m_shape23 % m_shape1) *
     * m_stride1
     */
    template <typename ctype>
    class ParamVectVisitor<4, ctype, BCAST_1010> {
        StridedDivSeq2 m_shape123;
        StridedDivSeq<false> m_shape3;
        int m_stride3, m_stride1;
        ctype* __restrict m_ptr;

    public:
        static const int NDIM = 4;
        using rwtype = typename VectTypeTrait<ctype>::vect_type;
        static const int packed_size = sizeof(rwtype) / sizeof(ctype);

        void host_init(const TensorND& rv, int grid_size, int block_size);

#if MEGDNN_CC_CUDA
        devfunc void thread_init(uint32_t idx) {
            m_shape123.device_init(idx);
            m_shape3.device_init(idx);
        }

        devfunc void next() {
            m_shape123.next();
            m_shape3.next();
        }

        devfunc int offset(uint32_t idx) {
            return m_shape3.r() * m_stride3 + m_shape123.get() * m_stride1;
        }

        devfunc ctype* ptr() { return m_ptr; }
        devfunc rwtype& at(uint32_t idx) {
            return *(rwtype*)(&m_ptr[offset(idx)]);
        }
#endif
    };

    /* f}}} */


#if MEGDNN_CC_CUDA

    /* f{{{ user operator callers */

    /*
     * OpCaller is used to invoke user operator with loaded element arguments.
     *
     * device interface:
     *      void thread_init(uint32_t idx);
     *
     *      void on(uint32_t idx);
     *
     *      void next();
     */

    /*!
     * \brief call user op directly without visiting any params (i.e. arity ==
     *      0)
     */
    template<class Op>
    struct OpCallerNull {
        Op op;

        devfunc void thread_init(uint32_t) {
        }

        devfunc void on(uint32_t idx) {
            op(idx);
        }

        devfunc void next() {
        }
    };

    /*!
     * \brief call an operator whose each param are promted to the same ndim and
     *      brdcast_mask
     * \tparam PVis ParamElemVisitor class
     */
    template<class Op, int arity, class PVis>
    struct OpCallerUniform;

    //! specialization for arity == 1
    template<class Op, class PVis>
    struct OpCallerUniform<Op, 1, PVis> {
        Op op;
        PVis par[1];
        static const uint32_t packed_size = PVis::packed_size;

        devfunc void thread_init(uint32_t idx) {
            idx = idx * packed_size;
            par[0].thread_init(idx);
        }

        devfunc void on(uint32_t idx) {
            idx = idx * packed_size;
            op(idx, par[0].at(idx));
        }

        devfunc void on(uint32_t idx, uint32_t remain) {
            idx = idx * packed_size;
            if (remain >= packed_size) {
                op(idx, par[0].at(idx));
            } else {
                auto ptr0 = par[0].ptr();
                for (int i = 0; i < remain; i++) {
                    op(idx + i, ptr0[par[0].offset(idx + i)]);
                }
            }
        }

        devfunc void next() {
            par[0].next();
        }
    };
    //! specialization for arity == 2
    template<class Op, class PVis>
    struct OpCallerUniform<Op, 2, PVis> {
        Op op;
        PVis par[2];
        static const uint32_t packed_size = PVis::packed_size;

        devfunc void thread_init(uint32_t idx) {
            idx = idx * packed_size;
            par[0].thread_init(idx);
            par[1].thread_init(idx);
        }

        devfunc void on(uint32_t idx) {
            idx = idx * packed_size;
            op(idx, par[0].at(idx), par[1].at(idx));
        }

        devfunc void on(uint32_t idx, uint32_t remain) {
            idx = idx * packed_size;
            if (remain >= packed_size) {
                op(idx, par[0].at(idx), par[1].at(idx));
            } else {
                auto ptr0 = par[0].ptr();
                auto ptr1 = par[1].ptr();
                for (int i = 0; i < remain; i++) {
                    op(idx + i, ptr0[par[0].offset(idx + i)],
                       ptr1[par[1].offset(idx + i)]);
                }
            }
        }

        devfunc void next() {
            par[0].next();
            par[1].next();
        }
    };
    //! specialization for arity == 3
    template<class Op, class PVis>
    struct OpCallerUniform<Op, 3, PVis> {
        Op op;
        PVis par[3];
        static const uint32_t packed_size = PVis::packed_size;

        devfunc void thread_init(uint32_t idx) {
            idx = idx * packed_size;
            par[0].thread_init(idx);
            par[1].thread_init(idx);
            par[2].thread_init(idx);
        }

        devfunc void on(uint32_t idx) {
            idx = idx * packed_size;
            op(idx, par[0].at(idx), par[1].at(idx), par[2].at(idx));
        }

        devfunc void on(uint32_t idx, uint32_t remain) {
            idx = idx * packed_size;
            if (remain >= packed_size) {
                op(idx, par[0].at(idx), par[1].at(idx), par[2].at(idx));
            } else {
                auto ptr0 = par[0].ptr();
                auto ptr1 = par[1].ptr();
                auto ptr2 = par[2].ptr();
                for (int i = 0; i < remain; i++) {
                    op(idx + i, ptr0[par[0].offset(idx + i)],
                       ptr1[par[1].offset(idx + i)],
                       ptr2[par[2].offset(idx + i)]);
                }
            }
        }

        devfunc void next() {
            par[0].next();
            par[1].next();
            par[2].next();
        }
    };

    /*!
     * \brief call binary (i.e. arity == 2) operator with different param
     *      visitors
     */
    template <class Op, class PVis0, class PVis1>
    struct OpCallerBinary {
        Op op;
        PVis0 par0;
        PVis1 par1;
        MEGDNN_STATIC_ASSERT(PVis0::packed_size == PVis1::packed_size,
                             "vector size mismatch")

        static const uint32_t packed_size = PVis0::packed_size;

        devfunc void thread_init(uint32_t idx) {
            idx = idx * packed_size;
            par0.thread_init(idx);
            par1.thread_init(idx);
        }

        devfunc void on(uint32_t idx) {
            idx = idx * packed_size;
            op(idx, par0.at(idx), par1.at(idx));
        }

        devfunc void next() {
            par0.next();
            par1.next();
        }
    };

    /* f}}} */

    template <class OpCaller>
    __global__ void cuda_kern(OpCaller op_caller, uint32_t size) {
        uint32_t idx = blockIdx.x * blockDim.x + threadIdx.x,
                 delta = blockDim.x * gridDim.x;
        // each thread works on at most 3 elements; see get_launch_spec
        op_caller.thread_init(idx);
        if (idx < size) {
            op_caller.on(idx);
            idx += delta;
            if (idx < size) {
                op_caller.next();
                op_caller.on(idx);
                idx += delta;
                if (idx < size) {
                    op_caller.next();
                    op_caller.on(idx);
                }
            }
        }
    }

    template <class Op, int arity, class PVis>
    __global__ void cuda_kern(OpCallerUniform<Op, arity, PVis> op_caller,
                              uint32_t size) {
        constexpr uint32_t packed_size = PVis::packed_size;
        const uint32_t size_packed = DIVUP(size, packed_size);
        uint32_t idx = blockIdx.x * blockDim.x + threadIdx.x,
                 delta = blockDim.x * gridDim.x;
        if (idx < size_packed) {
            op_caller.on(idx, size - packed_size * idx);
            idx += delta;
            if (idx < size_packed) {
                op_caller.on(idx, size - packed_size * idx);
                idx += delta;
                if (idx < size_packed) {
                    op_caller.on(idx, size - packed_size * idx);
                }
            }
        }
    }

    //! invoke a user Op passed to run_elemwise
    template<class Op, typename ctype, int arity>
    class UserOpInvoker;

    /* f{{{ UserOpInvoker specializations */

    //! run op by promoting all params to same ndim
    template<class Op, typename ctype, int arity>
    class UserOpInvokerToSameNdim {
        const ElemwiseOpParamN<arity> &m_param;
        cudaStream_t m_stream;
        const Op &m_op;

        void dispatch0() {
            switch(m_param.max_ndim) {
#define cb(ndim) \
                case ndim: return dispatch1<ndim>();
                MEGDNN_FOREACH_TENSOR_NDIM(cb)
#undef cb
            }
            on_bad_ndim(m_param.max_ndim);
        }

        template <int ndim>
        void dispatch1() {
            typedef OpCallerUniform<
                    Op, arity,
                    ParamElemVisitor<ndim, ctype, BCAST_OTHER>>
                    Caller;
            size_t size = m_param.size;
            int grid_size, block_size;
            void (*fptr)(Caller, uint32_t) = cuda_kern<Caller>;
            get_launch_spec(reinterpret_cast<const void*>(fptr), size,
                            &grid_size, &block_size);

            Caller caller;
            caller.op = m_op;
            for (int i = 0; i < arity; ++i)
                caller.par[i].host_init(m_param[i], grid_size, block_size);
            (*fptr)<<<grid_size, block_size, 0, m_stream>>>(caller, size);
            after_kernel_launch();
        }

    public:
        UserOpInvokerToSameNdim(const ElemwiseOpParamN<arity>& param,
                                cudaStream_t stream, const Op& op)
                : m_param(param), m_stream(stream), m_op(op) {
            dispatch0();
            }
    };

    template <class Op, typename ctype, int arity>
    class UserOpInvokerToSameNdimIByteHelper {
    public:
        UserOpInvokerToSameNdimIByteHelper(const ElemwiseOpParamN<arity>& param,
                                           cudaStream_t stream, const Op& op)
                : m_rw_size(param.size),
                  m_param(param),
                  m_stream(stream),
                  m_op(op) {
            if (!try_vect_load_store_contiguous() && !try_vect_load_store()) {
                dispatch0();
            }
        }

    private:
        const ElemwiseOpParamN<arity>& m_param;
        size_t m_rw_size;
        cudaStream_t m_stream;
        const Op& m_op;
        using vect_type = typename VectTypeTrait<ctype>::vect_type;
        static const size_t packed_size = VectTypeTrait<ctype>::packed_size;

        void dispatch0() {
            switch (m_param.max_ndim) {
#define cb(ndim) \
    case ndim:   \
        return dispatch1<ndim>();
                MEGDNN_FOREACH_TENSOR_NDIM(cb)
#undef cb
            }
            on_bad_ndim(m_param.max_ndim);
        }

        void dispatch0_vect() {
            switch (m_param.max_ndim) {
#define cb(ndim) \
    case ndim:   \
        return dispatch1_vect<ndim>();
                MEGDNN_FOREACH_TENSOR_NDIM(cb)
#undef cb
            }
            on_bad_ndim(m_param.max_ndim);
        }

        void dispatch_contiguous() {
            typedef ParamVectVisitor<1, ctype, BCAST_OTHER> PVis;
            typedef OpCallerUniform<Op, arity, PVis> Caller;
            size_t size = m_rw_size;
            int grid_size, block_size;
            void (*fptr)(Caller, uint32_t) = cuda_kern<Op, arity, PVis>;
            get_launch_spec(reinterpret_cast<const void*>(fptr), size,
                            &grid_size, &block_size);

            Caller caller;
            caller.op = m_op;
            for (int i = 0; i < arity; ++i)
                caller.par[i].host_init(m_param[i], grid_size, block_size);
            (*fptr)<<<grid_size, block_size, 0, m_stream>>>(caller,
                                                            m_param.size);
            after_kernel_launch();
        }

        template <int ndim>
        void dispatch1() {
            typedef ParamElemVisitor<ndim, ctype, BCAST_OTHER> PVis;
            typedef OpCallerUniform<Op, arity, PVis> Caller;
            size_t size = m_rw_size;
            int grid_size, block_size;
            void (*fptr)(Caller, uint32_t) = cuda_kern<Caller>;
            get_launch_spec(reinterpret_cast<const void*>(fptr), size,
                            &grid_size, &block_size);
            Caller caller;
            caller.op = m_op;
            for (int i = 0; i < arity; ++i)
                caller.par[i].host_init(m_param[i], grid_size, block_size);
            (*fptr)<<<grid_size, block_size, 0, m_stream>>>(caller, size);
            after_kernel_launch();
        }

        template <int ndim>
        void dispatch1_vect() {
            typedef ParamVectVisitor<ndim, ctype, BCAST_OTHER> PVis;
            typedef OpCallerUniform<Op, arity, PVis> Caller;
            size_t size = m_rw_size;
            int grid_size, block_size;
            void (*fptr)(Caller, uint32_t) = cuda_kern<Caller>;
            get_launch_spec(reinterpret_cast<const void*>(fptr), size,
                            &grid_size, &block_size);
            Caller caller;
            caller.op = m_op;
            for (int i = 0; i < arity; ++i)
                caller.par[i].host_init(m_param[i], grid_size, block_size);
            (*fptr)<<<grid_size, block_size, 0, m_stream>>>(caller, size);
            after_kernel_launch();
        }

        bool try_vect_load_store() {
            auto try_last_contig = [](const TensorLayout& layout) {
                return layout.stride[layout.ndim - 1] == 1 &&
                       layout[layout.ndim - 1] % packed_size == 0;
            };
            /*
             * \NOTE: remove try_scalar() to adapt multi-type tenary op
             */
            for (int i = 0; i < arity; ++i) {
                if (!try_last_contig(m_param[i].layout)) return false;
            }
            m_rw_size /= packed_size;
            dispatch0_vect();
            return true;
        }

        bool try_vect_load_store_contiguous() {
            auto try_contig = [](const TensorLayout& layout) {
                return (layout.is_contiguous());
            };
            for (int i = 0; i < arity; ++i) {
                if (!try_contig(m_param[i].layout))
                    return false;
            }
            m_rw_size = DIVUP(m_rw_size, packed_size);
            dispatch_contiguous();
            return true;
        }
    };

#define INST_DT_IBYTE(ctype)                                                \
    template <class Op, int arity>                                          \
    class UserOpInvokerToSameNdim<Op, ctype, arity>                         \
            : public UserOpInvokerToSameNdimIByteHelper<Op, ctype, arity> { \
        using Super = UserOpInvokerToSameNdimIByteHelper<Op, ctype, arity>; \
                                                                            \
    public:                                                                 \
        UserOpInvokerToSameNdim(const ElemwiseOpParamN<arity>& param,       \
                                cudaStream_t stream, const Op& op)          \
                : Super{param, stream, op} {}                               \
    }
    INST_DT_IBYTE(dt_int8);
    INST_DT_IBYTE(dt_uint8);
    INST_DT_IBYTE(dt_qint8);
    INST_DT_IBYTE(dt_quint8);
#undef INST_DT_IBYTE

    //! implement general case by UserOpInvokerToSameNdim
    template<class Op, typename ctype, int arity>
    class UserOpInvoker: public UserOpInvokerToSameNdim<Op, ctype, arity> {
        public:
            UserOpInvoker(
                    const ElemwiseOpParamN<arity> &param,
                    cudaStream_t stream,
                    const Op &op):
                UserOpInvokerToSameNdim<Op, ctype, arity>(param, stream, op)
            {
            }
    };

    //! specialization for arity == 0
    template<class Op, typename ctype>
    class UserOpInvoker<Op, ctype, 0> {
        public:
            UserOpInvoker(
                    const ElemwiseOpParamN<0> &param,
                    cudaStream_t stream,
                    const Op &op) {
                size_t size = param.size;
                typedef OpCallerNull<Op> Caller;
                Caller caller;
                caller.op = op;
                int grid_size, block_size;
                void (*fptr)(Caller, uint32_t) = cuda_kern<Caller>;
                get_launch_spec(reinterpret_cast<const void*>(fptr), size,
                                &grid_size, &block_size);
                (*fptr) <<< grid_size, block_size, 0, stream >>> (caller, size);
                after_kernel_launch();
            }
    };

#define DEFINE_BRDCAST_DISPATCH_RECEIVERS(_cb_header, _cb_dispatch, _stride) \
    _cb_header(1) { \
        const ptrdiff_t *stride = _stride; \
        if (!stride[0]) { \
            return _cb_dispatch(1, BCAST_FULL); \
        } \
        _cb_dispatch(1, BCAST_OTHER); \
    } \
    _cb_header(2) { \
        const ptrdiff_t *stride = _stride; \
        if (!stride[0] && stride[1]) { \
            return _cb_dispatch(2, BCAST_10); \
        } \
        if (stride[0] && !stride[1]) { \
            return _cb_dispatch(2, BCAST_01); \
        } \
        _cb_dispatch(2, BCAST_OTHER); \
    } \
    _cb_header(3) { \
        const ptrdiff_t *stride = _stride; \
        if (!stride[0] && stride[1] && !stride[2]) { \
            return _cb_dispatch(3, BCAST_101); \
        } \
        _cb_dispatch(3, BCAST_OTHER); \
    }

    //! specialization for binary opr
    template<class Op, typename ctype>
    class UserOpInvoker<Op, ctype, 2> {
        bool m_invoked;
        const ElemwiseOpParamN<2> &m_param;
        cudaStream_t m_stream;
        const Op &m_op;

        void fallback() {
            megdnn_assert(!m_invoked);
            UserOpInvokerToSameNdim<Op, ctype, 2>(m_param, m_stream, m_op);
            m_invoked = true;
        }

        void dispatch0() {
            switch(m_param[0].layout.ndim) {
#define cb(ndim) \
                case ndim: return dispatch1_##ndim();
                MEGDNN_FOREACH_TENSOR_NDIM_SMALL(cb)
#undef cb
            }
            fallback();
        }

#define cb_header(ndim) void dispatch1_##ndim()
#define cb_dispatch(ndim, brdcast_mask) \
        dispatch2<ParamElemVisitor<ndim, ctype, brdcast_mask> >()
DEFINE_BRDCAST_DISPATCH_RECEIVERS(cb_header, cb_dispatch,
        m_param[0].layout.stride)
#undef cb_header
#undef cb_dispatch


        template<class PVis0>
        void dispatch2() {
            switch(m_param[1].layout.ndim) {
#define cb(ndim) \
                case ndim: return dispatch3_##ndim<PVis0>();
                MEGDNN_FOREACH_TENSOR_NDIM_SMALL(cb)
#undef cb
            }
            fallback();
        }

#define cb_header(ndim) \
    template<class PVis0> \
    void dispatch3_##ndim()
#define cb_dispatch(ndim, brdcast_mask) \
        do_run<PVis0, ParamElemVisitor<ndim, ctype, brdcast_mask> >()
DEFINE_BRDCAST_DISPATCH_RECEIVERS(cb_header, cb_dispatch,
        m_param[1].layout.stride)
#undef cb_header
#undef cb_dispatch

        template<class PVis0, class PVis1>
        void do_run() {
            megdnn_assert(!m_invoked);
            m_invoked = true;
            typedef OpCallerBinary<Op, PVis0, PVis1> Caller;
            int grid_size, block_size;
            void (*fptr)(Caller, uint32_t) = cuda_kern<Caller>;
            size_t size = m_param.size;
            get_launch_spec(reinterpret_cast<const void*>(fptr),
                    size, &grid_size, &block_size);
            Caller caller;
            caller.op = m_op;
            caller.par0.host_init(m_param[0], grid_size, block_size);
            caller.par1.host_init(m_param[1], grid_size, block_size);
            (*fptr) <<< grid_size, block_size, 0, m_stream >>> (caller, size);
            after_kernel_launch();
        }

        public:
            UserOpInvoker(const ElemwiseOpParamN<2> &param, cudaStream_t stream,
                    const Op &op):
                m_param(param), m_stream(stream), m_op(op)
            {
                m_invoked = false;
                dispatch0();
                megdnn_assert(m_invoked);
            }
    };

#define DEFINE_VECT_BRDCAST_DISPATCH_RECEIVERS(_cb_header, _cb_dispatch, \
                                               _stride)                  \
    DEFINE_BRDCAST_DISPATCH_RECEIVERS(_cb_header, _cb_dispatch, _stride) \
    _cb_header(4) {                                                      \
        const ptrdiff_t* stride = _stride;                               \
        if (!stride[0] && stride[1] && !stride[2] && stride[3]) {        \
            return _cb_dispatch(4, BCAST_1010);                          \
        }                                                                \
        _cb_dispatch(4, BCAST_OTHER);                                    \
    }

    template <class Op, typename ctype>
    class UserOpInvokerBinaryIByteHelper {
    private:
        bool m_invoked;
        size_t m_rw_size;
        const ElemwiseOpParamN<2>& m_param;
        cudaStream_t m_stream;
        const Op& m_op;
        using vect_type = typename VectTypeTrait<ctype>::vect_type;
        static const size_t packed_size = VectTypeTrait<ctype>::packed_size;
        bool try_vect_load_store() {
            auto try_last_contig_or_scalar = [](const TensorLayout& layout) {
                return (layout.stride[layout.ndim - 1] == 1 &&
                        layout[layout.ndim - 1] % packed_size == 0) ||
                       (layout.ndim == 1 && layout.stride[0] == 0);
            };
            for (int i = 0; i < 2; ++i) {
                if (!try_last_contig_or_scalar(m_param[i].layout))
                    return false;
            }
            m_rw_size /= packed_size;
            dispatch0_vect();
            return true;
        }

        bool try_vect_load_store_contiguous() {
            auto try_contig = [](const TensorLayout& layout) {
                return (layout.is_contiguous());
            };
            for (int i = 0; i < 2; ++i) {
                if (!try_contig(m_param[i].layout))
                    return false;
            }
            m_rw_size = DIVUP(m_rw_size, packed_size);
            dispatch_contiguous();
            return true;
        }

        void dispatch0() {
            switch (m_param[0].layout.ndim) {
#define cb(ndim) \
    case ndim:   \
        return dispatch1_##ndim();
                MEGDNN_FOREACH_TENSOR_NDIM_SMALL(cb)
#undef cb
            }
            fallback();
        }

        void dispatch0_vect() {
            switch (m_param[0].layout.ndim) {
#define cb(ndim) \
    case ndim:   \
        return dispatch1_vect_##ndim();
                MEGDNN_FOREACH_TENSOR_NDIM_SMALL(cb)
#undef cb
                case 4:
                    return dispatch1_vect_4();
            }
            fallback();
        }

        void dispatch_contiguous() {
            m_invoked = true;
            typedef ParamVectVisitor<1, ctype, BCAST_OTHER> PVis;
            typedef OpCallerUniform<Op, 2, PVis> Caller;
            size_t size = m_rw_size;
            int grid_size, block_size;
            void (*fptr)(Caller, uint32_t) = cuda_kern<Op, 2, PVis>;
            get_launch_spec(reinterpret_cast<const void*>(fptr), size,
                            &grid_size, &block_size);

            Caller caller;
            caller.op = m_op;
            for (int i = 0; i < 2; ++i)
                caller.par[i].host_init(m_param[i], grid_size, block_size);
            (*fptr)<<<grid_size, block_size, 0, m_stream>>>(caller,
                                                            m_param.size);
            after_kernel_launch();
        }

        void fallback() {
            megdnn_assert(!m_invoked);
            UserOpInvokerToSameNdim<Op, ctype, 2>(m_param, m_stream, m_op);
            m_invoked = true;
        }

#define cb_header(ndim) void dispatch1_##ndim()
#define cb_dispatch(ndim, brdcast_mask) \
    dispatch2<ParamElemVisitor<ndim, ctype, brdcast_mask>>()
        DEFINE_BRDCAST_DISPATCH_RECEIVERS(cb_header, cb_dispatch,
                                          m_param[0].layout.stride)
#undef cb_header
#undef cb_dispatch

#define cb_header(ndim) void dispatch1_vect_##ndim()
#define cb_dispatch(ndim, brdcast_mask) \
    dispatch2_vect<ParamVectVisitor<ndim, ctype, brdcast_mask>>()
        DEFINE_VECT_BRDCAST_DISPATCH_RECEIVERS(cb_header, cb_dispatch,
                                               m_param[0].layout.stride)
#undef cb_header
#undef cb_dispatch

        template <class PVis0>
        void dispatch2() {
            switch (m_param[1].layout.ndim) {
#define cb(ndim) \
    case ndim:   \
        return dispatch3_##ndim<PVis0>();
                MEGDNN_FOREACH_TENSOR_NDIM_SMALL(cb)
#undef cb
            }
            fallback();
        }

        template <class PVis0>
        void dispatch2_vect() {
            switch (m_param[1].layout.ndim) {
#define cb(ndim) \
    case ndim:   \
        return dispatch3_vect_##ndim<PVis0>();
                MEGDNN_FOREACH_TENSOR_NDIM_SMALL(cb)
#undef cb
                case 4:
                    return dispatch3_vect_4<PVis0>();
            }
            fallback();
        }

#define cb_header(ndim)    \
    template <class PVis0> \
    void dispatch3_##ndim()
#define cb_dispatch(ndim, brdcast_mask) \
    do_run<PVis0, ParamElemVisitor<ndim, ctype, brdcast_mask>>()
        DEFINE_BRDCAST_DISPATCH_RECEIVERS(cb_header, cb_dispatch,
                                          m_param[1].layout.stride)
#undef cb_header
#undef cb_dispatch

#define cb_header(ndim)    \
    template <class PVis0> \
    void dispatch3_vect_##ndim()
#define cb_dispatch(ndim, brdcast_mask) \
    do_run<PVis0, ParamVectVisitor<ndim, ctype, brdcast_mask>>()
        DEFINE_VECT_BRDCAST_DISPATCH_RECEIVERS(cb_header, cb_dispatch,
                                               m_param[1].layout.stride)
#undef cb_header
#undef cb_dispatch

        template <class PVis0, class PVis1>
        void do_run() {
            megdnn_assert(!m_invoked);
            m_invoked = true;
            typedef OpCallerBinary<Op, PVis0, PVis1> Caller;
            int grid_size, block_size;
            void (*fptr)(Caller, uint32_t) = cuda_kern<Caller>;
            size_t size = m_rw_size;
            get_launch_spec(reinterpret_cast<const void*>(fptr), size,
                            &grid_size, &block_size);
            Caller caller;
            caller.op = m_op;
            caller.par0.host_init(m_param[0], grid_size, block_size);
            caller.par1.host_init(m_param[1], grid_size, block_size);
            (*fptr)<<<grid_size, block_size, 0, m_stream>>>(caller, size);
            after_kernel_launch();
        }

    public:
        UserOpInvokerBinaryIByteHelper(const ElemwiseOpParamN<2>& param,
                                       cudaStream_t stream, const Op& op)
                : m_rw_size(param.size),
                  m_param(param),
                  m_stream(stream),
                  m_op(op) {
            m_invoked = false;
            if (!try_vect_load_store_contiguous() && !try_vect_load_store()) {
                dispatch0();
            }
            megdnn_assert(m_invoked);
        }
    };

#define INST_DT_IBYTE(ctype)                                                 \
    template <class Op>                                                      \
    class UserOpInvoker<Op, ctype, 2>                                        \
            : public UserOpInvokerBinaryIByteHelper<Op, ctype> {             \
        using Super = UserOpInvokerBinaryIByteHelper<Op, ctype>;             \
                                                                             \
    public:                                                                  \
        UserOpInvoker(const ElemwiseOpParamN<2>& param, cudaStream_t stream, \
                      const Op& op)                                          \
                : Super{param, stream, op} {}                                \
    }
    INST_DT_IBYTE(dt_int8);
    INST_DT_IBYTE(dt_uint8);
    INST_DT_IBYTE(dt_qint8);
    INST_DT_IBYTE(dt_quint8);
#undef INST_DT_IBYTE
#endif

#undef DEFINE_BRDCAST_DISPATCH_RECEIVERS
#undef DEFINE_VECT_BRDCAST_DISPATCH_RECEIVERS

    /* f}}} */

#undef devfunc
} // namespace elemwise_intl

/*!
 * \brief general element-wise kernel launcher
 *
 * \tparam arity number of params for the operator
 * \param param param values for the operator; must have been initialized (i.e.
 *      by calling ElemwiseOpParamN::init_from_given_tensor). The params
 *      can have arbitrary layouts, as long as they share the same total number
 *      of elements.
 * \param op callable with a signature compatible with
 *      `void op(uint32_t idx, ctype& param0, ..., ctype& param[arity - 1])`
 *      if arity == 0, there is only an `idx` input
 *      if ctype=dt_int8, dt_uint8, dt_qint8, dt_quint8, a signature compatible
 * with `void op(uint32_t idx, vect_type& param0, ..., ctype& param[arity - 1])`
 * should be implemented
 */
template <class Op, typename ctype, int arity>
void run_elemwise(const ElemwiseOpParamN<arity>& param, cudaStream_t stream,
                  const Op& op = Op());

#if MEGDNN_CC_CUDA
template<class Op, typename ctype, int arity>
void run_elemwise(
        const ElemwiseOpParamN<arity> &param, cudaStream_t stream,
        const Op &op) {
    param.assert_initialized();
    elemwise_intl::UserOpInvoker<Op, ctype, arity>(param, stream, op);
}

/*!
 * \brief explicit instantialization of run_elemwise for given template params;
 *      used in .cu files, so corresponding run_elemwise can be called from .cpp
 */
#define INST_RUN_ELEMWISE(Op, ctype, arity) \
template void run_elemwise<Op, ctype, arity>( \
        const ElemwiseOpParamN<arity>&, cudaStream_t, const Op&)

#endif

} // namespace cuda
} // namespace megdnn

// vim: ft=cpp syntax=cpp.doxygen foldmethod=marker foldmarker=f{{{,f}}}