comp_node.h 28.8 KB
Newer Older
1 2 3 4
#pragma once

#include "megbrain/utils/hash.h"
#include "megbrain/utils/metahelper.h"
M
Megvii Engine Team 已提交
5
#include "megbrain/utils/thin/function.h"
6 7 8 9 10 11
#include "megbrain/utils/thin/hash_table.h"
#include "megbrain/utils/thread.h"
#include "megdnn/thin/function.h"

#include <cstddef>
#include <memory>
M
Megvii Engine Team 已提交
12
#include <string>
13 14 15 16 17 18 19 20 21 22

namespace mgb {

// forward declaration; defined in comp_node_env.h
class CompNodeEnv;

namespace cg {
class ComputingGraph;
}

23
class CompNodeSeqRecorder;
24 25 26 27 28 29 30 31 32 33

/*!
 * \brief identifier for a memory node
 *
 * MemNode is comparable. CompNodes with the same MemNode can access memory of
 * each other directly
 */
class MemNode {
    const void* m_id = nullptr;

M
Megvii Engine Team 已提交
34 35
public:
    MemNode() = default;
36

M
Megvii Engine Team 已提交
37
    explicit MemNode(const void* id) : m_id{id} {}
38

M
Megvii Engine Team 已提交
39
    bool operator==(const MemNode& rhs) const { return m_id == rhs.m_id; }
40

M
Megvii Engine Team 已提交
41
    bool operator!=(const MemNode& rhs) const { return m_id != rhs.m_id; }
42

M
Megvii Engine Team 已提交
43
    operator bool() const { return m_id != nullptr; }
44 45 46 47 48 49 50 51 52 53
};

/*!
 * \brief abstraction of a streaming computing resource on localhost (a
 *      thread on CPU, a cuda stream, etc.)
 *
 * Note that most of the operations are asynchronous with respect to the caller
 * thread
 */
class CompNode {
M
Megvii Engine Team 已提交
54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70
public:
    //! computing device type
    enum class DeviceType {
        //! for "xpu" comp node that would mapped to available cn on
        //! current system
        UNSPEC = 0,

        CUDA = 1,
        CPU = 2,
        CAMBRICON = 3,
        ROCM = 8,
        ATLAS = 9,
        MULTITHREAD = 11,
        MAX_DEVICE_ID,
    };
    static constexpr size_t NR_DEVICE_TYPE =
            static_cast<size_t>(DeviceType::MAX_DEVICE_ID);
71

72 73 74 75 76 77 78 79 80 81 82 83 84 85
    struct DeviceProperties {
        DeviceProperties() {
            name = "unspec";
            total_memory = major = minor = 0;
        }

        std::string name;
        size_t total_memory;

        //! for cuda
        int major;
        int minor;
    };

M
Megvii Engine Team 已提交
86 87 88 89 90 91 92 93 94 95 96
    /*!
     * \brief an identifier to specify a computing node
     *
     * Note: logical locator is directly parsed from a string identifier
     * given by user; it should be translated to physical locator by calling
     * to_physical() before actual use.
     *
     * Unless explicitly specified otherwise, all locators are physical
     * locators.
     */
    struct Locator {
97
        /*!
M
Megvii Engine Team 已提交
98 99
         * \brief special device number for the "cpu default" comp node,
         *      which dispatches all tasks in the caller thread
100
         */
M
Megvii Engine Team 已提交
101 102 103 104 105 106 107
        static constexpr int DEVICE_CPU_DEFAULT = -1024;
        /*!
         * \brief special device number for the "multithread_default"
         * comp node, which dispatches all tasks to thread pool and the
         * caller thread is the main thread of thread pool
         */
        static constexpr int DEVICE_MULTITHREAD_DEFAULT = -1025;
108

M
Megvii Engine Team 已提交
109
        DeviceType type = DeviceType::UNSPEC;
110

M
Megvii Engine Team 已提交
111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126
        /*!
         * corresponding to a physical computing device; memories between
         * different devices are not shared.
         *
         * device == -1 means logical default device (maps to 0 by default,
         * and can be changed by set_device_map)
         *
         */
        int device = -1;

        //! multiple streams can execute on one computing device and share
        //! memory, when compnode type is multithread the field also stand
        //! for nr_threads
        union {
            int stream = 0;
            int nr_threads;
127 128 129
        };

        /*!
M
Megvii Engine Team 已提交
130 131 132 133
         * \brief parse a string identifier
         *
         * currently supported ID format: (gpu|cpu)<n>[:m] where n is the
         * device number, possibly with m as the stream id.
134
         */
135
        MGE_WIN_DECLSPEC_FUC static Locator parse(const std::string& id);
136 137

