tensor.cpp 41.4 KB
Newer Older
1 2 3 4
/**
 * \file imperative/python/src/tensor.cpp
 * 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
 *
 * 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 "megbrain/common.h"
13
#include "megbrain/dtype.h"
14
#include "megbrain/imperative/ops/autogen.h"
15 16
#include "megbrain/imperative/ops/backward_graph.h"
#include "megbrain/imperative/ops/utility.h"
17
#include "megbrain/imperative/profiler.h"
18
#include "megbrain/imperative/transformations/dim_expansion.h"
19
#include "megbrain/imperative/transformations/dtype_promote.h"
20 21 22 23 24 25
#include "megbrain/imperative/transformations/eval.h"
#include "megbrain/imperative/transformations/lazy.h"
#include "megbrain/imperative/transformations/scalar.h"
#include "megbrain/imperative/transformations/symbol.h"
#include "megbrain/imperative/transformations/trace.h"
#include "megbrain/imperative/utils/map.h"
26
#include "megbrain/imperative/utils/stats.h"
27
#include "megbrain/opr/io.h"
28
#include "megbrain/plugin/profiler.h"
29

30
#include "./common.h"
31
#include "./grad.h"
32
#include "./graph_rt.h"
33
#include "./helper.h"
34 35 36
#include "./module_trace.h"
#include "./numpy_dtypes.h"
#include "./tensor.h"
37
#include "./tensor_utils.h"
38
#include "./transformation.h"
39

40
#include <object.h>
41 42
#include <pybind11/numpy.h>
#include <pybind11/operators.h>
43 44
#include <pybind11/pytypes.h>
#include <pyerrors.h>
45
#include <range/v3/all.hpp>
46
#include <string>
47 48 49

#include <unordered_map>

50 51
#include "../../src/impl/mgb_cg_impl.h"

52
namespace py = pybind11;
53
namespace views = ranges::views;
54 55 56

namespace mgb::imperative::python {

57 58
namespace {
WeakKeyMap<ValueWeakRef, py::object> module_trace_info_map;
59 60 61

struct SymbolVarContext {
    TransformationContext context;
62 63
    std::shared_ptr<SymbolTransformation> symbol_tsf;
    std::shared_ptr<ScalarTransformation> scalar_tsf;
64
    std::shared_ptr<DTypePromoteTransformation> dtype_promote_tsf;
65
    std::shared_ptr<DimExpansionTransformation> dim_expansion_tsf;
66

67 68 69
    SymbolVarContext(cg::ComputingGraph* graph) {
        symbol_tsf = std::make_shared<SymbolTransformation>(graph);
        scalar_tsf = std::make_shared<ScalarTransformation>();
70
        dtype_promote_tsf = std::make_shared<DTypePromoteTransformation>();
71
        dim_expansion_tsf = std::make_shared<DimExpansionTransformation>();
72 73 74 75
        Transformation::swap_context(context);
    }

    void init() {
76 77
        symbol_tsf->register_at(Transformation::top());
        scalar_tsf->register_at(Transformation::top());
78
        dtype_promote_tsf->register_at(Transformation::top());
79
        dim_expansion_tsf->register_at(Transformation::top());
80 81
    }

82 83 84 85 86 87 88 89
    ValueRef symvar2val(py::handle py_symbol_var) {
        auto* symbol_var = py_symbol_var.cast<PySymbolVar*>();
        ValueRef value = symbol_tsf->value_type().make(symbol_var->m_node);
        if (symbol_var->is_scalar) {
            value = scalar_tsf->value_type().make(value);
        }
        return value;
    }
90

91 92 93 94 95 96 97 98 99 100 101
    py::object val2symvar(py::handle typeobj, ValueRef value) {
        bool is_scalar = false;
        if (auto* scalar_value = value.as(scalar_tsf->value_type())) {
            value = scalar_value->value();
            is_scalar = true;
        }
        auto* node = value.cast(symbol_tsf->value_type()).node();
        auto py_symbol_var =
                typeobj(pybind11::cast(node, pybind11::return_value_policy::automatic));
        py_symbol_var.cast<PySymbolVar*>()->is_scalar = is_scalar;
        return py_symbol_var;
102 103
    }

104 105
    ~SymbolVarContext() { Transformation::swap_context(context); }
};
106

107 108
}  // namespace

109 110
interpreter::Interpreter::Channel* interpreter_for_py = nullptr;
PyTypeObject* py_tensor_type = nullptr;
111
PyObject* cpp_use_symbolic_shape;
112 113 114 115 116 117 118

#define REGISTE_APPLY_FUNC(mode) \
    void set_##mode(py::object pyf) { mode = pyf.ptr(); }

REGISTE_APPLY_FUNC(cpp_use_symbolic_shape)

#undef REGISTE_APPLY_FUNC
119

120 121 122
PyArray_Descr* _dtype_promotion(PyObject* const* args, size_t nargs);
CompNode _get_device(PyObject* const* args, size_t nargs);

123 124
PyObject* py_apply(
        PyObject* self, PyObject* const* args, size_t nargs /* , PyObject* kwnames */) {
125 126 127 128 129
    try {
        // if (kwnames && PyTuple_GET_SIZE(kwnames)) {
        //     PyErr_SetString(PyExc_TypeError, "keyword argument not allowed");
        //     return nullptr;
        // }
130
        if (nargs < 2) {
131 132 133 134
            PyErr_SetString(
                    PyExc_TypeError,
                    "py_apply expects one Op and at least one tensor "
                    "as argument");
135 136
            return nullptr;
        }
137

138
        auto* py_op = args[0];
139

140 141 142
        ++args;
        --nargs;

143
        auto op = py::handle(py_op).cast<std::shared_ptr<OpDef>>();
144
        SmallVector<ValueRef, 8> tensors(nargs);
145

146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183
        SmallVector<bool, 8> is_symbol_var(nargs, false);
        ComputingGraph* cg = nullptr;
        for (size_t i = 0; i < nargs; ++i) {
            if ((!TensorWrapper::try_cast(args[i])) &&
                py::isinstance<PySymbolVar>(py::handle(args[i]))) {
                is_symbol_var[i] = true;
                ComputingGraph* cur_cg =
                        py::handle(args[i]).cast<PySymbolVar*>()->m_node->owner_graph();
                if (cg == nullptr) {
                    cg = cur_cg;
                } else {
                    mgb_assert(cg == cur_cg);
                }
            }
        }

        mgb::CompNode target_cn;
        mgb::DType target_dtype;

        auto convert_pyinput_to_tensor = [&](size_t i) -> ValueRef {
            if (!target_dtype.valid()) {
                target_dtype = npy::dtype_np2mgb_descr(_dtype_promotion(args, nargs));
                target_cn = _get_device(args, nargs);
            }
            HostTensorND ht(target_cn);
            ht = npy::np2tensor(args[i], npy::Meth::copy_into(&ht), target_dtype);
            if (PyArray_Check(args[i])) {  // non scaler
                return imperative::apply(
                        CreateTensor(CreateTensor::Const, target_cn, ht.layout()),
                        HostStorage::make(ht.storage()))[0];
            } else {  // scaler
                return imperative::apply(
                        CreateTensor(CreateTensor::Const, target_cn, target_dtype, {}),
                        HostStorage::make(ht.storage()))[0];
            }
        };

        if (cg != nullptr) {
184
            // swap to a special context to reuse scalar handle
185 186
            size_t symbol_var_idx = 8;
            SymbolVarContext context(cg);
187
            context.init();
188
            for (size_t i = 0; i < nargs; ++i) {
189 190 191 192 193 194
                if (is_symbol_var[i]) {
                    symbol_var_idx = i;
                    tensors[i] = context.symvar2val(args[i]);
                } else {
                    tensors[i] = convert_pyinput_to_tensor(i);
                }
195
            }
196
            auto outputs = imperative::apply(*op, tensors);
197
            auto ret = pybind11::tuple(outputs.size());
198
            auto typeobj = py::handle(args[symbol_var_idx]).get_type();
199
            for (size_t i = 0; i < outputs.size(); ++i) {
200
                ret[i] = context.val2symvar(typeobj, outputs[i]);
201 202 203
            }
            return ret.release().ptr();
        }
204 205

        for (size_t i = 0; i < nargs; ++i) {
206
            if (TensorWrapper* tw = TensorWrapper::try_cast(args[i])) {
207
                tensors[i] = tw->m_tensor->data();
208
            } else {
209
                tensors[i] = convert_pyinput_to_tensor(i);
210 211 212
            }
        }

213
        auto outputs = [&] { return imperative::apply(*op, tensors); }();
214 215 216
        size_t nout = outputs.size();
        auto ret = py::tuple(nout);
        for (size_t i = 0; i < nout; ++i) {
217
            ret[i] = TensorWrapper::make(py_tensor_type, std::move(outputs[i]));
218 219
        }
        return ret.release().ptr();
220 221
    }
    PYEXT17_TRANSLATE_EXC_RET(nullptr)
222 223 224 225 226 227 228 229 230 231 232
}

