interpreter_impl.cpp 29.4 KB
Newer Older
M
Megvii Engine Team 已提交
1
/**
2
 * \file imperative/src/impl/interpreter/interpreter_impl.cpp
M
Megvii Engine Team 已提交
3 4
 * MegEngine is Licensed under the Apache License, Version 2.0 (the "License")
 *
5
 * Copyright (c) 2014-2021 Megvii Inc. All rights reserved.
M
Megvii Engine Team 已提交
6 7 8 9 10 11
 *
 * 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.
 */

12
#include "./interpreter_impl.h"
13

14
#include "megbrain/common.h"
15 16 17
#include "megbrain/imperative/opr_utility.h"
#include "megbrain/imperative/ops/backward_graph.h"
#include "megbrain/imperative/ops/autogen.h"
18 19 20
#include "megbrain/imperative/utils/to_string.h"

#include "../op_trait.h"
21 22 23 24 25 26 27 28 29 30 31 32 33 34 35

using namespace mgb;
using namespace imperative;
using namespace interpreter;
using namespace interpreter::intl;

std::unique_ptr<Interpreter::Channel> InterpreterImpl::create_channel() {
    return std::make_unique<ChannelImpl>();
}

Interpreter& Interpreter::inst() {
    static InterpreterImpl inst_;
    return inst_;
}

36
Handle ChannelImpl::put(const HostTensorND& value, bool no_cache) {
37 38 39 40
    auto info = alloc();
    info->desc.layout = value.layout();
    info->desc.comp_node = value.comp_node();
    info->desc.value = value.proxy_to_default_cpu();
41
    info->h_value = value;
42
    m_buffer.enqueue(Put{info, value, no_cache});
43 44 45 46
    if (m_async_level == 0) {
        sync();
        info->desc.comp_node.sync();
    }
47 48 49
    return info;
}

50
Handle ChannelImpl::put(const DeviceTensorND& data) {
M
Megvii Engine Team 已提交
51 52 53 54
    auto info = alloc();
    info->desc.layout = data.layout();
    info->desc.comp_node = data.comp_node();
    info->ptr = Tensor::make(data);
55
    m_channel_state.profiler->record_host<TensorProduceEvent>(info->id, info->desc.layout, info->desc.comp_node);
M
Megvii Engine Team 已提交
56 57 58
    return info;
}

59
void ChannelImpl::del(Handle handle) {
60 61 62 63 64 65
    mgb_assert(m_valid_handle.count(handle), "invalid handle: %p", handle);
    auto* info = reinterpret_cast<TensorInfo*>(handle);
    detach_users(info);
    info->detach_producer();
    m_valid_handle.erase(handle);
    m_buffer.enqueue(Del{info});
66 67
}

68
void ChannelImpl::swap_in(Handle handle) {
69
    if (m_worker_state.options.enable_swap) {
70 71
        mgb_assert(m_valid_handle.find(handle) != m_valid_handle.end(),
                "invalid handle: %p", handle);
72 73 74
        auto* info = reinterpret_cast<TensorInfo*>(handle);
        m_buffer.enqueue(SwapIn{info});
        info->evict_type = NONE;
75 76 77
    }
}

78
void ChannelImpl::swap_out(Handle handle) {
79
    if (m_worker_state.options.enable_swap) {
80 81
        mgb_assert(m_valid_handle.find(handle) != m_valid_handle.end(),
                "invalid handle: %p", handle);
82 83 84
        auto* info = reinterpret_cast<TensorInfo*>(handle);
        m_buffer.enqueue(SwapOut{info});
        info->evict_type = SWAP;
85 86 87
    }
}

88
void ChannelImpl::drop(Handle handle) {
89
    if (m_worker_state.options.enable_drop) {
90 91
        mgb_assert(m_valid_handle.find(handle) != m_valid_handle.end(),
                "invalid handle: %p", handle);
92 93 94 95 96 97 98
        auto* info = reinterpret_cast<TensorInfo*>(handle);
        if (!info->producer) {
            mgb_log_warn("the input that produced tensor %p has been deleted, this drop operation will be ignored", info);
            return;
        }
        info->evict_type = DROP;
        m_buffer.enqueue(Drop{info});
99 100 101
    }
}