        /*!
M
Megvii Engine Team 已提交
138
         * \brief set mapping between device numbers of a device type
139
         */
140 141
        MGE_WIN_DECLSPEC_FUC static void set_device_map(
                DeviceType type, int from, int to);
142 143

        /*!
M
Megvii Engine Team 已提交
144 145
         * \brief set the actual device type to be used for
         *      DeviceType::UNSPEC
146
         */
147
        MGE_WIN_DECLSPEC_FUC static void set_unspec_device_type(DeviceType type);
148 149

        /*!
M
Megvii Engine Team 已提交
150
         * \brief get corresponding physical Locator
151
         *
M
Megvii Engine Team 已提交
152 153
         * DeviceType::UNSPEC would be resolved, and device map would be
         * applied on device number
154
         */
155
        MGE_WIN_DECLSPEC_FUC Locator to_physical() const;
156 157

        /*!
M
Megvii Engine Team 已提交
158 159
         * \brief get string description of this locator that can be parsed
         *      again
160
         */
161
        MGE_WIN_DECLSPEC_FUC std::string to_string() const;
162

M
Megvii Engine Team 已提交
163 164
        bool operator==(const Locator& rhs) const {
            return type == rhs.type && device == rhs.device && stream == rhs.stream;
165
        }
M
Megvii Engine Team 已提交
166
    };
167

M
Megvii Engine Team 已提交
168 169
    struct LocatorPairHashKey {
        Locator locator, locator_logical;
170

M
Megvii Engine Team 已提交
171 172
        bool operator==(const LocatorPairHashKey& rhs) const {
            return locator == rhs.locator && locator_logical == rhs.locator_logical;
173 174
        }

M
Megvii Engine Team 已提交
175 176 177 178 179 180 181
        struct Hash {
            size_t operator()(const LocatorPairHashKey& k) const {
                return hash_pair_combine(
                        mgb::hash(k.locator), mgb::hash(k.locator_logical));
            }
        };
    };
182

M
Megvii Engine Team 已提交
183 184 185 186
    //! predefined special streams
    struct Stream {
        static constexpr int COPY = -1, REMOTE_SEND = -2, LOOP_SWAP = -3;
    };
187

M
Megvii Engine Team 已提交
188
    CompNode() = default;
189

M
Megvii Engine Team 已提交
190 191 192
    /*!
     * \brief manually destroy all comp node resources
     */
193
    MGE_WIN_DECLSPEC_FUC static void finalize();
194

M
Megvii Engine Team 已提交
195 196 197 198 199
    /*!
     * \brief load a computing node from logical locator ID;
     * \see Locator::parse
     */
    static CompNode load(const std::string& id) { return load(Locator::parse(id)); }
200

M
Megvii Engine Team 已提交
201 202 203 204 205 206
    /*!
     * \brief create a CompNode object from **logical** locator
     */
    static CompNode load(const Locator& locator) {
        return load(locator.to_physical(), locator);
    }
207

208
    MGE_WIN_DECLSPEC_FUC static CompNode load(
M
Megvii Engine Team 已提交
209
            const Locator& locator_physical, const Locator& locator_logical);
210

M
Megvii Engine Team 已提交
211
    /* =================== memory management ======================== */
212

M
Megvii Engine Team 已提交
213 214 215 216 217 218 219 220 221 222
    /*!
     * \brief allocate memory on this computing node
     *
     * Note: allocation of device memory is synchronous with the host,
     * meaning that the memory can be used immediately; however deallocation
     * is asynchronous to ensure that the memory can be used by
     * already-launched kernels on the computing node.
     *
     * Exception should be raised if allocation fails.
     */
223
    MGE_WIN_DECLSPEC_FUC void* alloc_device(size_t size) const;
224

M
Megvii Engine Team 已提交
225
    //! deallocate device buffer; see alloc_device() for more details
226
    MGE_WIN_DECLSPEC_FUC void free_device(void* ptr) const;
227

M
Megvii Engine Team 已提交
228 229 230 231 232 233
    /*!
     * \brief allocate memory on host that is associated with the device,
     *      which may accelerate I/O
     *
     * Both allocation and deallocation on host are synchronous.
     */
234
    MGE_WIN_DECLSPEC_FUC void* alloc_host(size_t size) const;
235

236
    MGE_WIN_DECLSPEC_FUC void free_host(void* ptr) const;
237

M
Megvii Engine Team 已提交
238 239 240 241
    //! copy from underlying device to host
    void copy_to_host(void* host_ptr, const void* device_ptr, size_t size) const {
        return m_impl->copy_to_host(host_ptr, device_ptr, size);
    }
242

M
Megvii Engine Team 已提交
243 244 245 246
    //! copy from host to underlying device
    void copy_to_device(void* device_ptr, const void* host_ptr, size_t size) const {
        return m_impl->copy_to_device(device_ptr, host_ptr, size);
    }
247

248 249 250 251 252 253 254 255 256 257 258 259 260 261
    //! copy from underlying device to host
    void copy_to_host_ref(
            megdnn::RefPtr& host_ref_ptr, megdnn::RefPtr& device_ref_ptr,
            size_t size) const {
        return m_impl->copy_to_host_ref(host_ref_ptr, device_ref_ptr, size);
    }

