comp_node.h 25.6 KB
Newer Older
1 2 3 4
/**
 * \file src/core/include/megbrain/comp_node.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 9 10 11 12 13 14 15
 *
 * 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 "megbrain/utils/hash.h"
#include "megbrain/utils/metahelper.h"
M
Megvii Engine Team 已提交
16
#include "megbrain/utils/thin/function.h"
17 18 19 20 21 22
#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 已提交
23
#include <string>
24 25 26 27 28 29 30 31 32 33

namespace mgb {

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

namespace cg {
class ComputingGraph;
}

34
class CompNodeSeqRecorder;
35 36 37 38 39 40 41 42 43 44

/*!
 * \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 已提交
45 46
public:
    MemNode() = default;
47

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

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

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

M
Megvii Engine Team 已提交
54
    operator bool() const { return m_id != nullptr; }
55 56 57 58 59 60 61 62 63 64
};

/*!
 * \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 已提交
65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81
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);
82

M
Megvii Engine Team 已提交
83 84 85 86 87 88 89 90 91 92 93
    /*!
     * \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 {
94
        /*!
M
Megvii Engine Team 已提交
95 96
         * \brief special device number for the "cpu default" comp node,
         *      which dispatches all tasks in the caller thread
97
         */
M
Megvii Engine Team 已提交
98 99 100 101 102 103 104
        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;
105

M
Megvii Engine Team 已提交
106
        DeviceType type = DeviceType::UNSPEC;
107

M
Megvii Engine Team 已提交
108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123
        /*!
         * 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;
124 125 126
        };

        /*!
M
Megvii Engine Team 已提交
127 128 129 130
         * \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.
131
         */
132
        MGE_WIN_DECLSPEC_FUC static Locator parse(const std::string& id);
133 134

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

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

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

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

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

M
Megvii Engine Team 已提交
165 166
    struct LocatorPairHashKey {
        Locator locator, locator_logical;
167

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

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

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

M
Megvii Engine Team 已提交
185
    CompNode() = default;
186

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

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

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

205
    MGE_WIN_DECLSPEC_FUC static CompNode load(
M
Megvii Engine Team 已提交
206
            const Locator& locator_physical, const Locator& locator_logical);
207

M
Megvii Engine Team 已提交
208
    /* =================== memory management ======================== */
209

M
Megvii Engine Team 已提交
210 211 212 213 214 215 216 217 218 219
    /*!
     * \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.
     */
220
    MGE_WIN_DECLSPEC_FUC void* alloc_device(size_t size) const;
221

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

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

233
    MGE_WIN_DECLSPEC_FUC void free_host(void* ptr) const;
234

M
Megvii Engine Team 已提交
235 236 237 238
    //! 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);
    }
239

M
Megvii Engine Team 已提交
240 241 242 243
    //! 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);
    }
244

M
Megvii Engine Team 已提交
245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 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);
    }

    //! 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
     */
273
    MGE_WIN_DECLSPEC_FUC static void try_coalesce_all_free_memory();
M
Megvii Engine Team 已提交
274 275 276 277 278

    /*
     * \brief specifies how to pre-allocate from raw dev allocator
     *
     */
279
    MGE_WIN_DECLSPEC_FUC static void set_prealloc_config(
M
Megvii Engine Team 已提交
280 281 282 283 284
            size_t alignment, size_t min_req, size_t max_overhead, double growth_factor,
            DeviceType device_type);
    /*!
     * \brief get compute capability of the specified device
     */
285 286
    MGE_WIN_DECLSPEC_FUC static size_t get_compute_capability(
            int dev, DeviceType device_type);
M
Megvii Engine Team 已提交
287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308

    /* =================== 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
     */
309
    MGE_WIN_DECLSPEC_FUC static void sync_all();