102
void ChannelImpl::dispatch_default_cpu(
103
        std::shared_ptr<OpDef> op,
104 105 106 107
        const SmallVector<TensorInfo*>& input_infos,
        const SmallVector<LogicalTensorDesc>& input_descs,
        SmallVector<Handle>* outputs) {
    auto [output_descs, validated] = OpDef::infer_output_attrs_fallible(*op, input_descs);
108
    MGB_MARK_USED_VAR(validated);
109

110 111 112
    SmallVector<DeviceTensorND> input_tensornds;
    input_tensornds.reserve(input_descs.size());
    CompNode output_cn;
113 114
    {
        MGB_LOCK_GUARD(m_mutex);
115
        for (auto&& info : input_infos) {
116
            auto input_cn = info->desc.comp_node;
117
            if (!output_cn.valid()) {
118 119 120 121 122 123 124
                output_cn = input_cn;
            } else {
                mgb_assert(output_cn == input_cn, "cannot decide output comp node");
            }

            if (info->ptr && info->ptr->try_get_value()) {
                input_tensornds.emplace_back(info->ptr->get_value().proxy_to_default_cpu());
125
            } else {
126 127
                mgb_assert(!info->h_value.empty(), "inp->h_value is empty!");
                input_tensornds.emplace_back(info->h_value.proxy_to_default_cpu());
128 129 130 131 132 133 134 135 136 137 138 139 140 141
            }
        }
    }

    outputs->reserve(output_descs.size());
    SmallVector<DeviceTensorND> output_tensornds;
    output_tensornds.reserve(output_descs.size());
    for (auto&& desc : output_descs) {
        // TODO: may conflict with condtake, which need alloc inside
        mgb_assert(!desc.layout.is_empty());
        // use HostTensorND alloc_host for cuda pinned memory
        output_tensornds.emplace_back(HostTensorND(output_cn, desc.layout).proxy_to_default_cpu());
    }

142 143 144 145 146 147 148 149 150 151 152
    auto tinfo_to_tid = [&](SmallVector<TensorInfo*> tinfo) {
        SmallVector<uint64_t> tid;
        for (auto* ptinfo: tinfo) {
            tid.push_back(ptinfo->id);
        }
        return tid;
    };
    OpEvent event_data = {++m_last_id, op, tinfo_to_tid(input_infos), {}};

    m_channel_state.profiler->record_host<HostOpExecuteEvent>(event_data);

153 154 155 156 157 158 159
    OpDef::apply_on_device_tensornd(*op, input_tensornds, &output_tensornds);

    SmallVector<TensorInfo*> output_infos;
    output_infos.reserve(output_descs.size());
    for (auto&& tensornd : output_tensornds) {
        HostTensorND host_tensornd = HostTensorND::make_proxy(tensornd)
            .proxy_to_comp_node(output_cn);
160 161 162
        // use `put` for consistency
        auto info = reinterpret_cast<TensorInfo*>(put(host_tensornd, false));
        mgb_assert(info->desc.layout.ndim != 0);
163 164 165
        output_infos.push_back(info);
        outputs->push_back(info);
    }
166 167

    if (m_channel_state.options.enable_drop) {
168
        TensorInfo::ComputePath::make(op, input_infos, output_infos);
169
    }
170 171 172 173

    event_data.outputs = tinfo_to_tid(output_infos);

    m_channel_state.profiler->record_host<HostOpFinishEvent>(event_data);
174
}
175

176 177 178 179 180
void ChannelImpl::dispatch_kernel(
        std::shared_ptr<OpDef> op,
        const SmallVector<TensorInfo*>& input_infos,
        const SmallVector<LogicalTensorDesc>& input_descs,
        SmallVector<Handle>* outputs) {
181
    auto [output_descs, validated] = OpDef::infer_output_attrs_fallible(*op, input_descs);
182

183
    ApplyOp cmd{std::move(op)};
184
    cmd.inputs = std::move(input_infos);
185
    cmd.outputs.reserve(output_descs.size());
186 187
    outputs->reserve(output_descs.size());
    for (auto&& desc : output_descs) {
188 189
        auto info = alloc();
        info->desc = desc;
190 191 192 193 194
        // make sure desc's value is consistent with h_value
        if (!info->desc.value.empty()) {
            info->h_value = HostTensorND::make_proxy(desc.value)
                .proxy_to_comp_node(desc.comp_node);
        }
195
        cmd.outputs.push_back(info);
196
        outputs->push_back(info);
197
    }
198
    if (m_channel_state.options.enable_drop) {
199
        TensorInfo::ComputePath::make(cmd.op, cmd.inputs, cmd.outputs);
200
    }
201
    m_buffer.enqueue(std::move(cmd));
202
    if (!validated && m_channel_state.options.async_level == 1) {
203
        sync();
204
    } else if (m_channel_state.options.async_level == 0) {
205
        sync();
206
        // check device error
207
        for (auto&& oup : *outputs) {
208 209
            auto info = reinterpret_cast<TensorInfo*>(oup);
            info->ptr->comp_node().sync();
210
        }
211
    }
212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231
}