    //! copy from host to underlying device
    void copy_to_device_ref(
            megdnn::RefPtr& device_ref_ptr, megdnn::RefPtr& host_ref_ptr,
            size_t size) const {
        return m_impl->copy_to_device_ref(device_ref_ptr, host_ref_ptr, size);
    }

M
Megvii Engine Team 已提交
262 263 264 265 266 267 268 269 270 271 272
    /*!
     * \brief copy from this device to another device; would use the
     *      computing resource on dest_node
     * \param src source memory that must be allocated on this device
     */
    void peer_copy_to(
            CompNode dest_node, void* dest, const void* src, size_t size) const {
        return m_impl->peer_copy_to(
                reinterpret_cast<Impl*>(dest_node.m_impl), dest, src, size);
    }

273 274 275 276 277 278 279 280
    void peer_copy_to_ref(
            CompNode dest_node, megdnn::RefPtr& dst_ref_ptr,
            megdnn::RefPtr& src_ref_ptr, size_t size) const {
        return m_impl->peer_copy_to_ref(
                reinterpret_cast<Impl*>(dest_node.m_impl), dst_ref_ptr, src_ref_ptr,
                size);
    }

M
Megvii Engine Team 已提交
281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297
    //! get alignment requiement in bytes; guaranteed to be power of 2
    size_t get_mem_addr_alignment() const { return m_impl->get_mem_addr_alignment(); }

    /*!
     * \brief get the size of the paddings which must be reserved at the
     * end of memory chunk; guaranteed to be power of 2
     */
    size_t get_mem_padding() const {
        size_t padding = m_impl->get_mem_padding();
        mgb_assert(!(padding & (padding - 1)), "mem padding should be power of 2");
        return padding;
    }

    /*!
     * \brief release consecutive free chunks on all devices to defragment;
     *      see DevMemAlloc::try_coalesce_free
     */
298
    MGE_WIN_DECLSPEC_FUC static void try_coalesce_all_free_memory();
M
Megvii Engine Team 已提交
299 300 301 302 303

    /*
     * \brief specifies how to pre-allocate from raw dev allocator
     *
     */
304
    MGE_WIN_DECLSPEC_FUC static void set_prealloc_config(
M
Megvii Engine Team 已提交
305 306
            size_t alignment, size_t min_req, size_t max_overhead, double growth_factor,
            DeviceType device_type);
307

M
Megvii Engine Team 已提交
308
    /*!
309
     * \brief get device property of the specified device
M
Megvii Engine Team 已提交
310
     */
311
    MGE_WIN_DECLSPEC_FUC static DeviceProperties get_device_prop(
312
            int dev, DeviceType device_type);
M
Megvii Engine Team 已提交
313

314 315 316 317 318 319 320 321 322 323 324 325 326 327 328
    /*!
     * \brief get control of host ptr to user
     */
    MGE_WIN_DECLSPEC_FUC void map_to_cpu(void* ptr, size_t size, bool blocking = false);

    /*!
     * \brief release control of host ptr to system
     */
    MGE_WIN_DECLSPEC_FUC void unmap_to_gpu(void* ptr, size_t size);

    /*!
     * \brief get logical address by host ptr
     */
    MGE_WIN_DECLSPEC_FUC void* get_logical_addr_by_host_ptr(void* ptr, size_t size);

329 330 331 332 333 334 335 336 337 338 339 340 341 342
    /*!
     * \brief register user external device ptr, which means not malloc by MegEngine
     * case 1: cpu and cuda compnode will do nothing, just return args ptr
     * case 2: OpenCL(ION) compnode will do real register, OpenCL(map/svm) compnode will
     * trigger assert, caused by OpenCL only can use extern ION ptr, can not use map/svm
     * with different OpenCL context.
     */
    MGE_WIN_DECLSPEC_FUC void* register_external_device_ptr(void* ptr, size_t size);

    /*!
     * \brief unregister user external device ptr, which means not malloc by MegEngine
     */
    MGE_WIN_DECLSPEC_FUC void* unregister_external_device_ptr(void* ptr, size_t size);

M
Megvii Engine Team 已提交
343 344 345 346 347 348 349 350 351 352 353 354 355 356 357 358 359 360 361 362 363
    /* =================== synchronization ======================== */

    class Event;
    class EventPool;

