comp_node.h 27.5 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

245 246 247 248 249 250 251 252 253 254 255 256 257 258
    //! 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 已提交
259 260 261 262 263 264 265 266 267 268 269
    /*!
     * \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);
    }

270 271 272 273 274 275 276 277
    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 已提交
278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294
    //! 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
     */
295
    MGE_WIN_DECLSPEC_FUC static void try_coalesce_all_free_memory();
M
Megvii Engine Team 已提交
296 297 298 299 300

    /*
     * \brief specifies how to pre-allocate from raw dev allocator
     *
     */
301
    MGE_WIN_DECLSPEC_FUC static void set_prealloc_config(
M
Megvii Engine Team 已提交
302 303 304 305 306
            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
     */
307 308
    MGE_WIN_DECLSPEC_FUC static size_t get_compute_capability(
            int dev, DeviceType device_type);
M
Megvii Engine Team 已提交
309 310 311 312 313 314 315 316 317 318 319 320 321 322 323 324 325 326 327 328 329 330

    /* =================== 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
     */
331
    MGE_WIN_DECLSPEC_FUC static void sync_all();
M
Megvii Engine Team 已提交
332 333 334 335 336 337 338 339 340 341 342 343 344 345 346 347 348 349 350

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

352
#if !MGB_BUILD_SLIM_SERVING
M
Megvii Engine Team 已提交
353 354 355 356
    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);
    }
357

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

360 361 362 363 364 365
    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 已提交
366 367 368
    size_t get_max_block_size_available() const {
        return m_impl->get_max_block_size_available();
    }
369 370

    size_t get_free_mem() const { return m_impl->get_free_mem(); }
371 372 373 374 375 376

    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(); }
377 378
#endif

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

M
Megvii Engine Team 已提交
382 383 384 385 386 387 388 389
    //! 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";
    }
390

M
Megvii Engine Team 已提交
391 392 393 394
    //! get string representation of physical device
    std::string to_string_physical() const {
        return m_impl ? m_impl->locator().to_string() : "invalid";
    }
395

M
Megvii Engine Team 已提交
396 397 398 399
    //! get string representation of logical device
    std::string to_string_logical() const {
        return m_impl ? m_impl->locator_logical().to_string() : "invalid";
    }
400

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

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

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

M
Megvii Engine Team 已提交
409
    //! see CompNodeEnv::activate
410
    MGE_WIN_DECLSPEC_FUC void activate() const;
411

M
Megvii Engine Team 已提交
412
    //! get device type of this comp node
413
    MGE_WIN_DECLSPEC_FUC DeviceType device_type() const;
414

M
Megvii Engine Team 已提交
415 416 417 418 419 420 421 422 423
    /*!
     * \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
424
    MGE_WIN_DECLSPEC_FUC std::unique_ptr<MegBrainError> check_async_error() const;
425

M
Megvii Engine Team 已提交
426 427 428 429 430 431 432 433 434 435 436 437 438 439
    /*!
     * \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);
    }
440

M
Megvii Engine Team 已提交
441 442 443 444 445 446 447 448 449
    /*!
     *  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));
    }
450

M
Megvii Engine Team 已提交
451 452 453 454
    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,
455

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

M
Megvii Engine Team 已提交
462 463 464 465 466 467
        //! 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,
468

M
Megvii Engine Team 已提交
469 470 471
        //! Whether this comp node supports copy stream, so computation and
        //! I/O can be parallelized
        HAS_COPY_STREAM = 1 << 3,
472

M
Megvii Engine Team 已提交
473 474 475 476
        //! 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,
477

M
Megvii Engine Team 已提交
478 479 480 481
        //! 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,
482

M
Megvii Engine Team 已提交
483 484 485 486
        //! Whether this comp node supports unified address. i.e. CPU and
        //! CUDA supports unified address.
        SUPPORT_UNIFIED_ADDRESS = 1 << 6,
    };
487

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

490
    MGE_WIN_DECLSPEC_FUC static bool contain_flag(DeviceType device_type, Flag flag);
491

M
Megvii Engine Team 已提交
492
    using UnorderedSet = ThinHashSet<CompNode>;
493

M
Megvii Engine Team 已提交
494 495 496 497
    template <typename T>
    using UnorderedMap = ThinHashMap<CompNode, T>;

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

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

    /* =================== specialized ======================== */
505

M
Megvii Engine Team 已提交
506 507
    //! get default CPU comp node
    // implemented in comp_node/cpu/comp_node.cpp
508
    MGE_WIN_DECLSPEC_FUC static CompNode default_cpu();
509

M
Megvii Engine Team 已提交
510 511 512 513 514 515 516 517 518 519 520
    /*!
     * \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
     */
521
    MGE_WIN_DECLSPEC_FUC static bool enable_affinity_for_cpu(bool flag);
522

M
Megvii Engine Team 已提交
523 524 525
protected:
    //! ImplBase with env(); defined in CompNodeEnv
    class Impl;
526

M
Megvii Engine Team 已提交
527 528 529 530 531 532 533 534 535 536 537 538 539 540 541
    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;
542 543 544 545 546 547 548 549 550 551
        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 已提交
552 553
        virtual void peer_copy_to(
                Impl* dest_impl, void* dest, const void* src, size_t size) = 0;
554

555 556 557 558 559 560
        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);
        }