M
Megvii Engine Team 已提交
310 311 312 313 314 315 316 317 318 319 320 321 322 323 324 325 326 327 328

    /* =================== 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();
    }
329

330
#if !MGB_BUILD_SLIM_SERVING
M
Megvii Engine Team 已提交
331 332 333 334
    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);
    }
335

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

M
Megvii Engine Team 已提交
338 339 340
    size_t get_max_block_size_available() const {
        return m_impl->get_max_block_size_available();
    }
341 342

    size_t get_free_mem() const { return m_impl->get_free_mem(); }
343 344
#endif

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

M
Megvii Engine Team 已提交
348 349 350 351 352 353 354 355
    //! 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";
    }
356

M
Megvii Engine Team 已提交
357 358 359 360
    //! get string representation of physical device
    std::string to_string_physical() const {
        return m_impl ? m_impl->locator().to_string() : "invalid";
    }
361

M
Megvii Engine Team 已提交
362 363 364 365
    //! get string representation of logical device
    std::string to_string_logical() const {
        return m_impl ? m_impl->locator_logical().to_string() : "invalid";
    }
366

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

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

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

M
Megvii Engine Team 已提交
375
    //! see CompNodeEnv::activate
376
    MGE_WIN_DECLSPEC_FUC void activate() const;
377

M
Megvii Engine Team 已提交
378
    //! get device type of this comp node
379
    MGE_WIN_DECLSPEC_FUC DeviceType device_type() const;
380

M
Megvii Engine Team 已提交
381 382 383 384 385 386 387 388 389
    /*!
     * \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
390
    MGE_WIN_DECLSPEC_FUC std::unique_ptr<MegBrainError> check_async_error() const;
391

M
Megvii Engine Team 已提交
392 393 394 395 396 397 398 399 400 401 402 403 404 405
    /*!
     * \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);
    }
406

M
Megvii Engine Team 已提交
407 408 409 410 411 412 413 414 415
    /*!
     *  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));
    }
416

417 418 419 420 421 422 423 424 425 426
    /*!
     *  enable dispatcher
     */
    void enable_dispatch() { m_impl->enable_dispatch(); }

    /*!
     *  disable dispatcher so that task will be done inplace
     */
    void disable_dispatch(bool* flag) { m_impl->disable_dispatch(flag); }

M
Megvii Engine Team 已提交
427 428 429 430
    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,
431

M
Megvii Engine Team 已提交
432 433 434 435 436
        //! 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,
437

M
Megvii Engine Team 已提交
438 439 440 441 442 443
        //! 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,
444

M
Megvii Engine Team 已提交
445 446 447
        //! Whether this comp node supports copy stream, so computation and
        //! I/O can be parallelized
        HAS_COPY_STREAM = 1 << 3,
448

M
Megvii Engine Team 已提交
449 450 451 452
        //! 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,
453

M
Megvii Engine Team 已提交
454 455 456 457
        //! 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,
458

M
Megvii Engine Team 已提交
459 460 461 462
        //! Whether this comp node supports unified address. i.e. CPU and
        //! CUDA supports unified address.
        SUPPORT_UNIFIED_ADDRESS = 1 << 6,
    };
463

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

466
    MGE_WIN_DECLSPEC_FUC static bool contain_flag(DeviceType device_type, Flag flag);
467

M
Megvii Engine Team 已提交
468
    using UnorderedSet = ThinHashSet<CompNode>;
469

M
Megvii Engine Team 已提交
470 471 472 473
    template <typename T>
    using UnorderedMap = ThinHashMap<CompNode, T>;

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

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

    /* =================== specialized ======================== */
481

M
Megvii Engine Team 已提交
482 483
    //! get default CPU comp node
    // implemented in comp_node/cpu/comp_node.cpp
484
    MGE_WIN_DECLSPEC_FUC static CompNode default_cpu();
485

M
Megvii Engine Team 已提交
486 487 488 489 490 491 492 493 494 495 496
    /*!
     * \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
     */
497
    MGE_WIN_DECLSPEC_FUC static bool enable_affinity_for_cpu(bool flag);
498

M
Megvii Engine Team 已提交
499 500 501
protected:
    //! ImplBase with env(); defined in CompNodeEnv
    class Impl;
502

M
Megvii Engine Team 已提交
503 504 505 506 507 508 509 510 511 512 513 514 515 516 517 518 519
    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;
        virtual void peer_copy_to(
                Impl* dest_impl, void* dest, const void* src, size_t size) = 0;
520

M
Megvii Engine Team 已提交
521 522
        virtual size_t get_mem_addr_alignment() = 0;
        virtual size_t get_mem_padding();
523

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

M
Megvii Engine Team 已提交
526 527 528 529
        virtual void sync() = 0;

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

531
#if !MGB_BUILD_SLIM_SERVING
M
Megvii Engine Team 已提交
532 533 534 535 536
        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; }
        virtual size_t get_max_block_size_available() { return 0; }
537
        virtual size_t get_free_mem() { return 0; }
538 539
#endif

M
Megvii Engine Team 已提交
540 541 542 543 544
        virtual Locator locator() = 0;
        virtual Locator locator_logical() = 0;

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

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

548 549 550 551
        virtual void enable_dispatch();

        virtual void disable_dispatch(bool* flag);