    std::unique_ptr<Event> create_event(size_t flags = 0) const {
        return m_impl->create_event(flags);
    }

    //! wait for an event created on another CompNode
    inline void device_wait_event(Event& event) const;

    /*!
     * \brief block host thread to wait for all previous operations on this
     *      computing node to finish
     */
    void sync() const { return m_impl->sync(); }

    /*!
     * \brief synchronize all computing nodes
     */
364
    MGE_WIN_DECLSPEC_FUC static void sync_all();
M
Megvii Engine Team 已提交
365 366 367 368 369 370 371 372 373 374 375 376 377 378 379 380 381 382 383

    /* =================== misc ======================== */

    /*!
     * \brief get id of underlying memory node; comp nodes that share the
     *      same mem node can access memory allocated by each other.
     */
    MemNode mem_node() const { return m_impl->mem_node(); }

    bool operator==(const CompNode& rhs) const { return m_impl == rhs.m_impl; }

    bool operator!=(const CompNode& rhs) const { return !this->operator==(rhs); }

    bool valid() const { return m_impl; }

    //! get total and free memory on the computing device in bytes
    std::pair<size_t, size_t> get_mem_status_bytes() const {
        return m_impl->get_mem_status_bytes();
    }
384

385
#if !MGB_BUILD_SLIM_SERVING
M
Megvii Engine Team 已提交
386 387 388 389
    std::pair<size_t, size_t> get_free_left_and_right(
            size_t begin_ptr, size_t end_ptr) {
        return m_impl->get_free_left_and_right(begin_ptr, end_ptr);
    }
390

M
Megvii Engine Team 已提交
391
    size_t get_used_memory() const { return m_impl->get_used_memory(); }
392

393 394 395 396 397 398
    size_t get_reserved_memory() const { return m_impl->get_reserved_memory(); }

    size_t get_max_reserved_memory() const { return m_impl->get_max_reserved_memory(); }

    size_t get_max_used_memory() const { return m_impl->get_max_used_memory(); }

M
Megvii Engine Team 已提交
399 400 401
    size_t get_max_block_size_available() const {
        return m_impl->get_max_block_size_available();
    }
402 403

    size_t get_free_mem() const { return m_impl->get_free_mem(); }
404 405 406 407 408 409

    void reset_max_reserved_memory() const {
        return m_impl->reset_max_reserved_memory();
    }

