MegEngine/dnn/include/megdnn/oprs/utils.h

/**
 * \file dnn/include/megdnn/oprs/utils.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/internal/opr_header_prologue.h"

namespace megdnn {

//! base class for random number generators
class RNGBase: public OperatorBase {
    DEF_OPR_IMPL_CTOR(RNGBase, OperatorBase);
    public:
        virtual void exec(_megdnn_tensor_out dst,
                _megdnn_workspace workspace) = 0;
        virtual size_t get_workspace_in_bytes(const TensorLayout &dst) = 0;
    protected:
        virtual void check_exec(const TensorLayout &dst, size_t workspace_in_bytes) = 0;
};

//! sample from poisson distribution
class PoissonRNG: public OperatorBase {
    DEF_OPR_IMPL(PoissonRNG, OperatorBase, 1, 1);
    DEF_OPR_PARAM(PoissonRNG);
    public:
        virtual void exec(_megdnn_tensor_in lam, 
                  _megdnn_tensor_out dst,
                  _megdnn_workspace workspace) = 0;
        virtual size_t get_workspace_in_bytes(const TensorLayout &lam, 
                const TensorLayout &dst) = 0;
    protected:
        void check_exec(const TensorLayout &lam, const TensorLayout &dst, 
                    size_t workspace_in_bytes);
};

//! sample from beta distribution
class BetaRNG: public OperatorBase {
    DEF_OPR_IMPL(BetaRNG, OperatorBase, 2, 1);
    DEF_OPR_PARAM(BetaRNG);
    public:
        virtual void exec(_megdnn_tensor_in alpha, 
                        _megdnn_tensor_in beta, 
                        _megdnn_tensor_out dst,
                        _megdnn_workspace workspace) = 0;
        virtual size_t get_workspace_in_bytes(const TensorLayout &alpha, 
                    const TensorLayout &beta, const TensorLayout &dst) = 0;
    protected:
        void check_exec(const TensorLayout &alpha, const TensorLayout &beta, 
                        const TensorLayout &dst, size_t workspace_in_bytes);
};

//! sample from gamma distribution
class GammaRNG: public OperatorBase {
    DEF_OPR_IMPL(GammaRNG, OperatorBase, 2, 1);
    DEF_OPR_PARAM(GammaRNG);
    public:
        virtual void exec(_megdnn_tensor_in shape, 
                        _megdnn_tensor_in scale, 
                        _megdnn_tensor_out dst,
                        _megdnn_workspace workspace) = 0;
        virtual size_t get_workspace_in_bytes(const TensorLayout &shape, 
                    const TensorLayout &scale, const TensorLayout &dst) = 0;
    protected:
        void check_exec(const TensorLayout &shape, const TensorLayout &scale,
                        const TensorLayout &dst, size_t workspace_in_bytes);
};

//! sample from uniform distribution on the interval (0, 1]
class UniformRNG: public RNGBase {
    DEF_OPR_IMPL(UniformRNG, RNGBase, 0, 1);
    DEF_OPR_PARAM(UniformRNG);
    protected:
        void check_exec(const TensorLayout &dst, size_t workspace_in_bytes);
};

//! sample from gaussian distribution
class GaussianRNG: public RNGBase {
    DEF_OPR_IMPL(GaussianRNG, RNGBase, 0, 1);
    DEF_OPR_PARAM(GaussianRNG);
    protected:
        void check_exec(const TensorLayout &dst, size_t workspace_in_bytes);
};

class PermutationRNG: public RNGBase {
    DEF_OPR_IMPL(PermutationRNG, RNGBase, 0, 1);
    DEF_OPR_PARAM(PermutationRNG);
    protected:
        void check_exec(const TensorLayout &dst, size_t workspace_in_bytes);
};

class ShuffleRNGForward : public OperatorBase {
    DEF_OPR_IMPL(ShuffleRNGForward, OperatorBase, 1, 2);
    DEF_OPR_PARAM(ShuffleRNG);

public:
    virtual void exec(_megdnn_tensor_in src, _megdnn_tensor_out dst,
                      _megdnn_tensor_out indices,
                      _megdnn_workspace workspace) = 0;
    void deduce_layout(const TensorLayout& src, TensorLayout& dst,
                       TensorLayout& indices);
    virtual size_t get_workspace_in_bytes(const TensorLayout& src,
                                          const TensorLayout& dst,
                                          const TensorLayout& indices) = 0;

protected:
    void check_exec(const TensorLayout& src, const TensorLayout& dst,
                    const TensorLayout& indices, size_t workspace_in_bytes);
};
using ShuffleRNG = ShuffleRNGForward;

class ShuffleRNGBackward : public OperatorBase {
    DEF_OPR_IMPL(ShuffleRNGBackward, OperatorBase, 2, 1);
    DEF_OPR_PARAM(ShuffleRNG);

public:
    virtual void exec(_megdnn_tensor_in diff, _megdnn_tensor_in indices,
                      _megdnn_tensor_out grad, _megdnn_workspace workspace) = 0;
    virtual size_t get_workspace_in_bytes(const TensorLayout& diff,
                                          const TensorLayout& indices,
                                          const TensorLayout& grad) = 0;

protected:
    void check_exec(const TensorLayout& diff, const TensorLayout& indices,
                    const TensorLayout& grad, size_t workspace_in_bytes);
};

/*!
 * \brief sleep for specific time on the computing device; useful for testing
 *      async problems
 */
class SleepForward: public OperatorBase {
    DEF_OPR_IMPL(SleepForward, OperatorBase, 0, 0);
    DEF_OPR_PARAM(Sleep);

    public:
        virtual void exec() = 0;
};
using Sleep = SleepForward;

/*!
 * \brief calculating checksum of a tensor
 *
 * data must be a one-dimensional contiguous tensor with dtype byte
 */
class ChecksumForward: public OperatorBase {
    DEF_OPR_PARAM(Empty);
    DEF_OPR_IMPL(ChecksumForward, OperatorBase, 0, 1);

    public:
        using Result = opr_result::Checksum;

        virtual size_t get_workspace_in_bytes(const TensorLayout &data) = 0;

        virtual Result exec(_megdnn_tensor_in data,
                _megdnn_workspace workspace) = 0;

    protected:
        void check_exec(const TensorLayout &layout, size_t workspace_in_bytes);
};
using Checksum = ChecksumForward;

/*!
 * \brief calculating max absolute difference of the two input tensors
 *
 * src1 and src2 must be a one-dimensional contiguous tensor.
 */
class MaxTensorDiff : public OperatorBase {
    DEF_OPR_PARAM(Empty);
    DEF_OPR_IMPL(MaxTensorDiff, OperatorBase, 0, 2);

    public:
        virtual size_t get_workspace_in_bytes(const TensorLayout& layout1,
                                              const TensorLayout& layout2) = 0;

        virtual float exec(_megdnn_tensor_in src1, _megdnn_tensor_in src2,
                           _megdnn_workspace workspace) = 0;

    protected:
        void check_exec(const TensorLayout& layout1,
                        const TensorLayout& layout2, size_t workspace_in_bytes);
};


bool check_bias_share_in_channel(const TensorLayout& bias,
                                 const param::ConvBias::Format format);

}  // namespace megdnn

#include "megdnn/internal/opr_header_epilogue.h"

// vim: syntax=cpp.doxygen