diff --git a/python_module/src/cpp/megbrain_wrap.cpp b/python_module/src/cpp/megbrain_wrap.cpp index 4d4ab7a1b02c234730fb1b5bc5078dcb3da94c40..efcea280c225d42d0e2f9f69f5d0d577c66635c5 100644 --- a/python_module/src/cpp/megbrain_wrap.cpp +++ b/python_module/src/cpp/megbrain_wrap.cpp @@ -880,6 +880,11 @@ SymbolVar SharedScalar::_as_sym_var(CompGraph &cg, mgb::CompNode &cn) { ssprintf("SharedScalar@%p", m_val.get())); } +/* =============== Operator =============== */ + +const std::unique_ptr Operator::sm_opr_footprint_ptr{ + std::make_unique()}; + /* ================= misc ================= */ SymbolVar fill_retain_dtype(SymbolVar var, PyObject *value) { diff --git a/python_module/src/cpp/megbrain_wrap.h b/python_module/src/cpp/megbrain_wrap.h index 37c773e3b4b8022eefd7ab1aa899dd6145e4096f..e1484e4e14e37af6c3b256c4f01b18ddd264711c 100644 --- a/python_module/src/cpp/megbrain_wrap.h +++ b/python_module/src/cpp/megbrain_wrap.h @@ -17,6 +17,8 @@ #include "megbrain/graph.h" #include "megbrain/opr/io.h" +#include "megbrain/plugin/opr_footprint.h" + #include #include @@ -441,16 +443,24 @@ class SharedScalar { */ class Operator { mgb::cg::OperatorNodeBase* m_operator_node; + std::string m_params; + + static const std::unique_ptr sm_opr_footprint_ptr; public: Operator() : m_operator_node(nullptr){}; Operator(mgb::cg::OperatorNodeBase* operator_node) - : m_operator_node(operator_node) {} + : m_operator_node(operator_node), + m_params(std::move( + (sm_opr_footprint_ptr->calc_footprint(m_operator_node)).param->to_string())) + {} size_t id() const { return m_operator_node->id(); } const std::string& name() const { return m_operator_node->name(); } + const std::string& params() const { return m_params; } + const std::shared_ptr get_owner_graph() const { return m_operator_node->owner_graph()->shared_from_this(); } diff --git a/python_module/src/swig/operator.i b/python_module/src/swig/operator.i index 3248caa333232e4f1762a6099014d347cfe9c2d6..e7bd0911868b9c87ac975d858e69321fe318e9c0 100644 --- a/python_module/src/swig/operator.i +++ b/python_module/src/swig/operator.i @@ -58,6 +58,10 @@ public: return $self->name(); } + const std::string& _get_params() const { + return $self->params(); + } + SymbolVarArray _get_inputs() { return $self->inputs(); } diff --git a/python_module/src/swig/operator.py b/python_module/src/swig/operator.py index 1de6aff82531239839e8cf73e33a60c63f1a5b9a..a1eb095e72f68ad119fac950c2ccce42d4033692 100644 --- a/python_module/src/swig/operator.py +++ b/python_module/src/swig/operator.py @@ -19,6 +19,11 @@ def id(self): def name(self): return self._get_name() +@property +def params(self): + import json + return json.loads(self._get_params()) + @property def inputs(self): return tuple(self._get_inputs())