    void reset_max_used_memory() const { return m_impl->reset_max_used_memory(); }
410 411
#endif

M
Megvii Engine Team 已提交
412
    //! change to another stream on the same memory node
413
    MGE_WIN_DECLSPEC_FUC CompNode change_stream(int dest_stream) const;
414

M
Megvii Engine Team 已提交
415 416 417 418 419 420 421 422
    //! get string representation
    std::string to_string() const {
        return m_impl ? mgb::ssprintf(
                                "CompNode(\"%s\" from \"%s\")",
                                to_string_physical().c_str(),
                                to_string_logical().c_str())
                      : "invalid";
    }
423

M
Megvii Engine Team 已提交
424 425 426 427
    //! get string representation of physical device
    std::string to_string_physical() const {
        return m_impl ? m_impl->locator().to_string() : "invalid";
    }
428

M
Megvii Engine Team 已提交
429 430 431 432
    //! get string representation of logical device
    std::string to_string_logical() const {
        return m_impl ? m_impl->locator_logical().to_string() : "invalid";
    }
433

M
Megvii Engine Team 已提交
434
    uint64_t get_uid() { return m_impl->get_uid(); }
435

M
Megvii Engine Team 已提交
436 437
    //! get the physical locator that created this comp node
    Locator locator() const { return m_impl->locator(); }
438

M
Megvii Engine Team 已提交
439 440
    //! get the logical locator that created this comp node
    Locator locator_logical() const { return m_impl->locator_logical(); }
441

M
Megvii Engine Team 已提交
442
    //! see CompNodeEnv::activate
443
    MGE_WIN_DECLSPEC_FUC void activate() const;
444

M
Megvii Engine Team 已提交
445
    //! get device type of this comp node
446
    MGE_WIN_DECLSPEC_FUC DeviceType device_type() const;
447

M
Megvii Engine Team 已提交
448 449 450 451 452 453 454 455 456
    /*!
     * \brief check for error on the asynchronous computing stream
     *
     * This is used for devices with limited error handling such as CUDA.
     *
     * It will return MegBrainError with error messages rather than
     * directly throw exception; return nullptr if no error.
     */
    MGB_WARN_UNUSED_RESULT
457
    MGE_WIN_DECLSPEC_FUC std::unique_ptr<MegBrainError> check_async_error() const;
458

M
Megvii Engine Team 已提交
459 460 461 462 463 464 465 466 467 468 469 470 471 472
    /*!
     * \brief create a CompNodeSeqRecorder associated with this computing
     * node
     *
     * Note: the implementation must be thread safe: simultaneous calls to
     * create_seq_recorder() must block until existing CompNodeSeqRecorder
     * objects are either destructed or stopped.
     *
     * \return the recorder object; nullptr is returned if recording is not
     *      supported
     */
    std::unique_ptr<CompNodeSeqRecorder> create_seq_recorder(cg::ComputingGraph* cg) {
        return m_impl->create_seq_recorder(cg);
    }
473

M
Megvii Engine Team 已提交
474 475 476 477 478 479 480 481 482
    /*!
     *  insert callback into current compute stream.
     *  The callack is to be called after all currently enqueued
     *  iterms in the stream have completed. And the later tasks
     *  in the stream must wait for the callback to finish.
     */
    void add_callback(megdnn::thin_function<void()>&& cb) {
        return m_impl->add_callback(std::move(cb));
    }
483

M
Megvii Engine Team 已提交
484 485 486 487
    enum class Flag : uint32_t {
        //! Whether computing recorder is supported on this comp node (i.e.
        //! whether non-zero comp_node_seq_record_level is allowed)
        SUPPORT_RECORDER = 1 << 0,
488

M
Megvii Engine Team 已提交
489 490 491 492 493
        //! Whether dynamic memory allocation is supported in seq recorder.
        //! If this flag is not setted, ComputingSequence::do_execute()
        //! would skip the warm up and allow seq recorder to start
        //! immediately
        RECORDER_SUPPORT_DYNAMIC_ALLOC = 1 << 1,
494

M
Megvii Engine Team 已提交
495 496 497 498 499 500
        //! Whether the capacity of the asynchronous execution queue on this
        //! comp node is limited.
        //! If this flag is set, tasks on multiple comp nodes would be
        //! dispatched from multiple cpu threads.
        //! \see ComputingGraph::Options::async_exec_level
        QUEUE_LIMITED = 1 << 2,
501

M
Megvii Engine Team 已提交
502 503 504
        //! Whether this comp node supports copy stream, so computation and
        //! I/O can be parallelized
        HAS_COPY_STREAM = 1 << 3,
505

M
Megvii Engine Team 已提交
506 507 508 509
        //! Destructing an event is unsafe if the comp node is not
        //! synchronized; setting this flag would cause computing sequence
        //! to sync the comp node in its dtor.
        EVENT_DTOR_UNSAFE = 1 << 4,
510

M
Megvii Engine Team 已提交
511 512 513 514
        //! CompNode is available even there is no thread support, i.e.
        //! MGB_HAVE_THREAD=0. Usually this means that execution on the
        //! CompNode is synchronous, i.e. behaves like cpu:default
        SUPPORT_NO_THREAD = 1 << 5,
515

M
Megvii Engine Team 已提交
516 517 518 519
        //! Whether this comp node supports unified address. i.e. CPU and
        //! CUDA supports unified address.
        SUPPORT_UNIFIED_ADDRESS = 1 << 6,
    };
520

M
Megvii Engine Team 已提交
521
    bool contain_flag(Flag flag) { return contain_flag(device_type(), flag); }
522

523
    MGE_WIN_DECLSPEC_FUC static bool contain_flag(DeviceType device_type, Flag flag);
524

M
Megvii Engine Team 已提交
525
    using UnorderedSet = ThinHashSet<CompNode>;
526

M
Megvii Engine Team 已提交
527 528 529 530
    template <typename T>
    using UnorderedMap = ThinHashMap<CompNode, T>;

    //! apply function to each initialized comp node
531
    MGE_WIN_DECLSPEC_FUC static void foreach (thin_function<void(CompNode)> callback);
M
Megvii Engine Team 已提交
532 533

    //! get total number of specific devices on this system
534 535
    MGE_WIN_DECLSPEC_FUC static size_t get_device_count(
            DeviceType type, bool warn = true);
M
Megvii Engine Team 已提交
536 537

    /* =================== specialized ======================== */
538

M
Megvii Engine Team 已提交
539 540
    //! get default CPU comp node
    // implemented in comp_node/cpu/comp_node.cpp
541
    MGE_WIN_DECLSPEC_FUC static CompNode default_cpu();
542

M
Megvii Engine Team 已提交
543 544 545 546 547 548 549 550 551 552 553
    /*!
     * \brief set whether to enable affinity setting for CPU comp nodes
     *
     * If enabled, computation on cpux would be bound to the x'th CPU.
     *
     * This is disabled by default.
     *
     * (implemented in comp_node/cpu/comp_node.cpp)
     *
     * \return original setting
     */
554
    MGE_WIN_DECLSPEC_FUC static bool enable_affinity_for_cpu(bool flag);
555

M
Megvii Engine Team 已提交
556 557 558
protected:
    //! ImplBase with env(); defined in CompNodeEnv
    class Impl;
559

M
Megvii Engine Team 已提交
560 561 562 563 564 565 566 567 568 569 570 571 572 573 574
    class ImplBase : public NonCopyableObj, public DynTypeObj {
    public:
        typedef void (*free_func_t)(ImplBase* self, void* ptr);
        //! memory free might be called after finalize(); so we should
        //! not rely on virtual function for this
        const free_func_t free_device;
        const free_func_t free_host;

