opr_impl.h 5.0 KB
Newer Older
1 2 3 4
/**
 * \file dnn/src/cuda/rng/opr_impl.h
 * MegEngine is Licensed under the Apache License, Version 2.0 (the "License")
 *
5
 * Copyright (c) 2014-2021 Megvii Inc. All rights reserved.
6 7 8
 *
 * Unless required by applicable law or agreed to in writing,
 * software distributed under the License is distributed on an
9 10
 * "AS IS" BASIS, WITHOUT ARRANTIES OR CONDITIONS OF ANY KIND, either express or
 * implied.
11 12 13
 */
#pragma once

14
#include <curand.h>
15 16 17 18 19 20 21 22 23 24 25
#include "megdnn/oprs.h"
#include "src/cuda/handle.h"

namespace megdnn {
namespace cuda {

class CuRandHandle {
    curandGenerator_t m_gen;
    uint64_t m_seed;

    CuRandHandle(const CuRandHandle&) = delete;
26
    CuRandHandle& operator=(const CuRandHandle&) = delete;
27

28 29 30
public:
    CuRandHandle(cudaStream_t stream, uint64_t seed = 0);
    ~CuRandHandle();
31

32
    void seed(uint64_t seed);
33

34
    curandGenerator_t gen() const { return m_gen; }
35

36 37 38
    void ensure_seed(uint64_t seed) {
        if (m_seed != seed) {
            this->seed(seed);
39
        }
40
    }
41 42
};

43
class UniformRNGImpl : public UniformRNG {
44 45
    CuRandHandle m_curand_handle;

46 47 48
public:
    UniformRNGImpl(Handle* handle);
    void exec(_megdnn_tensor_inout dst, _megdnn_workspace) override;
49

50
    size_t get_workspace_in_bytes(const TensorLayout&) override { return 0; }
51 52
};

53
class GaussianRNGImpl : public GaussianRNG {
54 55
    CuRandHandle m_curand_handle;

56 57 58 59 60 61 62 63 64 65 66 67 68 69
public:
    GaussianRNGImpl(Handle* handle);

    void exec(_megdnn_tensor_inout dst, _megdnn_workspace) override;

    size_t get_workspace_in_bytes(const TensorLayout& layout) override;
};

class GammaRNGImpl : public GammaRNG {
    uint64_t m_seed, m_offset;
    cudaStream_t m_stream;

public:
    GammaRNGImpl(Handle* handle);
70

71 72
    void exec(_megdnn_tensor_in shape,_megdnn_tensor_in scale, 
              _megdnn_tensor_out dst, _megdnn_workspace) override;
73

74 75 76 77
    size_t get_workspace_in_bytes(const TensorLayout&, const TensorLayout&,
                                  const TensorLayout&) override {
        return 0;
    }
78

79 80 81 82 83 84 85
    void seed(uint64_t seed) { m_seed = seed; }

    void ensure_seed(uint64_t seed) {
        if (m_seed != seed) {
            this->seed(seed);
        }
    }
86 87
};

88 89 90
class BetaRNGImpl : public BetaRNG {
    uint64_t m_seed, m_offset;
    cudaStream_t m_stream;
91

92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110
public:
    BetaRNGImpl(Handle* handle);

    void exec(_megdnn_tensor_in alpha,_megdnn_tensor_in beta, 
              _megdnn_tensor_out dst, _megdnn_workspace) override;

    size_t get_workspace_in_bytes(const TensorLayout&, const TensorLayout&,
                                  const TensorLayout&) override {
        return 0;
    }

    void seed(uint64_t seed) { m_seed = seed; }

    void ensure_seed(uint64_t seed) {
        if (m_seed != seed) {
            this->seed(seed);
        }
    }
};
111

112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155
class PoissonRNGImpl : public PoissonRNG {
    uint64_t m_seed, m_offset;
    cudaStream_t m_stream;

public:
    PoissonRNGImpl(Handle* handle);

    void exec(_megdnn_tensor_in lam, _megdnn_tensor_out dst,
              _megdnn_workspace) override;

    size_t get_workspace_in_bytes(const TensorLayout&,
                                  const TensorLayout&) override {
        return 0;
    }

    void seed(uint64_t seed) { m_seed = seed; }

    void ensure_seed(uint64_t seed) {
        if (m_seed != seed) {
            this->seed(seed);
        }
    }
};

class PermutationRNGImpl : public PermutationRNG {
    uint64_t m_seed, m_offset;
    cudaStream_t m_stream;

public:
    PermutationRNGImpl(Handle* handle);

    void exec(_megdnn_tensor_inout dst, _megdnn_workspace) override;

    size_t get_workspace_in_bytes(const TensorLayout& layout) override;

    void seed(uint64_t seed) { m_seed = seed; }

    void ensure_seed(uint64_t seed) {
        if (m_seed != seed) {
            this->seed(seed);
        }
    }
};

156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194
class ShuffleRNGForwardImpl : public ShuffleRNGForward {
    uint64_t m_seed, m_offset;
    cudaStream_t m_stream;

public:
    using ShuffleRNGForward::ShuffleRNGForward;
    ShuffleRNGForwardImpl(Handle* handle);

    void exec(_megdnn_tensor_in src, _megdnn_tensor_out dst,
              _megdnn_tensor_out indices, _megdnn_workspace workspace) override;

    size_t get_workspace_in_bytes(const TensorLayout& src,
                                  const TensorLayout& dst,
                                  const TensorLayout& indices) override;

    void seed(uint64_t seed) { m_seed = seed; }

    void ensure_seed(uint64_t seed) {
        if (m_seed != seed) {
            this->seed(seed);
        }
    }
};

class ShuffleRNGBackwardImpl : public ShuffleRNGBackward {
    cudaStream_t m_stream;

public:
    using ShuffleRNGBackward::ShuffleRNGBackward;
    ShuffleRNGBackwardImpl(Handle* handle);
    void exec(_megdnn_tensor_in diff, _megdnn_tensor_in indices,
              _megdnn_tensor_out grad, _megdnn_workspace workspace) override;
    size_t get_workspace_in_bytes(const TensorLayout&,
                                  const TensorLayout&,
                                  const TensorLayout&) override {
        return 0;
    }
};

195 196 197
}  // namespace cuda
}  // namespace megdnn
// vim: syntax=cpp.doxygen