SmallVector<Handle> ChannelImpl::apply_op(
        std::shared_ptr<OpDef> op,
        const SmallVector<Handle>& inputs) {
    for (auto i : inputs) {
        mgb_assert(m_valid_handle.find(i) != m_valid_handle.end(),
                "invalid handle: %p", i);
    }
    SmallVector<TensorInfo*> input_infos;
    input_infos.reserve(inputs.size());
    SmallVector<LogicalTensorDesc> input_descs;
    input_descs.reserve(inputs.size());
    {
        MGB_LOCK_GUARD(m_mutex);
        for (auto i : inputs) {
            auto info = reinterpret_cast<TensorInfo*>(i);
            mgb_assert(!info->invalid, "Invalid tensor, unable to apply_op!");
            input_infos.push_back(info);
            input_descs.push_back(info->desc);
232
            regenerate(info);
233 234 235 236
        }
    }

    SmallVector<Handle> outputs;
237 238 239 240
    DispatchMode dispatch_mode = m_channel_state.options.enable_host_compute
            ? OpDef::decide_dispatch_mode(*op, input_descs)
            : DispatchMode::KERNEL;
    switch (dispatch_mode) {
241 242 243 244 245 246 247 248 249
        case DEFAULT_CPU: {
            dispatch_default_cpu(op, input_infos, input_descs, &outputs);
            break;
        }
        case KERNEL: {
            dispatch_kernel(op, input_infos, input_descs, &outputs);
            break;
        }
    }
250 251 252
    return outputs;
}

253
HostTensorND ChannelImpl::get_value(Handle handle) {
254
    // TODO: maybe get_value should be done on host. i.e. delete GetValue
255 256 257 258
    mgb_assert(m_valid_handle.find(handle) != m_valid_handle.end(),
               "invalid handle: %p", handle);
    auto info = reinterpret_cast<TensorInfo*>(handle);
    mgb_assert(!m_waitee);
259 260
    // donnot use info->value_fetched, it's unsafe
    mgb_assert(!info->invalid, "Invalid tensor, unable to get_value!");
261
    std::unique_lock<decltype(m_mutex)> lock(m_mutex);
262 263 264 265 266
    TensorPtr tensor_ptr = info->ptr;
    auto value_fetched = [&]() {
        return tensor_ptr && tensor_ptr->value_fetched();
    };
    if (!value_fetched()) {
267
        m_waitee = info;
268
        regenerate(info);
269
        m_buffer.enqueue(GetValue{info});
270
        m_channel_state.profiler->record_host<TensorWaitPropEvent>(info->id, TensorInfo::HostValue);
271 272
        m_cv.wait(lock, [&]() {
            check_worker_exc_unsafe();
273 274
            tensor_ptr = info->ptr;
            return value_fetched();
275
        });
276
        m_channel_state.profiler->record_host<TensorWaitPropFinishEvent>(info->id, TensorInfo::HostValue);
277 278
        m_waitee = nullptr;
    }
279
    return tensor_ptr->get_value();
280 281
}

282
TensorShape ChannelImpl::get_shape(Handle handle) {
283 284 285 286 287 288 289 290 291
    mgb_assert(m_valid_handle.find(handle) != m_valid_handle.end(),
               "invalid handle: %p", handle);
    auto info = reinterpret_cast<TensorInfo*>(handle);
    if (info->desc.layout.ndim != 0) {
        return info->desc.layout;
    }
    std::unique_lock<decltype(m_mutex)> lock(m_mutex);
    mgb_assert(!m_waitee);
    m_waitee = info;
292 293
    m_buffer.flush();
    m_channel_state.profiler->record_host<TensorWaitPropEvent>(info->id, TensorInfo::Shape);
294 295
    m_cv.wait(lock, [&]() {
        check_worker_exc_unsafe();
296
        return static_cast<bool>(info->ptr);
297
    });
298
    m_channel_state.profiler->record_host<TensorWaitPropFinishEvent>(info->id, TensorInfo::Shape);
299 300 301 302 303 304
    m_waitee = nullptr;
    TensorShape ret = info->ptr->layout();
    mgb_assert(ret.ndim != 0);
    return ret;
}