        virtual void* alloc_device(size_t size) = 0;
        virtual void* alloc_host(size_t size) = 0;

        virtual void copy_to_host(
                void* host_ptr, const void* device_ptr, size_t size) = 0;
        virtual void copy_to_device(
                void* device_ptr, const void* host_ptr, size_t size) = 0;
575 576 577 578 579 580 581 582 583 584
        virtual void copy_to_host_ref(
                megdnn::RefPtr& host_ref_ptr, megdnn::RefPtr& device_ref_ptr,
                size_t size) {
            copy_to_host(host_ref_ptr.get_ptr(), device_ref_ptr.get_ptr(), size);
        }
        virtual void copy_to_device_ref(
                megdnn::RefPtr& device_ref_ptr, megdnn::RefPtr& host_ref_ptr,
                size_t size) {
            copy_to_device(device_ref_ptr.get_ptr(), host_ref_ptr.get_ptr(), size);
        }
M
Megvii Engine Team 已提交
585 586
        virtual void peer_copy_to(
                Impl* dest_impl, void* dest, const void* src, size_t size) = 0;
587

588 589 590 591 592 593
        virtual void peer_copy_to_ref(
                Impl* dest_impl, megdnn::RefPtr& dest, megdnn::RefPtr& src,
                size_t size) {
            peer_copy_to(dest_impl, dest.get_ptr(), src.get_ptr(), size);
        }

594 595 596 597 598 599
        virtual void map_to_cpu(void* ptr, size_t size, bool blocking = false);

        virtual void unmap_to_gpu(void* ptr, size_t size);

        virtual void* get_logical_addr_by_host_ptr(void* ptr, size_t size);

600 601 602 603
        virtual void* register_external_device_ptr(void* ptr, size_t size);

        virtual void* unregister_external_device_ptr(void* ptr, size_t size);

M
Megvii Engine Team 已提交
604 605
        virtual size_t get_mem_addr_alignment() = 0;
        virtual size_t get_mem_padding();
606

M
Megvii Engine Team 已提交
607
        virtual std::unique_ptr<Event> create_event(size_t flags) = 0;
608

M
Megvii Engine Team 已提交
609 610 611 612
        virtual void sync() = 0;

        virtual MemNode mem_node() = 0;
        virtual std::pair<size_t, size_t> get_mem_status_bytes() = 0;
613

614
#if !MGB_BUILD_SLIM_SERVING
M
Megvii Engine Team 已提交
615 616 617 618
        virtual std::pair<size_t, size_t> get_free_left_and_right(size_t x, size_t y) {
            return {x - x, y - y};
        }
        virtual size_t get_used_memory() { return 0; }
619 620 621
        virtual size_t get_reserved_memory() { return 0; }
        virtual size_t get_max_reserved_memory() { return 0; }
        virtual size_t get_max_used_memory() { return 0; }
M
Megvii Engine Team 已提交
622
        virtual size_t get_max_block_size_available() { return 0; }
623
        virtual size_t get_free_mem() { return get_mem_status_bytes().second; }
624 625
        virtual void reset_max_reserved_memory() {}
        virtual void reset_max_used_memory() {}
626 627
#endif

M
Megvii Engine Team 已提交
628 629 630 631 632
        virtual Locator locator() = 0;
        virtual Locator locator_logical() = 0;