TensorWrapper::TensorWrapper(PyObject* args, PyObject* kwargs) {
    if (kwargs && PyDict_Size(kwargs)) {
        throw py::type_error("keyword argument not allowed");
    }
    auto nargs = PyTuple_Size(args);
    auto tup = py::reinterpret_borrow<py::tuple>(args);
    if (nargs == 0) {
        throw py::type_error("too few arguments");
    }
233
    if (auto* t = try_cast(tup[0].ptr())) {
234 235 236
        if (nargs > 1) {
            throw py::type_error("expect 1 argument");
        }
237
        m_tensor = t->m_tensor->copy();
238
    } else {
239 240
        if (nargs == 1) {
            auto arg0 = PyTuple_GetItem(args, 0);
241 242 243 244 245 246
            // for DeviceTensorND
            if (strstr(arg0->ob_type->tp_name, "DeviceTensorND")) {
                auto dv = py::handle(arg0).cast<DeviceTensorND>();
                m_tensor = std::make_shared<Tensor>(imperative::apply(
                        CreateTensor(CreateTensor::Common, dv.comp_node(), dv.layout()),
                        DeviceStorage::make(dv.storage()))[0]);
247
            } else {
248 249
                throw py::type_error(
                        "single argument is not tensor, varnode or devicetensor");
250
            }
251
        } else {
252
            py::detail::loader_life_support life_sup;  // FIXME!!!required to cast DType
253 254
            if (nargs != 5 && nargs != 6) {
                throw py::type_error("expect 5 or 6 arguments");
255
            }
256 257 258 259
            auto data = tup[0].cast<py::array>();
            DType dtype = tup[1].cast<DType>();
            CompNode cn = tup[2].cast<CompNode>();
            bool is_const = tup[3].cast<bool>();
260
            bool no_cache = nargs == 6 ? tup[4].cast<bool>() : false;
261
            std::string name;
262 263
            if (tup[nargs - 1].ptr() != Py_None)
                name = tup[nargs - 1].cast<std::string>();
264 265

            // const op
266
            {
267 268 269
                CreateTensor::Kind kind = is_const ? CreateTensor::Const
                                        : no_cache ? CreateTensor::Unique
                                                   : CreateTensor::Common;
270
                HostTensorND ret(cn);
271 272 273 274 275 276 277 278 279 280 281
                ret = npy::np2tensor(data.ptr(), npy::Meth::copy_into(&ret), dtype);
                mgb_assert(
                        ret.layout().is_empty() || ret.layout().is_contiguous(),
                        "host value should be continuous");
                ValueShape shape;
                for (size_t i = 0; i < data.ndim(); ++i) {
                    shape[shape.ndim++] = data.shape(i);
                }
                m_tensor = std::make_shared<Tensor>(imperative::apply(
                        CreateTensor(kind, cn, ret.dtype(), shape),
                        HostStorage::make(ret.storage()))[0]);
282 283
            }

284 285 286
            if (!name.empty()) {
                m_tensor->reset(
                        imperative::apply(RenameValue(name), m_tensor->data())[0]);
287
            }
288 289
        }
    }
290
    mgb_assert(m_tensor->data());
291 292
}