305
DType ChannelImpl::get_dtype(Handle handle) {
306 307 308
    mgb_assert(m_valid_handle.find(handle) != m_valid_handle.end(),
               "invalid handle: %p", handle);
    auto info = reinterpret_cast<TensorInfo*>(handle);
309
    m_channel_state.profiler->record_host<TensorGetPropEvent>(info->id, TensorInfo::DType);
310 311 312 313 314
    auto ret = info->desc.layout.dtype;
    mgb_assert(ret.valid());
    return ret;
}

315
CompNode ChannelImpl::get_device(Handle handle) {
316 317 318
    mgb_assert(m_valid_handle.find(handle) != m_valid_handle.end(),
               "invalid handle: %p", handle);
    auto info = reinterpret_cast<TensorInfo*>(handle);
319
    m_channel_state.profiler->record_host<TensorGetPropEvent>(info->id, TensorInfo::Device);
320 321 322 323 324
    auto ret = info->desc.comp_node;
    mgb_assert(ret.valid());
    return ret;
}

325
DeviceTensorND ChannelImpl::get_dev_tensor(Handle handle) {
326 327 328 329 330 331
    mgb_assert(m_valid_handle.find(handle) != m_valid_handle.end(),
               "invalid handle: %p", handle);
    auto info = reinterpret_cast<TensorInfo*>(handle);
    std::unique_lock<decltype(m_mutex)> lock(m_mutex);
    mgb_assert(!m_waitee);
    m_waitee = info;
332
    regenerate(info);
333 334
    m_buffer.flush();
    m_channel_state.profiler->record_host<TensorWaitPropEvent>(info->id, TensorInfo::DevValue);
335 336
    m_cv.wait(lock, [&]() {
        check_worker_exc_unsafe();
337
        return static_cast<bool>(info->ptr);
338
    });
339
    m_channel_state.profiler->record_host<TensorWaitPropFinishEvent>(info->id, TensorInfo::DevValue);
340 341 342 343 344
    m_waitee = nullptr;
    return info->ptr->dev_tensor();
}

void ChannelImpl::sync() {
345 346
    m_buffer.flush();
    m_channel_state.profiler->record_host<SyncStartEvent>();
347
    m_worker.wait_all_task_finish();
348 349
    CompNode::sync_all();
    m_channel_state.profiler->record_host<SyncFinishEvent>();
350 351 352 353 354 355 356 357
    MGB_LOCK_GUARD(m_mutex);
    check_worker_exc_unsafe();
}

void ChannelImpl::close() {
    sync();
}

358 359
int ChannelImpl::get_option(std::string name) {
    return m_channel_state.options.get_option(name);
360 361
}

362 363 364
void ChannelImpl::set_option(std::string name, int value) {
    m_channel_state.options.set_option(name, value);
    m_buffer.enqueue(SetOption{name, value});
365 366 367 368
}

TensorInfo* ChannelImpl::alloc() {
    MGB_LOCK_GUARD(m_mutex);
369
    auto info = m_pool.alloc();
370
    m_valid_handle.insert(info);
371 372
    info->id = m_last_id++;
    m_channel_state.profiler->record_host<TensorDeclareEvent>(info->id);
373
    return info;
374 375 376 377
}

void ChannelImpl::free(TensorInfo* ptr) {
    MGB_LOCK_GUARD(m_mutex);
378
    m_channel_state.profiler->record_host<TensorEraseEvent>(ptr->id);
379 380 381
    m_pool.free(ptr);
}

382 383 384 385
ChannelImpl::ChannelImpl() : m_worker(this), m_buffer(this){
    m_channel_state.tid = std::this_thread::get_id();
}

386 387 388
ChannelImpl::~ChannelImpl() {
    close();
}
389

390 391
void ChannelImpl::produce_tensor(TensorInfo* dest, TensorPtr ptr) {
    MGB_LOCK_GUARD(m_mutex);
392
    m_worker_state.profiler->record_host<TensorProduceEvent>(dest->id, ptr->layout(), ptr->comp_node());
393 394 395 396 397
    dest->value_fetched = ptr->value_fetched();
    // update tensor desc for static infer
    dest->desc.layout = ptr->layout();
    dest->desc.comp_node = ptr->comp_node();
    dest->ptr = std::move(ptr);
398
    if (m_waitee == dest) {
399
        m_cv.notify_all();
400 401 402
    }
}