        virtual std::unique_ptr<CompNodeSeqRecorder> create_seq_recorder(
                cg::ComputingGraph* cg);
633

M
Megvii Engine Team 已提交
634
        virtual void add_callback(megdnn::thin_function<void()>&&);
635

M
Megvii Engine Team 已提交
636 637 638
        virtual uint64_t get_uid() {
            mgb_throw(MegBrainError, "get_uid is not impl yet");
        };
639

M
Megvii Engine Team 已提交
640 641
    protected:
        ImplBase(free_func_t fd, free_func_t fh) : free_device{fd}, free_host{fh} {}
642

M
Megvii Engine Team 已提交
643 644
        ~ImplBase() = default;
    };
645

M
Megvii Engine Team 已提交
646 647 648
    //! implementations are allocated statically, so no memory management
    //! is needed
    ImplBase* m_impl = nullptr;
649

M
Megvii Engine Team 已提交
650 651 652 653
    friend class CompNodeEnv;
    friend struct HashTrait<CompNode>;
    friend struct HashTrait<CompNode::Locator>;
    friend class CompNodeImplHelper;
654

M
Megvii Engine Team 已提交
655 656
public:
    CompNode(ImplBase* impl) : m_impl{impl} {}
657 658 659 660
};

MGB_DEF_ENUM_CLASS_BIT_OPR(CompNode::Flag)

661 662 663 664 665 666 667 668 669 670 671 672 673 674 675 676 677 678 679 680 681 682 683 684 685 686 687 688 689 690 691 692 693 694 695 696
/*!
 * \brief record computation operations on a computing node
 *
 * This is used for fast execution of an identical computation sequence where
 * only input/output data differ.
 *
 * When this object is created from a comp node, recording starts immediately.
 * Call stop() when computation finishes, and call replay() when it needs to be
 * re-executed.
 *
 * Implementations should consider thread safe in comp_node, in order to support
 * multi threads reording in the same comp_node simultaneously, using thread
 * local recorder in comp_node.
 *
 * Note. When recording is over, the recorder is independent with comp_node, so
 * the task dispatched into recorder should not related to the comp_node
 * methord, and the thread of recorder replay is the user thread.
 */
class CompNodeSeqRecorder {
public:
    virtual ~CompNodeSeqRecorder() noexcept = default;

    /*!
     * \brief Enter fake-exec mode
     *
     * Memory allocation/free is only allowed in fake-exec mode, and kernels
     * should not be actually recorded in this mode.
     *
     * This should be paired with exit_fake_exec()
     */
    virtual void enter_fake_exec(const CompNode& comp_node) = 0;

    //! Exit fake-exec mode
    virtual void exit_fake_exec(const CompNode& comp_node) = 0;

    virtual void stop(const CompNode& comp_node) = 0;
697

698 699 700
    virtual void replay() = 0;
};

701 702 703 704
/*!
 * \brief event associated with a CompNode node, used for cross-device
 *      synchronization
 */
M
Megvii Engine Team 已提交
705 706 707
class CompNode::Event : public NonCopyableObj {
protected:
    static int sm_cpu_sync_level;
708

M
Megvii Engine Team 已提交
709 710
    //! flags when this event is created
    size_t const m_create_flags;
711

M
Megvii Engine Team 已提交
712
    Event(size_t create_flags) : m_create_flags{create_flags} {}
713

M
Megvii Engine Team 已提交
714 715
public:
    enum Flags { NEED_TIMER = 1 };
716

M
Megvii Engine Team 已提交
717
    virtual ~Event() = default;
718

M
Megvii Engine Team 已提交
719 720 721 722 723 724 725 726 727
    /*!
     * \brief record this event on the comp node that creates it
     *
     * Note that if a comp node is recorded multiple times, then subsequent
     * calls would overwrite its internal state and other methods that
     * examine the status would only examine the completion of the most
     * recent call to record().
     */
    virtual void record() = 0;
728

M
Megvii Engine Team 已提交
729 730
    //! whether this event has finished; it must has been recorded
    virtual bool finished() = 0;
731

M
Megvii Engine Team 已提交
732 733
    //! block the host thread (caller thread) to wait for this event
    virtual void host_wait() = 0;
734

M
Megvii Engine Team 已提交
735 736 737
    //! get elapsed time in seconds from this to another event; the events
    //! must be finished
    virtual double elapsed_time_until(Event& end) = 0;
738

M
Megvii Engine Team 已提交
739 740 741
    //! record an action on another comp node so it would wait for this
    //! event
    virtual void device_wait_by(CompNode cn) = 0;
742

M
Megvii Engine Team 已提交
743 744
    //! get the comp node to which this event is associated
    virtual CompNode comp_node() const = 0;
745

M
Megvii Engine Team 已提交
746 747
    //! flags when this event is created
    size_t create_flags() const { return m_create_flags; }
748

M
Megvii Engine Team 已提交
749 750 751 752 753 754 755 756
    /*!
     * \brief set CPU resource usage level when performing synchronization
     * \param level CPU waiting level:
     *      0. condition var (the default)
     *      1. busy wait with yield
     *      2. busy wait
     */
    static void set_cpu_sync_level(int level) { sm_cpu_sync_level = level; }
757 758 759 760 761 762 763 764 765 766
};

/*!
 * \brief pool of events that can be reused
 */
class CompNode::EventPool {
    CompNode m_cn;
    std::vector<std::unique_ptr<CompNode::Event>> m_allocated;
    std::vector<CompNode::Event*> m_free;
    Spinlock m_lock;
767
    size_t m_flags;
768

M
Megvii Engine Team 已提交
769
public:
770 771
    MGE_WIN_DECLSPEC_FUC explicit EventPool(CompNode cn, size_t flags = 0);
    MGE_WIN_DECLSPEC_FUC ~EventPool();
772

773
    MGE_WIN_DECLSPEC_FUC CompNode::Event* alloc();
774

775
    MGE_WIN_DECLSPEC_FUC void free(CompNode::Event* ev);
776

M
Megvii Engine Team 已提交
777
    //! assert that all allocated events have been freed
778
    MGE_WIN_DECLSPEC_FUC void assert_all_freed();
779 780
};

M
Megvii Engine Team 已提交
781
void CompNode::device_wait_event(Event& event) const {
782 783 784
    event.device_wait_by(*this);
}

M
Megvii Engine Team 已提交
785
template <>
786
struct HashTrait<CompNode> {
M
Megvii Engine Team 已提交
787
    static size_t eval(const CompNode& val) {
788 789 790 791 792
        static_assert(sizeof(size_t) == sizeof(void*), "bad hash type");
        return reinterpret_cast<size_t>(static_cast<void*>(val.m_impl));
    }
};

M
Megvii Engine Team 已提交
793
template <>
794
struct HashTrait<CompNode::Locator> {
M
Megvii Engine Team 已提交
795 796 797
    static size_t eval(const CompNode::Locator& val) {
        return static_cast<size_t>(val.device) + (static_cast<size_t>(val.type) << 4) +
               (static_cast<size_t>(val.stream) << 8);
798 799 800
    }
};

801 802 803 804 805 806 807 808 809 810 811 812 813 814 815 816 817 818 819 820 821
namespace comp_node_detail {

/*!
 * \brief an inplace doubly linked list for efficient inserting/deleting
 *
 * Note: do not use this directly; it is only for CompNodeDepedentObject
 */
class DepedentObjList {
    class Sentinel;