M
Megvii Engine Team 已提交
561 562
        virtual size_t get_mem_addr_alignment() = 0;
        virtual size_t get_mem_padding();
563

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

M
Megvii Engine Team 已提交
566 567 568 569
        virtual void sync() = 0;

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

571
#if !MGB_BUILD_SLIM_SERVING
M
Megvii Engine Team 已提交
572 573 574 575
        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; }
576 577 578
        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 已提交
579
        virtual size_t get_max_block_size_available() { return 0; }
580
        virtual size_t get_free_mem() { return get_mem_status_bytes().second; }
581 582
        virtual void reset_max_reserved_memory() {}
        virtual void reset_max_used_memory() {}
583 584
#endif

M
Megvii Engine Team 已提交
585 586 587 588 589
        virtual Locator locator() = 0;
        virtual Locator locator_logical() = 0;

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

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

M
Megvii Engine Team 已提交
593 594 595
        virtual uint64_t get_uid() {
            mgb_throw(MegBrainError, "get_uid is not impl yet");
        };
596

M
Megvii Engine Team 已提交
597 598
    protected:
        ImplBase(free_func_t fd, free_func_t fh) : free_device{fd}, free_host{fh} {}
599

M
Megvii Engine Team 已提交
600 601
        ~ImplBase() = default;
    };
602

M
Megvii Engine Team 已提交
603 604 605
    //! implementations are allocated statically, so no memory management
    //! is needed
    ImplBase* m_impl = nullptr;
606

M
Megvii Engine Team 已提交
607 608 609 610
    friend class CompNodeEnv;
    friend struct HashTrait<CompNode>;
    friend struct HashTrait<CompNode::Locator>;
    friend class CompNodeImplHelper;
611

M
Megvii Engine Team 已提交
612 613
public:
    CompNode(ImplBase* impl) : m_impl{impl} {}
614 615 616 617
};

MGB_DEF_ENUM_CLASS_BIT_OPR(CompNode::Flag)

618 619 620 621 622 623 624 625 626 627 628 629 630 631 632 633 634 635 636 637 638 639 640 641 642 643 644 645 646 647 648 649 650 651 652 653
/*!
 * \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;
654

655 656 657
    virtual void replay() = 0;
};

658 659 660 661
/*!
 * \brief event associated with a CompNode node, used for cross-device
 *      synchronization
 */
M
Megvii Engine Team 已提交
662 663 664
class CompNode::Event : public NonCopyableObj {
protected:
    static int sm_cpu_sync_level;
665

M
Megvii Engine Team 已提交
666 667
    //! flags when this event is created
    size_t const m_create_flags;
668

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

M
Megvii Engine Team 已提交
671 672
public:
    enum Flags { NEED_TIMER = 1 };
673

M
Megvii Engine Team 已提交
674
    virtual ~Event() = default;
675

M
Megvii Engine Team 已提交
676 677 678 679 680 681 682 683 684
    /*!
     * \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;
685

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

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

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

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

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

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

M
Megvii Engine Team 已提交
706 707 708 709 710 711 712 713
    /*!
     * \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; }
714 715 716 717 718 719 720 721 722 723
};

/*!
 * \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;
724
    size_t m_flags;
725

M
Megvii Engine Team 已提交
726
public:
727 728
    MGE_WIN_DECLSPEC_FUC explicit EventPool(CompNode cn, size_t flags = 0);
    MGE_WIN_DECLSPEC_FUC ~EventPool();
729

730
    MGE_WIN_DECLSPEC_FUC CompNode::Event* alloc();
731

732
    MGE_WIN_DECLSPEC_FUC void free(CompNode::Event* ev);
733

M
Megvii Engine Team 已提交
734
    //! assert that all allocated events have been freed
735
    MGE_WIN_DECLSPEC_FUC void assert_all_freed();
736 737
};

M
Megvii Engine Team 已提交
738
void CompNode::device_wait_event(Event& event) const {
739 740 741
    event.device_wait_by(*this);
}

M
Megvii Engine Team 已提交
742
template <>
743
struct HashTrait<CompNode> {
M
Megvii Engine Team 已提交
744
    static size_t eval(const CompNode& val) {
745 746 747 748 749
        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 已提交
750
template <>
751
struct HashTrait<CompNode::Locator> {
M
Megvii Engine Team 已提交
752 753 754
    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);
755 756 757
    }
};

758 759 760 761 762 763 764 765 766 767 768 769 770 771 772 773 774 775 776 777 778
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:
779
    MGE_WIN_DECLSPEC_FUC virtual std::shared_ptr<void> callback() = 0;
780 781
    ~DepedentObjList() = default;

782 783
    MGE_WIN_DECLSPEC_FUC static void add(DepedentObjList* ptr);
    MGE_WIN_DECLSPEC_FUC static void remove(DepedentObjList* ptr);
784 785

public:
786
    MGE_WIN_DECLSPEC_FUC static void invoke_callback_and_clean();
787 788 789 790 791 792 793 794 795 796 797 798 799 800 801 802 803 804 805 806 807 808 809 810
};

}  // 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;
811
    MGE_WIN_DECLSPEC_FUC std::shared_ptr<void> callback() override final;
812 813 814 815 816 817 818 819 820 821 822 823 824 825 826 827 828 829 830 831 832 833 834

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 已提交
835
}  // namespace mgb
836 837

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