403 404 405 406 407
void ChannelImpl::release_tensor(TensorInfo* dest) {
    MGB_LOCK_GUARD(m_mutex);
    dest->ptr.reset();
}

408
void ChannelImpl::regenerate(TensorInfo* dest) {
409
    if (dest->evict_type == DROP) {
410 411 412
        recompute(dest->producer);
    } else if (dest->evict_type == SWAP) {
        swap_in(dest);
413
    }
414
    mgb_assert(dest->evict_type == NONE);
415 416
}

417 418 419 420
void ChannelImpl::recompute(TensorInfo::ComputePath* path) {
    SmallVector<TensorInfo*> workspaces(path->outputs.size(), nullptr);
    for (auto&& input: path->inputs) {
        regenerate(input);
421
    }
422 423 424
    for (auto&& output: path->outputs) {
        if(output == nullptr) {
            continue;
425
        }
426
        output->evict_type = NONE;
427
    }
428
    m_buffer.enqueue(ApplyOp{path->op, path->inputs, path->outputs});
429 430
}

431 432 433 434 435 436 437 438 439 440
void ChannelImpl::detach_users(TensorInfo* dest) {
    SmallVector<TensorInfo::ComputePath*> users = dest->users;
    for (auto* user: users) {
        for (auto* output: user->outputs) {
            if (output == nullptr) {
                continue;
            }
            regenerate(output);
            output->detach_producer();
        }
441
    }
442 443
    mgb_assert(dest->users.size() == 0);
    //dest->users.clear();
444 445
}

446 447 448 449 450 451 452 453 454 455 456 457 458 459 460
void ChannelImpl::sync_device_scope(CompNode device) {
    auto& prev = m_worker_state.device_scope_map[device];
    auto& current = m_worker_state.scopes;
    auto push_scope = [&](std::string name) {
        m_worker_state.profiler->record_device<DeviceBeginScope>(device, name);
    };
    auto pop_scope = [&](std::string name) {
        m_worker_state.profiler->record_device<DeviceEndScope>(device, name);
    };
    size_t similarity = 0;
    for (size_t i = 0; i < prev.size() && i < current.size(); i++) {
        if (prev[i] == current[i]) {
            similarity++;
        } else {
            break;
461 462
        }
    }
463 464 465
    while (prev.size() > similarity) {
        pop_scope(prev.back());
        prev.pop_back();
466
    }
467 468 469
    while (prev.size() < current.size()) {
        prev.push_back(current[prev.size()]);
        push_scope(prev.back());
470 471 472
    }
}