M
Megvii Engine Team 已提交
552 553 554
        virtual uint64_t get_uid() {
            mgb_throw(MegBrainError, "get_uid is not impl yet");
        };
555

M
Megvii Engine Team 已提交
556 557
    protected:
        ImplBase(free_func_t fd, free_func_t fh) : free_device{fd}, free_host{fh} {}
558

M
Megvii Engine Team 已提交
559 560
        ~ImplBase() = default;
    };
561

M
Megvii Engine Team 已提交
562 563 564
    //! implementations are allocated statically, so no memory management
    //! is needed
    ImplBase* m_impl = nullptr;
565

M
Megvii Engine Team 已提交
566 567 568 569
    friend class CompNodeEnv;
    friend struct HashTrait<CompNode>;
    friend struct HashTrait<CompNode::Locator>;
    friend class CompNodeImplHelper;
570

M
Megvii Engine Team 已提交
571 572
public:
    CompNode(ImplBase* impl) : m_impl{impl} {}
573 574 575 576
};

MGB_DEF_ENUM_CLASS_BIT_OPR(CompNode::Flag)

577 578 579 580 581 582 583 584 585 586 587 588 589 590 591 592 593 594 595 596 597 598 599 600 601 602 603 604 605 606 607 608 609 610 611 612
/*!
 * \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;
613

614 615 616
    virtual void replay() = 0;
};

617 618 619 620
/*!
 * \brief event associated with a CompNode node, used for cross-device
 *      synchronization
 */
M
Megvii Engine Team 已提交
621 622 623
class CompNode::Event : public NonCopyableObj {
protected:
    static int sm_cpu_sync_level;
624

M
Megvii Engine Team 已提交
625 626
    //! flags when this event is created
    size_t const m_create_flags;
627

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

M
Megvii Engine Team 已提交
630 631
public:
    enum Flags { NEED_TIMER = 1 };
632

M
Megvii Engine Team 已提交
633
    virtual ~Event() = default;
634

M
Megvii Engine Team 已提交
635 636 637 638 639 640 641 642 643
    /*!
     * \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;
644

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

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

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

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

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

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

M
Megvii Engine Team 已提交
665 666 667 668 669 670 671 672
    /*!
     * \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; }
673 674 675 676 677 678 679 680 681 682
};

/*!
 * \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;
683
    size_t m_flags;
684

M
Megvii Engine Team 已提交
685
public:
686 687
    MGE_WIN_DECLSPEC_FUC explicit EventPool(CompNode cn, size_t flags = 0);
    MGE_WIN_DECLSPEC_FUC ~EventPool();
688

689
    MGE_WIN_DECLSPEC_FUC CompNode::Event* alloc();
690

691
    MGE_WIN_DECLSPEC_FUC void free(CompNode::Event* ev);
692

M
Megvii Engine Team 已提交
693
    //! assert that all allocated events have been freed
694
    MGE_WIN_DECLSPEC_FUC void assert_all_freed();
695 696
};

M
Megvii Engine Team 已提交
697
void CompNode::device_wait_event(Event& event) const {
698 699 700
    event.device_wait_by(*this);
}

M
Megvii Engine Team 已提交
701
template <>
702
struct HashTrait<CompNode> {
M
Megvii Engine Team 已提交
703
    static size_t eval(const CompNode& val) {
704 705 706 707 708
        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 已提交
709
template <>
710
struct HashTrait<CompNode::Locator> {
M
Megvii Engine Team 已提交
711 712 713
    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);
714 715 716
    }
};

717 718 719 720 721 722 723 724 725 726 727 728 729 730 731 732 733 734 735 736 737
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:
738
    MGE_WIN_DECLSPEC_FUC virtual std::shared_ptr<void> callback() = 0;
739 740
    ~DepedentObjList() = default;

741 742
    MGE_WIN_DECLSPEC_FUC static void add(DepedentObjList* ptr);
    MGE_WIN_DECLSPEC_FUC static void remove(DepedentObjList* ptr);
743 744

public:
745
    MGE_WIN_DECLSPEC_FUC static void invoke_callback_and_clean();
746 747 748 749 750 751 752 753 754 755 756 757 758 759 760 761 762 763 764 765 766 767 768 769
};

}  // 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;
770
    MGE_WIN_DECLSPEC_FUC std::shared_ptr<void> callback() override final;
771 772 773 774 775 776 777 778 779 780 781 782 783 784 785 786 787 788 789 790 791 792 793

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 已提交
794
}  // namespace mgb
795 796

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