293
PyObject* TensorWrapper::module_trace_info() {
294
    if (auto module_trace_info = module_trace_info_map.try_get(m_tensor->data())) {
295 296 297
        if (module_trace_info->ptr()) {
            return module_trace_info->inc_ref().ptr();
        }
298
    }
299 300 301 302 303
    PyErr_SetString(
            PyExc_AttributeError,
            "Has no attribute named \'_NodeMixin__node\', please "
            "set it first");
    return nullptr;
304 305 306
}

void TensorWrapper::set_module_trace_info(PyObject* obj) {
307
    // TODO: erase when obj == nullptr
308
    module_trace_info_map[m_tensor->data()] = py::reinterpret_borrow<py::object>(obj);
309 310
}

311 312 313 314 315
void TensorWrapper::_set_name(PyObject* dest) {
    auto py_dest = py::reinterpret_borrow<py::object>(dest);
    auto name = py_dest.cast<std::string>();
    m_tensor->set_name(name);
}
316

317 318
PyObject* TensorWrapper::_detail() {
    return py::str(m_tensor->data().unwrap().to_string()).release().ptr();
319 320
}

321 322
void TensorWrapper::_watch() {
    m_tensor->data().watch();
323 324
}

325
PyObject* TensorWrapper::shape() {
326
    auto shape = m_tensor->shape();
327

328
    if (!shape) {
329 330
        Py_RETURN_NONE;
    }
331 332 333
    py::tuple ret(shape->ndim);
    for (size_t i = 0; i < shape->ndim; ++i) {
        ret[i] = shape->at(i);
334 335 336 337 338 339 340 341 342 343 344 345 346
    }
    return ret.release().ptr();
}

PyObject* TensorWrapper::dtype() {
    return py::cast(m_tensor->dtype()).release().ptr();
}

PyObject* TensorWrapper::device() {
    return py::cast(m_tensor->comp_node()).release().ptr();
}

PyObject* TensorWrapper::numpy() {
347
    auto hv = m_tensor->numpy();
348
    if (!hv) {
349 350 351
        PyErr_SetString(PyExc_ValueError, "tensor invalid");
        return nullptr;
    }
352 353
    auto arr = py::reinterpret_steal<py::array>(
            npy::ndarray_from_tensor(hv->as_nd(true), npy::ShareType::TRY_SHARE));
354
    if (hv->shape().is_scalar()) {
355 356 357 358 359 360 361
        mgb_assert(PyArray_Check(arr.ptr()));
        return PyArray_Squeeze(reinterpret_cast<PyArrayObject*>(arr.ptr()));
    }
    return arr.release().ptr();
}

void TensorWrapper::reset(PyObject* tensor) {
362
    TensorWrapper* t = TensorWrapper::try_cast(tensor);
363 364 365
    if (!t) {
        throw py::type_error("expect Tensor");
    }
366
    m_tensor->reset(t->m_tensor->data());
367 368
}

369
PyObject* TensorWrapper::detach() {
370 371
    auto detached = imperative::apply(DetachGrad(), m_tensor->data())[0];
    return TensorWrapper::make(py_tensor_type, detached).release().ptr();
372 373
}

374
PyObject* TensorWrapper::_dev_tensor() {
375 376 377
    auto dv = m_tensor->data().dev_tensor();
    // TODO: handle scalar
    return py::cast(dv->as_nd(true)).release().ptr();
378 379 380
}

void TensorWrapper::_drop() {
381
    imperative::apply(DTRCommand(DTRCommand::Drop), m_tensor->data());
382 383
}

384
PyObject* TensorWrapper::isscalar() {
385
    if (m_tensor->is_scalar()) {
386 387 388 389 390 391 392 393 394 395 396 397 398
        Py_RETURN_TRUE;
    } else {
        Py_RETURN_FALSE;
    }
}

struct TensorWeakRef {
    std::weak_ptr<Tensor> wptr;

    TensorWeakRef(const TensorWrapper& tw) : wptr(tw.m_tensor) {}

    py::object operator()() {
        if (auto p = wptr.lock()) {
399
            return TensorWrapper::make(py_tensor_type, p);
400 401 402
        }
        return py::none();
    }
403
    int _use_cnt() { return wptr.use_count(); }
404 405
};

406 407 408 409 410 411 412 413 414 415 416 417 418
#ifdef METH_FASTCALL
#define MGE_PY_INTERFACE(NAME, FUNC) \
    { #NAME, (PyCFunction)FUNC, METH_FASTCALL, nullptr }
#else
#define WRAP_FUNC_PY35(FUNC)                                \
    PyObject* py35_##FUNC(PyObject* self, PyObject* args) { \
        auto* arr = &PyTuple_GET_ITEM(args, 0);             \
        auto size = PyTuple_GET_SIZE(args);                 \
        return FUNC(self, arr, size);                       \
    }