473 474 475 476 477 478 479 480 481 482
void ChannelImpl::process_one_task(IdentifiedCommand& icmd) {
    m_worker_state.profiler->record_host<CommandExecuteEvent>(icmd);
    bool finished = false;
    auto do_finish_command = [&]{
        if (finished) {
            return;
        }
        m_worker_state.profiler->record_host<CommandFinishEvent>(icmd);
        finished = true;
    };
483
    //TODO: remove std::visit for support osx 10.12
484 485
    auto cmd_visitor = [&](const auto& cmd) {
            using T = std::decay_t<decltype(cmd)>;
486
            if constexpr (std::is_same_v<T, Put>) {
487 488
                auto value = cmd.no_cache ? std::make_shared<Tensor>(cmd.value) : Tensor::make(cmd.value);
                produce_tensor(cmd.dest, std::move(value));
489
            } else if constexpr (std::is_same_v<T, ApplyOp>) {
490
                uint64_t apply_id = ++m_last_id;
491
                SmallVector<TensorPtr> tensor_inputs;
492
                SmallVector<CompNode> devices;
493
                tensor_inputs.reserve(cmd.inputs.size());
494
                // refcnt == 1, owners: [TensorInfo::ptr]
495
                for (auto i : cmd.inputs) {
496
                    mgb_assert(i->ptr, "Invalid input tensor ptr!");
497
                    // refcnt ++, owners: [i->ptr, tensor_inputs]
498 499
                    tensor_inputs.push_back(i->ptr);
                }
500 501 502 503 504 505 506 507 508 509 510 511 512 513 514 515 516
                // Begin profiling operator
                auto tinfo_to_tid = [&](SmallVector<TensorInfo*> tinfo) {
                    SmallVector<uint64_t> tid;
                    for (auto* ptinfo: tinfo) {
                        tid.push_back(ptinfo->id);
                    }
                    return tid;
                };
                OpEvent event_data = {apply_id, cmd.op, tinfo_to_tid(cmd.inputs), tinfo_to_tid(cmd.outputs)};
                // Collecting devices
                for (auto i : cmd.inputs) {
                    devices.push_back(i->desc.comp_node);
                }
                for (auto i : cmd.outputs) {
                    devices.push_back(i->desc.comp_node);
                }
                devices.erase(std::unique(devices.begin(), devices.end()), devices.end());
517 518 519 520 521 522 523 524
                // Fused by command buffer. @see: CommandBuffer::fuse_del
                // Now if dest is inplacable, it's refcnt would be decreased to 1 and owned by tensor_inputs after Del.
                // Note for exprs like 'y = x op x', inplace is unsupported yet but Del would be also fused.
                for (auto* del : cmd.dels) {
                    // refcnt --, owners: [tensor_inputs]
                    // if it's decreased to 1, would be detected at @see: proxy_graph_detail::apply_on_physical_tensor
                    free(del);
                }
525 526 527
                // Before wait
                //TODO: split operator wait and execute so that OpWait could be corrected recorded.
                // Before execute
528 529 530 531 532 533
                if (m_worker_state.profiler->is_profiling()) {
                    m_worker_state.profiler->record_host<HostOpExecuteEvent>(event_data);
                    for (auto&& device: devices) {
                        sync_device_scope(device);
                        m_worker_state.profiler->record_device<DeviceOpExecuteEvent>(device, event_data);
                    }
534 535
                }
                // Apply op
536 537 538
                // Here std::move is REQUIRED for removing duplicated references.
                auto tensor_outputs = OpDef::apply_on_physical_tensor(
                    *cmd.op, std::move(tensor_inputs));
539
                // After execute
540 541 542 543 544
                if (m_worker_state.profiler->is_profiling()) {
                    m_worker_state.profiler->record_host<HostOpFinishEvent>(event_data);
                    for (auto&& device: devices) {
                        m_worker_state.profiler->record_device<DeviceOpFinishEvent>(device, event_data);
                    }
545 546
                }
                // End profiling operator
547 548
                mgb_assert(tensor_outputs.size() == cmd.outputs.size());
                for (size_t i = 0; i < tensor_outputs.size(); ++i) {
549 550 551
                    if (cmd.outputs[i] == nullptr) {
                        continue;
                    }
552 553 554 555 556
                    produce_tensor(cmd.outputs[i], std::move(tensor_outputs[i]));
                }
            } else if constexpr (std::is_same_v<T, Del>) {
                free(cmd.dest);
            } else if constexpr (std::is_same_v<T, GetValue>) {
557
                mgb_assert(cmd.dest->ptr, "Invalid tensor ptr!");
558 559 560 561 562 563
                cmd.dest->ptr->fetch_value();
                MGB_LOCK_GUARD(m_mutex);
                cmd.dest->value_fetched = true;
                if (m_waitee == cmd.dest) {
                    m_cv.notify_all();
                }
564
            } else if constexpr (std::is_same_v<T, SwapIn>) {
565
                produce_tensor(cmd.dest, Tensor::make(cmd.dest->h_value));
566
            } else if constexpr (std::is_same_v<T, SwapOut>) {
567
                cmd.dest->h_value = cmd.dest->ptr->get_value();
568
                release_tensor(cmd.dest);
569
            } else if constexpr (std::is_same_v<T, Drop>) {
570
                release_tensor(cmd.dest);
571 572 573 574 575 576 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
            } else if constexpr (std::is_same_v<T, SetOption>) {
                m_worker_state.options.set_option(cmd.key, cmd.value);
            } else if constexpr (std::is_same_v<T, StartProfile>) {
                CompNode::sync_all();
                m_worker_state.profiler.reset(cmd.profiler);
            } else if constexpr (std::is_same_v<T, StopProfile>) {
                for (auto&& [device, scopes]: m_worker_state.device_scope_map) {
                    MGB_MARK_USED_VAR(scopes);
                    sync_device_scope(device);
                }
                do_finish_command();
                auto profiler = std::make_unique<InterpreterProfiler>();
                std::swap(profiler, m_worker_state.profiler);
                auto records = profiler->stop();
                auto host_map = [this](std::thread::id tid) {
                    if (tid == m_channel_state.tid) {
                        return "channel";
                    } else if (tid == m_worker_state.tid) {
                        return "worker";
                    } else {
                        return "unknown";
                    }
                };
                InterpreterProfiler::dump_data(cmd.basename, cmd.format, records, profiler->get_option(), host_map);
            } else if constexpr (std::is_same_v<T, PushScope>) {
                m_worker_state.scopes.push_back(cmd.scope_name);
                do_finish_command();
                m_worker_state.profiler->record_host<WorkerBeginScope>(cmd.scope_name);
            } else if constexpr (std::is_same_v<T, PopScope>) {
                mgb_assert(m_worker_state.scopes.back() == cmd.scope_name, "scope name mismatch");
                m_worker_state.scopes.pop_back();
                do_finish_command();
                m_worker_state.profiler->record_host<WorkerEndScope>(cmd.scope_name);
604
            } else {
605
                static_assert(!std::is_same_v<T, T>);
606
            }
607
    };
608
    std::visit([&](const auto& cmd){
609 610 611 612 613 614 615
        using T = std::decay_t<decltype(cmd)>;
        if (!m_worker_state.options.catch_worker_execption) {
            cmd_visitor(cmd);
            return;
        }
        try {
            cmd_visitor(cmd);
616 617
        } catch (...) {
            MGB_LOCK_GUARD(m_mutex);
618 619 620 621 622 623 624
            if constexpr (std::is_same_v<T, ApplyOp>) {
                for (auto oup : cmd.outputs) {
                    oup->invalid = true;
                }
            } else if constexpr (std::is_same_v<T, Put>) {
                cmd.dest->invalid = true;
            }
625 626 627
            m_worker_exc = std::current_exception();
            m_cv.notify_all();
        }
628 629
    }, icmd.second);
    do_finish_command();
630 631 632 633
}