    struct StaticInfo;
    static StaticInfo sm_info;

    DepedentObjList *m_prev = nullptr, *m_next = nullptr;

    static void link(DepedentObjList* a, DepedentObjList* b) {
        a->m_next = b;
        b->m_prev = a;
    }

protected:
822
    MGE_WIN_DECLSPEC_FUC virtual std::shared_ptr<void> callback() = 0;
823 824
    ~DepedentObjList() = default;

825 826
    MGE_WIN_DECLSPEC_FUC static void add(DepedentObjList* ptr);
    MGE_WIN_DECLSPEC_FUC static void remove(DepedentObjList* ptr);
827 828

public:
829
    MGE_WIN_DECLSPEC_FUC static void invoke_callback_and_clean();
830 831 832 833 834 835 836 837 838 839 840 841 842 843 844 845 846 847 848 849 850 851 852 853
};

}  // namespace comp_node_detail

/*!
 * \brief base class for objects that depend on CompNode
 *
 * There is a CompNode::finalize() method that destorys all global comp nodes.
 * Therefore objects that depend on CompNode should all be marked as invalid at
 * that time.
 *
 * CompNode::finalize() is called in atexit() because some external libraries
 * that CompNode depends on seems to be registering exit handlers. It is also
 * impractical to require a correct destruction order because, for example, in
 * python atexit() handlers are invoked before global python objects get
 * reclaimed.
 *
 * As a result we give up enforcing a correct destruction order, but rather
 * require all CompNode-dependent objects to derive from this class so they can
 * get notified possibly do most of the cleanup when CompNode is finalized.
 */
class CompNodeDepedentObject : private comp_node_detail::DepedentObjList {
    //! 1: in on_comp_node_finalize(); 2: after on_comp_node_finalize()
    int m_state = 0;
854
    MGE_WIN_DECLSPEC_FUC std::shared_ptr<void> callback() override final;
855 856 857 858 859 860 861 862 863 864 865 866 867 868 869 870 871 872 873 874 875 876 877

protected:
    CompNodeDepedentObject() { add(this); }
    ~CompNodeDepedentObject() { remove(this); }

    /*!
     * \brief overwritten by subclasses to perform clean up jobs
     *
     * Note: in case the object has nested objects which hold a reference to the
     * object itself, a reference to this object must be kept so it would not be
     * released during the call of on_comp_node_finalize().
     */
    virtual std::shared_ptr<void> on_comp_node_finalize() = 0;

    //! exception would thrown if on_comp_node_finalize() has been called (do
    //! not raise if invoked from on_comp_node_finalize())
    void check_not_finalized() const;

    //! whether on_comp_node_finalize() has been called (true when invoked
    //! from on_comp_node_finalize())
    bool is_finalized() const { return m_state; }
};

M
Megvii Engine Team 已提交
878
}  // namespace mgb
879 880

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