WRAP_FUNC_PY35(py_apply);
WRAP_FUNC_PY35(dtype_promotion);
WRAP_FUNC_PY35(get_device);
419 420 421
WRAP_FUNC_PY35(make_shape_tuple);
WRAP_FUNC_PY35(getitem_cpp);
WRAP_FUNC_PY35(setitem_cpp);
422
WRAP_FUNC_PY35(split_cpp);
423
WRAP_FUNC_PY35(expand_dims_cpp);
424
WRAP_FUNC_PY35(squeeze_cpp);
425
WRAP_FUNC_PY35(transpose_cpp);
426 427
WRAP_FUNC_PY35(broadcast_cpp);
WRAP_FUNC_PY35(reshape_cpp);
428
WRAP_FUNC_PY35(Const);
429 430 431
WRAP_FUNC_PY35(astype_cpp);
WRAP_FUNC_PY35(convert_single_value_cpp);
WRAP_FUNC_PY35(convert_inputs_cpp);
432
WRAP_FUNC_PY35(astensor1d_cpp);
433 434 435 436 437
#undef WRAP_FUNC_PY35
#define MGE_PY_INTERFACE(NAME, FUNC) \
    { #NAME, (PyCFunction)py35_##FUNC, METH_VARARGS, nullptr }
#endif

438
void init_tensor(py::module m) {
439
    imperative::Tensor::static_initialize();
440 441 442 443 444

    static auto& transformations = TransformationManager::get_instance();

    using Segment = TransformationManager::Segment;

445 446 447 448 449 450
    using Channel = interpreter::Interpreter::Channel;

    auto* channel =
            imperative::ResourceManager::create_global<std::unique_ptr<Channel>>(
                    interpreter::Interpreter::inst().create_channel())
                    ->get();
451 452 453
    interpreter_for_py = channel;
    transformations.register_at<Segment::Eval>(
            std::make_shared<InterpreterTransformation>(
454
                    std::shared_ptr<Channel>(channel, [](Channel*) {})));
455 456
    transformations.register_at<Segment::Scalar>(
            std::make_shared<ScalarTransformation>());
457 458
    transformations.register_at<Segment::DTypePromote>(
            std::make_shared<DTypePromoteTransformation>());
459 460
    transformations.register_at<Segment::DimExpansion>(
            std::make_shared<DimExpansionTransformation>());
461

462 463
    static py::exception<interpreter::AsyncError> py_async_error(
            m, "AsyncError", PyExc_RuntimeError);
464 465
    py::register_exception_translator([](std::exception_ptr p) {
        try {
466 467
            if (p)
                std::rethrow_exception(p);
468 469 470 471 472 473 474 475 476 477
        } catch (const interpreter::AsyncError& e) {
            pyext17::pybind11_translate_exception(e.nested_ptr());
            if (PyErr_Occurred()) {
                PyObject *exc, *val, *tb;
                PyErr_Fetch(&exc, &val, &tb);
                PyErr_NormalizeException(&exc, &val, &tb);
                if (tb) {
                    PyException_SetTraceback(val, tb);
                }
                auto val2 = py_async_error.py::object::operator()(
478 479
                        "An async error is reported. See above for the actual cause."
                        " Hint: This is where it is reported, not where it happened."
480
                        " You may call `megengine.config.async_level = 0 "
481 482 483
                        "to get better error reporting.");
                PyException_SetCause(
                        val2.ptr(), val);  // PyException_SetCause steals reference
484 485
                Py_XDECREF(exc);
                Py_XDECREF(tb);
486 487
                PyErr_Restore(
                        py_async_error.inc_ref().ptr(), val2.release().ptr(), nullptr);
488 489 490 491 492 493
            } else {
                py_async_error("Unkown async error");
            }
        }
    });

494 495 496 497 498 499 500 501 502
    auto* tensor_type =
            TensorWrapper::wrap_t::type()
                    .def<&TensorWrapper::numpy>("numpy")
                    .def_getset<&TensorWrapper::shape>("shape")
                    .def_getset<&TensorWrapper::dtype>("dtype")
                    .def_getset<&TensorWrapper::device>("device")
                    .def<&TensorWrapper::reset>("_reset")
                    .def<&TensorWrapper::isscalar>("_isscalar")
                    .def<&TensorWrapper::detach>("detach")
503
                    // TODO: remove this
504 505 506
                    .def<&TensorWrapper::_dev_tensor>("_dev_tensor")
                    .def<&TensorWrapper::_drop>("_drop")
                    .def<&TensorWrapper::_use_cnt>("_use_cnt")
507 508 509
                    .def<&TensorWrapper::_detail>("_detail")
                    .def<&TensorWrapper::_set_name>("_set_name")
                    .def<&TensorWrapper::_watch>("_watch")
510 511 512 513 514 515
                    .def_getset<
                            &TensorWrapper::module_trace_info,
                            &TensorWrapper::set_module_trace_info>("_NodeMixin__node")
                    .finalize();
    if (!tensor_type)
        throw py::error_already_set();
516 517 518
    py::setattr(m, "Tensor", tensor_type);

    py::class_<TensorWeakRef>(m, "TensorWeakRef")
519 520 521
            .def(py::init<const TensorWrapper&>())
            .def("__call__", &TensorWeakRef::operator())
            .def("_use_cnt", &TensorWeakRef::_use_cnt);
522

523 524 525
    py::class_<PySymbolVar, std::shared_ptr<PySymbolVar>>(m, "SymbolVar")
            .def_property_readonly(
                    "dtype", [](PySymbolVar* v) { return v->m_node->dtype(); })
526 527 528
            .def_property(
                    "var", [](PySymbolVar* v) { return v->m_node; },
                    [](PySymbolVar* s, cg::VarNode* v) { s->m_node = v; })
529
            .def_property_readonly(
530
                    "device", [](PySymbolVar* v) { return v->m_node->comp_node(); })
531
            .def_property_readonly(
532
                    "graph", [](PySymbolVar* v) { return v->m_node->owner_graph(); })
533 534 535
            .def_property_readonly(
                    "shape",
                    [](PySymbolVar* v) -> const TensorShape* {
536
                        auto&& mgr = v->m_node->owner_graph()->static_infer_manager();
537 538
                        return mgr.infer_shape_fallible(v->m_node);
                    })
539 540 541 542 543 544 545 546 547 548 549 550 551 552 553
            .def("numpy",
                 [](PySymbolVar* v) {
                     auto&& mgr = v->m_node->owner_graph()->static_infer_manager();
                     auto&& type = mgr.get_infer_type(v->m_node);
                     using InferType = cg::static_infer::InferType;
                     if (!(type.value & (InferType::CONST | InferType::RT_STATIC))) {
                         throw py::value_error("value invalid!");
                     }
                     auto* val = mgr.infer_value_fallible(v->m_node);
                     if (!val) {
                         throw py::value_error("value invalid!");
                     }
                     auto np_val = py::cast(*val).attr("numpy")();
                     return np_val;
                 })
554 555 556 557 558 559
            .def("_isscalar", [](PySymbolVar* v) { return v->is_scalar; })
            .def(py::init([](cg::VarNode* node) {
                     return std::make_shared<PySymbolVar>(node);
                 }),
                 py::arg() = nullptr);

560
    static PyMethodDef method_defs[] = {
561 562 563
            MGE_PY_INTERFACE(apply, py_apply),
            MGE_PY_INTERFACE(dtype_promotion, dtype_promotion),
            MGE_PY_INTERFACE(get_device, get_device),
564 565 566
            MGE_PY_INTERFACE(make_shape_tuple, make_shape_tuple),
            MGE_PY_INTERFACE(getitem_cpp, getitem_cpp),
            MGE_PY_INTERFACE(setitem_cpp, setitem_cpp),
567
            MGE_PY_INTERFACE(split_cpp, split_cpp),
568
            MGE_PY_INTERFACE(expand_dims_cpp, expand_dims_cpp),
569
            MGE_PY_INTERFACE(squeeze_cpp, squeeze_cpp),
570
            MGE_PY_INTERFACE(transpose_cpp, transpose_cpp),
571 572
            MGE_PY_INTERFACE(broadcast_cpp, broadcast_cpp),
            MGE_PY_INTERFACE(reshape_cpp, reshape_cpp),
573
            MGE_PY_INTERFACE(Const, Const),
574 575 576
            MGE_PY_INTERFACE(astype_cpp, astype_cpp),
            MGE_PY_INTERFACE(convert_single_value_cpp, convert_single_value_cpp),
            MGE_PY_INTERFACE(convert_inputs_cpp, convert_inputs_cpp),
577
            MGE_PY_INTERFACE(astensor1d_cpp, astensor1d_cpp),
578
            {nullptr, nullptr, 0, nullptr}};
579
    for (auto&& def : method_defs) {
580 581
        if (def.ml_meth != nullptr) {
            auto* func = PyCFunction_NewEx(&def, nullptr, nullptr);
582 583
            if (!func)
                throw py::error_already_set();
584 585 586
            py::setattr(m, def.ml_name, func);
        }
    }
587

588 589 590 591
    static constexpr auto sync_py_task_q = [] {
        py::gil_scoped_release _;
        py_task_q.wait_all_task_finish();
    };
592

593
    m.def("clear_candidates", [channel]() { channel->clear_candidates(); });
594 595
    m.def("set_option", [channel](std::string name, size_t value) {
        channel->set_option(name, value);
596
    });
597
    m.def("get_option",
598 599 600 601 602 603 604 605 606 607 608 609 610 611 612 613 614 615 616 617 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
          [channel](std::string name) { return channel->get_option(name); });
    m.def("push_scope", [channel](std::string name) {
        Transformation::push_scope(name);
        channel->push_scope(name);
    });
    m.def("pop_scope", [channel](std::string name) {
        channel->pop_scope(name);
        Transformation::pop_scope(name);
    });
    m.def("start_profile", [channel](imperative::Profiler::options_t options) {
        channel->sync();
        imperative::Profiler::load_options(std::move(options));
        imperative::Profiler::start_profile();
        channel->start_profile();
    });
    m.def("stop_profile", [channel]() -> std::function<void(std::string, std::string)> {
        channel->stop_profile();
        channel->sync();
        imperative::Profiler::stop_profile();
        auto results = std::make_shared<imperative::Profiler::bundle_t>(
                imperative::Profiler::collect());
        return [results = results](std::string basename, std::string format) mutable {
            imperative::Profiler::dump_profile(basename, format, std::move(*results));
            results = nullptr;
        };
    });
    m.def("sync", [channel]() {
        if (channel->check_available()) {
            channel->sync();
        }
        sync_py_task_q();
    });
    m.def("full_sync", [channel]() {
        if (channel->check_available()) {
            channel->sync();
        }
        CompNode::sync_all();
        CompNode::foreach ([](CompNode cn) {
            auto err = cn.check_async_error();
            mgb_assert(!err, "%s", err->what());
        });
        sync_py_task_q();
    });
    m.def("close", [channel]() {
        channel->close();
        sync_py_task_q();
644 645 646 647 648 649 650 651
    });

    py::handle grad_key_type =
            GradKeyWrapper::wrap_t::type()
                    .def<&GradKeyWrapper::attach>("attach")
                    .def<&GradKeyWrapper::is_attached_to>("is_attached_to")
                    .def_getset<&GradKeyWrapper::get_name, &GradKeyWrapper::set_name>(
                            "name")
652 653 654 655
                    .def<&GradKeyWrapper::enter>("enter")
                    .def<&GradKeyWrapper::exit>("exit")
                    .def<&GradKeyWrapper::suppress>("suppress")
                    .def<&GradKeyWrapper::resume>("resume")
656 657 658
                    .finalize();
    if (!grad_key_type)
        throw py::error_already_set();
659
    py::setattr(m, "GradKey", grad_key_type);
660
    m.def("backward", &GradKeyWrapper::backward);
661
    m.def("get_backward_closure", &GradKeyWrapper::get_backward_closure);
662

663 664 665 666 667 668 669 670 671 672 673 674 675 676 677 678 679 680 681 682 683 684 685
    m.def("set_py_tensor_type", [](py::object type_obj) {
        py_tensor_type = reinterpret_cast<PyTypeObject*>(type_obj.inc_ref().ptr());
    });

    /**
     * \brief trace proxy
     *
     */
    struct Trace {
        bool symbolic = false;
        bool no_exec = false;
        bool capture_as_const = false;
        bool profile = false;
        bool record_input_shapes = false;
        py::function options_visitor;
        std::shared_ptr<TracingTransformation> tracing;
        std::shared_ptr<CompiledTransformation> compiled;
        std::shared_ptr<LazyEvalTransformation> lazy_eval;
        std::pair<size_t, std::shared_ptr<GraphProfiler>> profiler;
        std::optional<TraceResult> trace_result;
        std::function<bool(py::object, py::object)> array_comparator;

        bool compare_value(ValueRef lhs, ValueRef rhs) {
686 687
            auto lvalue = lhs.cast_ref<HostValue>();
            auto rvalue = rhs.cast_ref<HostValue>();
688
            if (lvalue->shape() != rvalue->shape()) {
689 690
                return false;
            }
691
            if (lvalue->shape().total_nr_elems() == 1) {
692 693 694 695
                return lvalue->item() == rvalue->item();
            }
            HostTensorND lnd = lvalue->as_nd(true);
            HostTensorND rnd = rvalue->as_nd(true);
696
            auto larr = py::reinterpret_steal<py::array>(
697
                    npy::ndarray_from_tensor(lnd, npy::ShareType::TRY_SHARE));
698
            auto rarr = py::reinterpret_steal<py::array>(
699
                    npy::ndarray_from_tensor(rnd, npy::ShareType::TRY_SHARE));
700 701 702 703 704 705 706 707 708 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 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 793 794 795 796 797 798 799 800 801 802 803 804 805 806 807 808 809 810 811 812 813 814 815 816 817 818 819 820 821 822 823 824 825 826 827 828 829 830 831 832 833 834 835
            return array_comparator(larr, rarr);
        }

        void enter() {
            auto& self = *this;
            if (!self.trace_result) {  // untraced
                self.tracing = std::make_shared<TracingTransformation>(
                        self.capture_as_const, self.record_input_shapes);
                if (self.symbolic) {
                    self.lazy_eval =
                            std::make_shared<LazyEvalTransformation>(self.no_exec);
                    self.options_visitor(py::cast(&self.lazy_eval->options()));
                }
            } else if (!self.compiled) {  // traced but not compiled
                using namespace std::placeholders;
                self.compiled = std::make_shared<CompiledTransformation>(
                        *self.trace_result, self.record_input_shapes);
                self.compiled->set_value_comparator(
                        std::bind(&Trace::compare_value, this, _1, _2));
                self.options_visitor(py::cast(&self.compiled->options()));
                self.compiled->compile();
            }
            // register transformations
            if (self.compiled) {
                if (self.profile) {
                    auto& current_graph = self.compiled->graph();
                    if (self.profiler.first != self.compiled->graph().id()) {
                        // graph changed
                        self.profiler = std::make_pair(
                                current_graph.id(),
                                std::make_shared<GraphProfiler>(&current_graph));
                    }
                }
                transformations.register_at<Segment::Trace>(self.compiled);
                // start execute because InputCallback depends
                self.compiled->execute();
            } else if (self.tracing) {
                transformations.register_at<Segment::Trace>(self.tracing);
                if (self.lazy_eval) {
                    transformations.register_at<Segment::Eval>(self.lazy_eval);
                }
            } else {
                mgb_throw(MegBrainError, "invalid state: neither tracing nor compiled");
            }
        }

        void exit() {
            auto& self = *this;
            if (self.tracing) {
                transformations.unregister<Segment::Trace>(self.tracing);
                self.trace_result = self.tracing->get_result();
                self.tracing.reset();
                if (self.lazy_eval) {
                    auto lazy_eval = std::move(self.lazy_eval);
                    transformations.unregister<Segment::Eval>(lazy_eval);
                    lazy_eval->check_exception();
                }
            } else if (self.compiled) {
                transformations.unregister<Segment::Trace>(self.compiled);
                self.compiled->wait();
            } else {
                mgb_throw(MegBrainError, "invalid state: neither tracing nor compiled");
            }
        }

        VarNodeArray dump(
                std::shared_ptr<ComputingGraph> graph,
                std::vector<std::tuple<std::string, std::string, TensorShape>> inputs,
                std::vector<std::pair<std::string, std::string>> outputs,
                bool prefer_input_names) {
            auto& self = *this;
            mgb_assert(self.trace_result);
            // mark is like "arg_0", "kwarg_xxx", "output_0" ...
            std::unordered_map<std::string, size_t> mark2var;
            for (size_t i = 0; i < self.trace_result->vars.size(); ++i) {
                auto& name = self.trace_result->vars[i].mark;
                if (!name.empty()) {
                    mark2var[name] = i;
                }
            }
            std::vector<std::tuple<size_t, std::string, TensorShape>> input_vars;
            std::vector<std::pair<size_t, std::string>> output_vars;
            for (auto&& [input_mark, input_name, input_shape] : inputs) {
                mgb_assert(input_shape.ndim, "input shape invalid");
                input_vars.push_back(
                        {mark2var.at(input_mark), input_name, input_shape});
            }
            for (auto&& [output_name, repr] : outputs) {
                output_vars.push_back({mark2var.at(output_name), repr});
            }
            self.options_visitor(py::cast(&graph->options()));
            auto vars = self.trace_result->dump(
                    *graph, input_vars, output_vars, prefer_input_names);
            return vars;
        }
    };

    py::class_<Trace>(m, "Trace")
            .def(py::init<>())
            .def_readwrite("record_input_shapes", &Trace::record_input_shapes)
            .def_readwrite("array_comparator", &Trace::array_comparator)
            .def_readwrite("profile", &Trace::profile)
            .def_property_readonly(
                    "options",
                    [](Trace& self) {
                        if (self.compiled) {
                            return &self.compiled->options();
                        } else {
                            return (ComputingGraph::Options*)nullptr;
                        }
                    })
            .def("get_profile",
                 [](Trace& self) -> py::object {
                     if (self.profiler.second && self.compiled) {
                         auto json = self.profiler.second->to_json_full(
                                 self.compiled->graph().current_comp_seq());
                         return py::str(json->to_string());
                     } else {
                         return py::none();
                     }
                 })
            .def_readwrite("symbolic", &Trace::symbolic)
            .def_readwrite("capture_as_const", &Trace::capture_as_const)
            .def_readwrite("no_exec", &Trace::no_exec)
            .def_readwrite("options_visitor", &Trace::options_visitor)
            .def("enter", &Trace::enter)
            .def("exit", &Trace::exit)
            .def("dump", &Trace::dump)
            .def("begin_excluded_region",
                 [](Trace& self) {
                     mgb_assert(bool(self.tracing) ^ bool(self.compiled));
                     if (self.tracing) {
                         transformations.unregister<Segment::Trace>(self.tracing);
                     } else if (self.compiled) {
                         transformations.unregister<Segment::Trace>(self.compiled);
                     }
836
                 })
837 838 839 840 841 842 843 844 845
            .def("end_excluded_region", [](Trace& self) {
                mgb_assert(bool(self.tracing) ^ bool(self.compiled));
                if (self.tracing) {
                    transformations.register_at<Segment::Trace>(self.tracing);
                } else if (self.compiled) {
                    transformations.register_at<Segment::Trace>(self.compiled);
                }
            });

846 847 848 849 850 851 852 853
    m.def("reduce_to_scalar", [](py::object op, py::object tensor) -> py::object {
        auto reduce_to_scalar = [](const OpDef& op, const ValueRef& input) {
            auto make_scalar_shape = [&](CompNode device) {
                return imperative::apply(
                        CreateTensor(CreateTensor::Const, device, dtype::Int32(), {0}),
                        HostStorage::make(device))[0];
            };
            return imperative::apply(op, input, make_scalar_shape(*input.device()))[0];
854
        };
855 856 857 858 859
        if (py::isinstance<PySymbolVar>(tensor)) {
            auto* graph = tensor.cast<PySymbolVar*>()->m_node->owner_graph();
            SymbolVarContext context(graph);
            context.init();
            auto output = reduce_to_scalar(
860
                    *op.cast<std::shared_ptr<OpDef>>(), context.symvar2val(tensor));
861
            auto typeobj = tensor.get_type();
862
            return context.val2symvar(typeobj, output);
863 864 865 866 867 868
        } else {
            auto* tw = TensorWrapper::try_cast(tensor.ptr());
            auto output = reduce_to_scalar(
                    *op.cast<std::shared_ptr<OpDef>>(), tw->m_tensor->data());
            return TensorWrapper::make(py_tensor_type, output);
        }
869 870
    });

871 872 873 874 875 876 877
    m.def("name_tensor", [](std::string name, py::object tensor) {
        auto* tw = TensorWrapper::try_cast(tensor.ptr());
        auto output = imperative::apply(TraceMarkVar(name), tw->m_tensor->data())[0];
        tw->m_tensor->reset(output);
    });

    m.def("is_grad_attached", [](std::vector<py::object> tensors) -> bool {
878
        SmallVector<ValueRef> values(tensors.size());
879 880
        for (size_t i = 0; i < tensors.size(); ++i) {
            values[i] = tensors[i].cast<TensorWrapper>().m_tensor->data();
881 882 883 884 885 886 887 888 889 890
        }
        auto outputs = imperative::apply(GetGradKey(), values);
        if (outputs[0].is<GradKeyValue>()) {
            return true;
        } else {
            return false;
        }
    });

    m.def("get_grad_key", [](std::vector<py::object> tensors) -> py::object {
891
        SmallVector<ValueRef> values(tensors.size());
892 893
        for (size_t i = 0; i < tensors.size(); ++i) {
            values[i] = tensors[i].cast<TensorWrapper>().m_tensor->data();
894
        }
895 896
        auto output = imperative::apply(GetGradKey(), values)[0];
        if (!output) {
897 898
            return py::none();
        }
899 900
        return py::reinterpret_borrow<py::object>(GradKeyWrapper::wrap_t::pycast(
                GradKeyWrapper::get(output.cast<GradKeyValue>())));
901 902 903 904 905 906 907 908
    });

    m.def("set_grad", [](py::object py_key, py::function backward_fn,
                         std::vector<py::object> inputs,
                         std::vector<py::object> outputs) {
        mgb_assert(GradKeyWrapper::wrap_t::type().isinstance(py_key.ptr()));
        auto* key = reinterpret_cast<GradKeyWrapper::wrap_t*>(py_key.ptr())->inst();
        GenericFunction generic_backward_fn =
909
                [backward_fn](Span<ValueRef> output_grads) -> ValueRefList {
910 911 912 913 914 915 916 917 918 919
            py::list output_grad_tws;
            for (auto&& output_grad : output_grads) {
                if (output_grad) {
                    output_grad_tws.append(
                            TensorWrapper::make(py_tensor_type, output_grad));
                } else {
                    output_grad_tws.append(py::none());
                }
            }
            py::tuple input_grad_tws = backward_fn(*output_grad_tws);
920 921 922
            ValueRefList input_grads(input_grad_tws.size());
            for (size_t i = 0; i < input_grad_tws.size(); ++i) {
                auto input_grad_tw = input_grad_tws[i];
923
                if (!input_grad_tw.is_none()) {
924 925
                    input_grads[i] =
                            py::cast<TensorWrapper>(input_grad_tw).m_tensor->data();
926
                } else {
927
                    input_grads[i] = {};
928 929 930 931
                }
            }
            return input_grads;
        };
932
        SmallVector<ValueRef> values(inputs.size() + outputs.size());
933 934
        for (size_t i = 0; i < inputs.size(); ++i) {
            values[i] = inputs[i].cast<TensorWrapper>().m_tensor->data();
935
        }
936 937 938
        for (size_t i = 0; i < outputs.size(); ++i) {
            values[i + inputs.size()] =
                    outputs[i].cast<TensorWrapper>().m_tensor->data();
939 940 941 942 943 944 945 946 947 948 949 950 951 952
        }
        auto wrapped_output_values = imperative::apply(
                SetGrad(key->m_key, generic_backward_fn, inputs.size()), values);
        std::vector<py::object> wrapped_outputs;
        mgb_assert(wrapped_output_values.size() == outputs.size());
        for (auto&& output_value : wrapped_output_values) {
            wrapped_outputs.push_back(
                    TensorWrapper::make(py_tensor_type, output_value));
        }
        return wrapped_outputs;
    });

    static py::function module_trace_hook;

953 954
    static auto get_module_trace = [] {
        static std::shared_ptr<ModuleTraceTransformation> module_trace_transformation;
955 956 957 958
        if (!module_trace_transformation) {
            mgb_assert(module_trace_hook);
            module_trace_transformation =
                    std::make_shared<ModuleTraceTransformation>(module_trace_hook);
959
            transformations.register_at<Segment::ModuleTrace>(
960 961
                    module_trace_transformation);
        }
962 963
        return module_trace_transformation;
    };
964

965 966
    m.def("set_cpp_use_symbolic_shape", &set_cpp_use_symbolic_shape);

967 968 969
    m.def("set_module_tracing", [=] { get_module_trace()->enable(); });

    m.def("unset_module_tracing", [=] { get_module_trace()->disable(); });
970

971
    m.def("is_tracing_module", [=] { return get_module_trace()->enabled(); });
972 973 974 975 976 977 978 979 980 981 982 983 984 985 986

    m.def("set_module_trace_hook",
          [](py::function function) { module_trace_hook = function; });

    m.def("begin_record_values", [] { Value::begin_record_values(); });

    m.def("end_record_values", [] {
        std::vector<std::pair<size_t, std::string>> reprs;
        auto values = Value::end_record_values();
        for (auto&& value : values) {
            reprs.push_back({value.id(), value.to_string()});
        }
        return reprs;
    });

987 988 989 990
    m.def("print_stats", [] { imperative::Stats::print(); });

    m.def("reset_stats", [] { imperative::Stats::reset(); });

991 992 993 994 995 996 997 998 999 1000 1001 1002 1003 1004 1005 1006 1007 1008 1009 1010 1011 1012 1013 1014 1015 1016 1017 1018 1019 1020 1021 1022 1023 1024 1025 1026 1027 1028 1029 1030 1031 1032 1033 1034 1035 1036 1037 1038 1039 1040 1041 1042 1043 1044 1045 1046 1047
    m.def("_get_convert_inputs",
          []() -> bool { return DTypePromoteCfg::convert_input_enabled; });
    m.def("_set_convert_inputs", [](bool flag) -> bool {
        bool ret = DTypePromoteCfg::convert_input_enabled;
        DTypePromoteCfg::convert_input_enabled = flag;
        return ret;
    });
    m.def("_get_amp_dtype_autocast",
          []() -> bool { return DTypePromoteCfg::amp_dtype_autocast_enabled; });
    m.def("_set_amp_dtype_autocast", [](bool flag) -> bool {
        bool ret = DTypePromoteCfg::amp_dtype_autocast_enabled;
        DTypePromoteCfg::amp_dtype_autocast_enabled = flag;
        return ret;
    });

    static auto get_amp_prec_dtype = [](bool is_high) -> std::string {
        DType& target = is_high ? DTypePromoteCfg::amp_high_prec_dtype
                                : DTypePromoteCfg::amp_low_prec_dtype;
        mgb_assert(target.category() == DTypeCategory::FLOAT);
        std::string ret = target.name();
        transform(ret.begin(), ret.end(), ret.begin(), ::tolower);
        return ret;
    };

    static auto set_amp_prec_dtype = [](bool is_high,
                                        std::string dtype_name) -> std::string {
        DType& target = is_high ? DTypePromoteCfg::amp_high_prec_dtype
                                : DTypePromoteCfg::amp_low_prec_dtype;
        std::string ret = target.name();

        if (dtype_name == "float32") {
            target = dtype::Float32();
        } else if (dtype_name == "float16") {
            target = dtype::Float16();
        } else if (dtype_name == "bfloat16") {
            target = dtype::BFloat16();
        } else {
            mgb_assert(
                    false, "casted type of amp should be float, but you give %s\n",
                    dtype_name.c_str());
        }

        transform(ret.begin(), ret.end(), ret.begin(), ::tolower);
        return ret;
    };

    m.def("_get_amp_high_prec_dtype",
          []() -> std::string { return get_amp_prec_dtype(true); });
    m.def("_set_amp_high_prec_dtype", [](std::string dtype_name) -> std::string {
        return set_amp_prec_dtype(true, dtype_name);
    });
    m.def("_get_amp_low_prec_dtype",
          []() -> std::string { return get_amp_prec_dtype(false); });
    m.def("_set_amp_low_prec_dtype", [](std::string dtype_name) -> std::string {
        return set_amp_prec_dtype(false, dtype_name);
    });

1048
    py::register_exception<TraceError>(m, "TraceError");
1049 1050
}

1051 1052
#undef MGE_PY_INTERFACE

1053
}  // namespace mgb::imperative::python
新手
引导
客服 返回
顶部