void ChannelImpl::check_worker_exc_unsafe() {
    if (m_worker_exc) {
634 635
        // for reuse interpreter_for_py after some exception tests
        m_waitee = nullptr;
636 637 638 639 640
        std::exception_ptr exc;
        std::swap(exc, m_worker_exc);
        std::rethrow_exception(exc);
    }
}
641 642 643 644 645

void ChannelImpl::CommandBuffer::enqueue(Command cmd) {
    if (std::get_if<Del>(&cmd) && fuse_del(std::get<Del>(cmd))) {
        return;
    }
646
    mgb_log_debug("%s Enqueued", to_string(cmd).c_str());
647 648 649 650 651
    m_commands.push_back(std::move(cmd));
    auto flush_pos = flush_pos_for(m_commands.back());
    flush(flush_pos);
}

652 653 654 655
void ChannelImpl::CommandBuffer::flush() {
    flush(m_commands.end());
}

656 657
void ChannelImpl::CommandBuffer::flush(Handle pos) {
    for (auto iter = m_commands.begin(); iter != pos; ++iter) {
658 659 660 661
        mgb_log_debug("%s Flushed", to_string(*iter).c_str());
        IdentifiedCommand icmd{++m_owner->m_last_id, std::move(*iter)};
        m_owner->m_channel_state.profiler->record_host<CommandEnqueueEvent>(icmd);
        m_owner->m_worker.add_task(std::move(icmd));
662 663 664 665 666 667 668 669 670 671 672 673 674 675 676 677 678 679 680 681
    }
    m_commands.erase(m_commands.begin(), pos);
}

auto ChannelImpl::CommandBuffer::flush_pos_for(const Command& cmd) -> Handle {
    return std::visit([this](const auto& cmd) {
        using T = std::decay_t<decltype(cmd)>;
        if constexpr (std::is_same_v<T, ApplyOp>) {
            auto* op_type = cmd.op->dyn_typeinfo();
            if (op_type == RemoteRecv::typeinfo() ||
                op_type == RemoteSend::typeinfo() ||
                op_type == CollectiveComm::typeinfo() ||
                op_type == opr::InputCallback::typeinfo() ||
                op_type == opr::OutputCallback::typeinfo() ||
                op_type == BackwardGraph::typeinfo()) {
                return m_commands.end();
            }
        } else if constexpr (std::is_same_v<T, GetValue>) {
            return m_commands.end();
        }
682 683 684
        size_t buffer_length = m_owner->m_channel_state.options.buffer_length;
        if (m_commands.size() > buffer_length) {
            return m_commands.begin() + (m_commands.size() - buffer_length);
685 686 687 688 689 690 691 692 693 694 695 696 697 698 699 700 701 702 703 704 705 706 707
        }
        return m_commands.begin();
    }, cmd);
}

/**
 * 1. Find ApplyOp(dest) in buffered commands
 * 2. Check if there are other usages between ApplyOp and Del, return false if not
 * 3. Fuse Del into ApplyOp, return true
 */
bool ChannelImpl::CommandBuffer::fuse_del(const Del& cmd) {
    auto* dest = cmd.dest;
    // TODO: eliminate Puts
    auto begin = m_commands.begin(), end = m_commands.end();
    auto apply_iter = std::find_if(begin, end, [dest](const Command& cmd){
        if (auto* apply = std::get_if<ApplyOp>(&cmd)) {
            return std::count(apply->inputs.begin(), apply->inputs.end(), dest) > 0;
        }
        return false;
    });
    if (apply_iter == end || find_last_usage(dest, {apply_iter+1, end}) != end) {
        return false;
    }
708
    mgb_log_debug("%s Fused", to_string(Command{cmd}).c_str());
709 710 711 712 713 714 715 716 717 718 719 720 721 722 723 724 725 726 727 728 729 730 731 732 733 734 735 736 737 738 739 740 741 742 743 744 745 746 747 748 749 750 751 752 753 754
    std::get<ApplyOp>(*apply_iter).dels.push_back(dest);
    return true;
}

auto ChannelImpl::CommandBuffer::find_last_usage(TensorInfo* dest, Range range)
        -> Handle {
    auto found = range[1];
    for (auto iter = range[0]; iter != range[1]; ++iter) {
        std::visit([&](const auto& cmd) {
            using T = std::decay_t<decltype(cmd)>;
            if constexpr (std::is_same_v<T, ApplyOp>) {
                if (std::count(cmd.inputs.begin(), cmd.inputs.end(),
                               dest) > 0) {
                    found = iter;
                }
            } else if constexpr (std::is_same_v<T, GetValue>) {
                if (cmd.dest == dest) {
                    found = iter;
                }
            } else if constexpr (std::is_same_v<T, SwapIn> ||
                    std::is_same_v<T, SwapOut> ||
                    std::is_same_v<T, Drop>) {
                //TODO: ignore swap-like commands, just remove them from buffer
                if (cmd.dest == dest) {
                    found = iter;
                }
            }
        }, *iter);
    };
    return found;
}

auto ChannelImpl::CommandBuffer::find_produce(TensorInfo* dest, Range range)
        -> Handle {
    return std::find_if(range[0], range[1], [dest](auto& cmd) {
        return std::visit([dest](const auto& cmd){
            using T = std::decay_t<decltype(cmd)>;
            if constexpr (std::is_same_v<T, ApplyOp>) {
                return std::count(cmd.outputs.begin(), cmd.outputs.end(), dest) > 0;
            } else if constexpr (std::is_same_v<T, Put>) {
                return cmd.dest == dest;
            }
            return false;
        }, cmd);
    });
}
755 756 757 758 759 760 761 762 763 764 765 766 767 768 769 770 771 772 773 774 775 776 777 778 779 780 781 782 783 784 785 786 787 788 789 790 791 792

void ChannelImpl::start_profile(std::unordered_map<std::string, int> option) {
    auto profiler_option = InterpreterProfiler::Option::from_dict(option);
    auto profiler = std::make_unique<InterpreterProfiler>();
    profiler->set_option(profiler_option);
    profiler->start(InterpreterProfiler::topic_to_mask(profiler_option.topic));
    std::swap(profiler, m_channel_state.profiler);
    m_buffer.enqueue(StartProfile{m_channel_state.profiler.get()});
}

void ChannelImpl::stop_profile(std::string basename, std::string format) {
    m_buffer.flush();
    auto profiler = std::make_unique<InterpreterProfiler>();
    std::swap(profiler, m_channel_state.profiler);
    profiler.release();
    m_buffer.enqueue(StopProfile{basename, format});
}

void ChannelImpl::push_scope(std::string name) {
    m_channel_state.profiler->record_host<ChannelBeginScope>(name);
    m_channel_state.scopes.push_back(name);
    m_buffer.enqueue(PushScope{name});
}

void ChannelImpl::pop_scope(std::string name) {
    mgb_assert((!m_channel_state.scopes.empty()) && m_channel_state.scopes.back() == name, "scope name mismatch");
    m_channel_state.scopes.pop_back();
    m_channel_state.profiler->record_host<ChannelEndScope>(name);
    m_buffer.enqueue(PopScope{name});
}

void ChannelImpl::assert_in_channel() {
    mgb_assert(m_channel_state.tid != std::this_thread::get_id());
}

void ChannelImpl::assert_in_worker() {
    mgb_assert(m_worker_state.tid == std::this_thread::get_id());
}