mindspore 2.6.0__cp39-cp39-manylinux1_x86_64.whl → 2.7.0rc1__cp39-cp39-manylinux1_x86_64.whl

mindspore/include/mindspore/ccsrc/pipeline/jit/pi/graph_capture/abstract_object.h

@@ -62,6 +62,32 @@ class AbstractObjectBase {
     std::unordered_map<const PyObject *, AObject *> obj_2_aobj_;
   };
 
+  template <typename T>
+  static AObject *ConstructAbstract(const py::object &obj) {
+    auto resource = Resource::Current();
+    MS_EXCEPTION_IF_NULL(resource);
+    auto pool = resource->pool();
+    MS_EXCEPTION_IF_NULL(pool);
+    return pool->New<T>(obj);
+  }
+
+  template <typename T>
+  static AObject *ConstructAbstract(const std::vector<AObject *> &elements) {
+    auto resource = Resource::Current();
+    MS_EXCEPTION_IF_NULL(resource);
+    auto pool = resource->pool();
+    MS_EXCEPTION_IF_NULL(pool);
+    return pool->New<T>(elements);
+  }
+
+  template <typename T>
+  static AObject *ConstructAbstract(const py::object &obj, const std::vector<AObject *> &elements) {
+    if (obj.ptr() == nullptr) {
+      return ConstructAbstract<T>(elements);
+    }
+    return ConstructAbstract<T>(obj);
+  }
+
  public:
   enum Type {
 #define ABSTRACT_TYPE_DEF(unit) kType##unit,
@@ -76,6 +102,14 @@ class AbstractObjectBase {
   static_assert(static_cast<int>(kTypeSlice) + 8 == static_cast<int>(kTypeType));  // builtin type
   static_assert(static_cast<int>(kTypeUnknown) == 0);
 
+  enum Scope {
+    SCOPE_NOT_SPECIFIED = 0,
+    SCOPE_LOCAL = 1,
+    SCOPE_PARAM = 1 << 1,
+    SCOPE_FREE_VAR = 1 << 2,
+    SCOPE_GLOBAL = 1 << 3
+  };
+
   // record PyObject and check self reference for list,tuple,dict
   static std::unordered_map<AObject::Type, PyTypeObject *> aobj_type_map;
   PyTypeObject *GetPyTypeObject(const Type &type) {
@@ -94,12 +128,11 @@ class AbstractObjectBase {
   virtual py::object GetPyObject() const { return py::object(); }
 
   virtual AObject *Binary(AObject *other, int op) { return MakeAObject(kTypeAnyValue); }
-  virtual AObject *Unary(int op) const { return MakeAObject(kTypeAnyValue); }
   virtual AObject *GetIter() const { return MakeAObject(kTypeAnyValue); }
 
   virtual AObject *GetAttr(const std::string &name);
   virtual AObject *GetItem(AObject *key) { return MakeAObject(kTypeAnyValue); }
-
+  AObject *GetItem(AObject *key, AObject *defalut_value);
   // return false if has an python exception
   virtual bool SetAttr(const std::string &name, AObject *value) { return true; }
   virtual bool SetItem(AObject *key, AObject *value) { return true; }
@@ -125,6 +158,8 @@ class AbstractObjectBase {
 
   static AObject *MakeFunction(const std::vector<AObject *> &args, const py::object &globals, int oparg);
 
+  static AObject *FuncAObjectUpdater(const py::object &func, const std::vector<AObject *> &args);
+
   /**
    * BUILD_SLICE,BUILD_STRING,BUILD_SET,BUILD_LIST,BUILD_TUPLE,BUILD_CONST_KEY_MAP,BUILD_MAP
    * \return a new AbstractObject if success, else a empty AbstractObject
@@ -144,7 +179,7 @@ class AbstractObjectBase {
   static const char *GetTypeDesc(AObject::Type type);
   static std::string ToString(PyObject *, bool print_type = true, size_t limit = SIZE_MAX);
   bool IsLatestVersion() const { return next_version_ == nullptr; }
-  AObject *GetLatestVersion();
+  AObject *GetLatestVersion() const;
   const AObject *GetPreVersion() const { return pre_version_; }
   void SetPreVersion(AObject *pre_version);
   const AObject *GetNextVersion() const { return next_version_; }
@@ -156,6 +191,10 @@ class AbstractObjectBase {
   void RemoveUser(AObject *user) { users_.erase(user); }
   bool HasMultiVersion() const { return pre_version_ != nullptr || next_version_ != nullptr; }
   virtual void CreateVersionWithNewValue() {}
+  void SetScope(Scope scope) { scope_ = scope; }
+  void AddScope(Scope scope) { scope_ = static_cast<Scope>(static_cast<int>(scope_) | static_cast<int>(scope)); }
+  Scope GetScope() const { return scope_; }
+  const std::string &GetScopeDesc() const;
 
  protected:
   static AObject *Convert(const abstract::AbstractBasePtr &abstract);
@@ -167,6 +206,7 @@ class AbstractObjectBase {
   AObject *next_version_{nullptr};
   std::set<AObject *> users_;
   AbstractWrapperPtr abstract_wrapper_{nullptr};
+  Scope scope_{SCOPE_NOT_SPECIFIED};
 };
 
 class AbstractObject : public AbstractObjectBase {
@@ -177,9 +217,6 @@ class AbstractObject : public AbstractObjectBase {
 
   py::object GetPyObject() const override { return value_; }
 
-  AObject *Binary(AObject *other, int op) override;
-  AObject *Unary(int op) const override;
-  AObject *UnaryValue(int op) const;
   AObject *GetIter() const override;
   AObject *GetAttr(const std::string &name) override;
   AObject *GetItem(AObject *key) override;
@@ -196,7 +233,7 @@ class AbstractString : public AbstractObject {
   AbstractString() : AbstractObject(kTypeString, py::object()), str_() {}
   explicit AbstractString(const py::object &str)
       : AbstractObject(kTypeString, str), str_(str.ptr() == nullptr ? std::string() : py::cast<std::string>(str)) {}
-  virtual ~AbstractString() = default;
+  ~AbstractString() override = default;
   AObject *GetItem(AObject *index) override;
 
  protected:
@@ -207,12 +244,12 @@ class AbstractType : public AbstractObject {
  public:
   explicit AbstractType(const py::object &cls)
       : AbstractObject(kTypeType, cls), type_type_(GetPyType(reinterpret_cast<PyTypeObject *>(cls.ptr()))) {}
-  virtual ~AbstractType() {}
+  ~AbstractType() override = default;
   bool IsMindSporeSupportedType() override { return false; }
 
   Type GetTypeType() const { return type_type_; }
   AObject *BuildAbstractInstance(const std::vector<AObject *> &args, int opcode);
-  py::object BuildInstance(const std::vector<py::object> &args, int opcode);
+  py::object BuildInstance(const std::vector<py::object> &args, int opcode, const py::object &kw);
 
  private:
   Type type_type_;
@@ -223,10 +260,9 @@ class AbstractSequence : public AbstractObject {
   explicit AbstractSequence(Type type, const py::object &obj)
       : AbstractObject(type, obj), element_type_(kTypeUnknown), elements_({}) {}
   explicit AbstractSequence(Type type, const std::vector<AObject *> &elements);
-  virtual ~AbstractSequence() {}
+  ~AbstractSequence() override = default;
 
   AObject *Binary(AObject *other, int op) override;
-  AObject *Unary(int op) const override;
   AObject *GetAttr(const std::string &name) override;
   bool SetAttr(const std::string &name, AObject *) override { return false; };
   AObject *GetItem(AObject *key) override;
@@ -243,6 +279,10 @@ class AbstractSequence : public AbstractObject {
   ///
   /// \return The vector of AObject objects.
   const std::vector<AObject *> &GetElements() const { return elements_; }
+  const std::vector<AObject *> &GetElementsWithInit() {
+    InitElementsListIfNeed();
+    return elements_;
+  }
   bool IsMindSporeSupportedType() override;
 
   void SetElementType(Type type) { element_type_ = type; }
@@ -262,7 +302,7 @@ class AbstractTuple : public AbstractSequence {
  public:
   explicit AbstractTuple(const py::object &tuple) : AbstractSequence(kTypeTuple, tuple) {}
   explicit AbstractTuple(const std::vector<AObject *> &elements) : AbstractSequence(kTypeTuple, elements) {}
-  virtual ~AbstractTuple() {}
+  ~AbstractTuple() override = default;
 };
 
 class AbstractNamedTuple : public AbstractObject {
@@ -288,13 +328,41 @@ class AbstractList : public AbstractSequence {
  public:
   explicit AbstractList(const py::object &list) : AbstractSequence(kTypeList, list) {}
   explicit AbstractList(const std::vector<AObject *> &elements) : AbstractSequence(kTypeList, elements) {}
-  virtual ~AbstractList() {}
+  ~AbstractList() override = default;
 
   AbstractList *ListAppend(AObject *item);
   AbstractList *ListExtend(AObject *list);
   AbstractTuple *ListToTuple();
 };
 
+class AbstractCellList : public AbstractSequence {
+ public:
+  explicit AbstractCellList(const py::object &cells) : AbstractSequence(kTypeNNCellList, cells) {}
+  explicit AbstractCellList(const std::vector<AObject *> &cells) : AbstractSequence(kTypeNNCellList, cells) {}
+  ~AbstractCellList() override = default;
+};
+
+class AbstractDictKeys : public AbstractSequence {
+ public:
+  explicit AbstractDictKeys(const py::object &keys) : AbstractSequence(kTypeDictKeys, keys) {}
+  explicit AbstractDictKeys(const std::vector<AObject *> &keys) : AbstractSequence(kTypeDictKeys, keys) {}
+  ~AbstractDictKeys() override = default;
+};
+
+class AbstractDictValues : public AbstractSequence {
+ public:
+  explicit AbstractDictValues(const py::object &values) : AbstractSequence(kTypeDictValues, values) {}
+  explicit AbstractDictValues(const std::vector<AObject *> &values) : AbstractSequence(kTypeDictValues, values) {}
+  ~AbstractDictValues() override = default;
+};
+
+class AbstractDictItems : public AbstractSequence {
+ public:
+  explicit AbstractDictItems(const py::object &items) : AbstractSequence(kTypeDictItems, items) {}
+  explicit AbstractDictItems(const std::vector<AObject *> &items) : AbstractSequence(kTypeDictItems, items) {}
+  ~AbstractDictItems() override = default;
+};
+
 using AObjectPair = std::pair<AObject *, AObject *>;
 using AObjectPairList = std::vector<AObjectPair>;
 
@@ -305,13 +373,15 @@ class AbstractDict : public AbstractObject {
   virtual ~AbstractDict() {}
 
   std::string ToString() const override;
-  AObject *Unary(int op) const override;
-  AObject *Binary(AObject *, int op) override;
   AObject *GetAttr(const std::string &name) override;
   bool SetAttr(const std::string &name, AObject *) override { return false; };
   AObject *GetItem(AObject *key) override;
   bool IsMindSporeSupportedType() override;
 
+  AObject *Keys();
+  AObject *Values();
+  AObject *Items();
+
   Type KeyType() const { return k_type_; }
   Type ValueType() const { return v_type_; }
 
@@ -327,6 +397,10 @@ class AbstractDict : public AbstractObject {
   ///
   /// \return The elements of AbstractDict object.
   const AObjectPairList &GetElements() const { return key_values_; }
+  const AObjectPairList &GetElementsWithInit() {
+    InitKeyValuesListIfNeed();
+    return key_values_;
+  }
 
   class ValueIter {
    public:
@@ -341,7 +415,9 @@ class AbstractDict : public AbstractObject {
     }
 
    private:
-    PyObject *map_, *key_, *val_;
+    PyObject *map_{nullptr};
+    PyObject *key_{nullptr};
+    PyObject *val_{nullptr};
     Py_ssize_t pos_;
   };
   auto begin() const { return ValueIter(this); }
@@ -365,8 +441,6 @@ class AbstractTensor : public AbstractObject {
  public:
   AbstractTensor(const py::object &o, bool is_stub);
   virtual ~AbstractTensor() {}
-  AObject *Binary(AObject *, int op) override;
-  AObject *Unary(int op) const override;
   AObject *GetAttr(const std::string &name) override;
   std::string ToString() const override;
 
@@ -375,7 +449,6 @@ class AbstractTensor : public AbstractObject {
   py::object GetTensor(bool sync);
 
   bool IsMindSporeSupportedType() override { return true; }
-  bool IsStubTensor() const { return is_stub_; }
 
  private:
   bool is_stub_;

mindspore/include/mindspore/ccsrc/pipeline/jit/pi/graph_capture/cfg.h

@@ -34,11 +34,14 @@ class CFG;
 class Block;
 class Instr {
  public:
-  Instr(const Instr &) = delete;
   Instr &operator=(const Instr &) = delete;
-  Instr(int op, int arg, int bci = -1, int line = -1) : bci_(bci), op_(op), arg_(arg), line_(line), jump_(nullptr) {}
+  Instr(const Instr &o)
+      : bci_(-1), op_(o.op_), arg_(o.arg_), name_(o.name_), cnst_(o.cnst_), jump_(nullptr), loc_(o.loc_) {}
+  Instr(int op, int arg) : bci_(-1), op_(op), arg_(arg), jump_(nullptr), loc_{-1, -1, -1, -1} {}
+  Instr(int op, int arg, const CodeLocation &loc) : Instr(op, arg) { loc_ = loc; }
   Instr(int op, int arg, const std::string &name) : Instr(op, arg) { name_ = name; }
   Instr(int op, int arg, const py::object &cnst) : Instr(op, arg) { cnst_ = cnst; }
+  Instr(int op, int arg, int bci, int line = -1) : Instr(op, arg) { bci_ = bci, loc_ = {line, line, -1, -1}; }
   explicit Instr(int op) : Instr(op, 0) {}
 
   int bci() const { return bci_; }
@@ -47,13 +50,16 @@ class Instr {
   void set_op(int op) { op_ = op; }
   int arg() const { return arg_; }
   void set_arg(int arg) { arg_ = arg; }
-  int line() const { return line_; }
-  void set_line(int l) { line_ = l; }
+  int line() const { return loc_.start_line_; }
+  void set_line(int l) { loc_.start_line_ = l; }
   Instr *extra_jump() const { return jump_; }
   void set_extra_jump(Instr *j) { jump_ = j; }
+  const auto &location() const { return loc_; }
+  void set_location(const CodeLocation &loc) { loc_ = loc; }
 
   const std::string &name() const { return name_; }
   void set_name(const std::string &n) { name_ = n; }
+  void set_name(const char *n) { name_ = n ? n : ""; }
   const py::object &cnst() const { return cnst_; }
   void set_cnst(const py::handle &cnst) { cnst_ = py::reinterpret_borrow<py::object>(cnst); }
 
@@ -64,11 +70,12 @@ class Instr {
   int bci_;
   int op_;
   int arg_;
-  int line_;
   // these field only one is valid, union this these field like this { const char *, PyObject *, Instr * } ?
   std::string name_;
+  // if python3.11 ~ python3.12 and opcode is call, `cnst_` is KW_NAMES
   py::object cnst_;
   Instr *jump_;
+  CodeLocation loc_;
 };
 
 class Block {
@@ -158,6 +165,10 @@ class CFG {
   std::vector<std::unique_ptr<Instr>> &instr_pool() { return instrs_; }
   std::vector<std::unique_ptr<Block>> &bb_pool() { return bb_pool_; }
   std::unique_ptr<Liveness> &liveness() { return liveness_; }
+  const ExceptionTable &exc_table() const { return exc_table_; }
+  // python3.11+ only, find first exception table item of try/with blocks
+  ExceptionTable::const_iterator FindTryWithBlock(int bci) const;
+  ExceptionTable::const_iterator FindExcTableItem(int bci) const;
   int GetLocalCount() const { return co_.LocalSize(); }
 
   const Liveness *GetLiveness();
@@ -172,14 +183,17 @@ class CFG {
   std::string ToString() const;
 
  private:
+  Instr *GetInstruction(int bci);
   void BuildInst(const uint8_t *begin, const uint8_t *end);
   std::map<int, Block *> BuildBB(const uint8_t *begin, const uint8_t *end);
   void BuildCFG(const std::map<int, Block *> &labels);
+  ExceptionTable::const_iterator FindTryWithStart(ExceptionTable::const_iterator iter) const;
 
   PyCodeWrapper co_;
   std::vector<std::unique_ptr<Instr>> instrs_;
   std::vector<std::unique_ptr<Block>> bb_pool_;
   std::unique_ptr<Liveness> liveness_;
+  ExceptionTable exc_table_;
 };
 }  // namespace pijit
 }  // namespace mindspore

mindspore/include/mindspore/ccsrc/pipeline/jit/pi/graph_capture/code_generator.h

@@ -48,9 +48,19 @@ struct GraphInputInfo {
   ValueNode *vargs = nullptr;
   ValueNode *kwargs = nullptr;
 };
+using InstructionList = std::vector<std::unique_ptr<Instr>>;
 
 class CodeGenerator {
  public:
+  struct ExceptionCodeItem {
+    // [begin, end]
+    Instr *begin_;
+    Instr *end_;
+    Instr *jump_;
+    int stack_;
+    bool lasti_;
+  };
+
   struct Code {
     int co_argcount;
     int co_kwonlyargcount;
@@ -63,18 +73,26 @@ class CodeGenerator {
     std::vector<std::string> co_freevars;
     std::string co_name;
     py::object co_filename;
-    py::object co_qualname;
-    py::object co_exceptiontable;
+    std::string co_qualname;
+    std::vector<ExceptionCodeItem> co_exceptiontable;
   };
 
-  explicit CodeGenerator(const NodeSet *nodes)
-      : nodes_(nodes), globals_(), code_(), nodes_alive_(), locals_map_(), missing_value_to_undefine_(false) {}
+  explicit CodeGenerator(const NodeSet *nodes, bool vm_mode = false)
+      : nodes_(nodes),
+        globals_(),
+        code_(),
+        nodes_alive_(),
+        locals_map_(),
+        missing_value_to_undefine_(false),
+        vm_mode_(vm_mode) {}
+  explicit CodeGenerator(Code &&ccode) : code_(std::move(ccode)) {}
 
   void set_missing_value_to_undefine(bool v) { missing_value_to_undefine_ = v; }
 
   void SetGlobals(const py::dict &dict) { globals_ = dict; }
   const py::dict &GetGlobals() const { return globals_; }
   const std::unordered_map<ValueNode *, int> &GetLocalsMap() const { return locals_map_; }
+  std::unordered_map<ValueNode *, int> &GetLocalsMap() { return locals_map_; }
   const Code &GetCode() const { return code_; }
   void SetArgsInfo(int argcount, int kwonlyargcount) {
     code_.co_argcount = argcount;
@@ -87,20 +105,22 @@ class CodeGenerator {
   void SetCellVariableNames(const std::vector<std::string> &names) { code_.co_cellvars = names; }
   void SetFreeVariableNames(const std::vector<std::string> &names) { code_.co_freevars = names; }
   void SetCodeName(const std::string &name) { code_.co_name = name; }
+  void SetQualName(const std::string &qualname) { code_.co_qualname = qualname; }
   void SetFileName(const py::object &file) { code_.co_filename = file; }
-  void SetQualName(const py::object &qualname) { code_.co_qualname = qualname; }
-  void SetExceptionTable(const py::object &exceptiontable) { code_.co_exceptiontable = exceptiontable; }
 
   void ClearAlive(ValueNode *node) { nodes_alive_.erase(node); }
   void ClearAlive() { nodes_alive_.clear(); }
   void MarkAlive(ValueNode *node) { MarkAlive(node, INT_MAX); }
   void MarkAlive();
   // make the node same as other node, use same local index, if the node not in locals, try to load it
-  void MakeSameLocal(ValueNode *new_node, ValueNode *old_node);
+  void MakeSameLocal(ValueNode *new_node, ValueNode *old_node, bool clear = false);
   void NewInstr(int op, int arg = 0, int line = -1);
   void AddInstrs(std::vector<std::unique_ptr<Instr>> &&list);
+  void AddInstructionWithExceptionTable(const CFG *cfg, int start, int end);
+  void CollectExceptionTableItem(const CFG *cfg, int start, int end);
   void AddInstr(std::unique_ptr<Instr> &&instr);
-  void EraseUnusedInstr();
+  void AddCallInstr(size_t load_args_offset, int oparg);
+  py::object NewCode();
 
   // initialize local map of parameters
   void Init();
@@ -122,33 +142,6 @@ class CodeGenerator {
   // add node to locals map
   int AllocLocal(ValueNode *node, int index = INT_MAX);
 
-  /**
-   * Transform code info to PyCodeObject
-   *
-   * \param ccode code info
-   * \return PyCodeObject
-   */
-  static py::object Transform(const Code &ccode);
-
-  /**
-   * Calculate max stack size
-   *
-   * \param list instruct nodes list
-   * \param sp start of stack depth
-   * \return max depth of stack, or -1 if stack out of bound
-   */
-  static int CalculateStackSize(const std::vector<std::unique_ptr<Instr>> &list, int sp = 0);
-
-  /**
-   * Convert instruction list to bytes object. generate line table.
-   *
-   * \param list instruct nodes list
-   * \param first_line first line
-   * \return first is co_code, second is co_lnotab
-   */
-  static std::pair<py::bytes, py::bytes> ConvertToCodeBytes(const std::vector<std::unique_ptr<Instr>> &list,
-                                                            int first_line);
-
   /**
    * Copy instruction list at range [start, end).
    * NOTE: reset opcode:
@@ -179,13 +172,6 @@ class CodeGenerator {
     const std::vector<std::unique_ptr<Instr>> &list, size_t start_bci, size_t end_bci,
     const std::vector<std::unique_ptr<Instr>> &replacement);
 
-  /**
-   * Erase unused instr
-   *
-   * \param list instruction list
-   */
-  static void EraseUnusedInstr(std::vector<std::unique_ptr<Instr>> *list);
-
   /**
    * generate rot instructions
    */
@@ -193,13 +179,27 @@ class CodeGenerator {
 
  private:
   void MarkAlive(ValueNode *node, int order);
+  void EraseUnusedInstr();
+  void ResetExceptionCodeItem(InstructionList::const_iterator erased);
+
+  static py::object Transform(const Code &ccode);
+  static std::pair<py::bytes, py::bytes> ConvertToCodeBytes(const Code &ccode);
+  static int CalculateStackSize(const Code &ccode, int sp = 0);
+
+  // python3.11+ only
+  std::vector<std::unique_ptr<Instr>> ByteCodePrefix() const;
+  void FixOffset();
+  void FixLocalOffset(int invalid_index);
+  void FixFreeOffset(const std::vector<std::string> &other_closure_names);
+  static int ExceptionStackRequired(const Code &ccode);
 
-  const NodeSet *nodes_;
+  const NodeSet *nodes_{nullptr};
   py::dict globals_;
   Code code_;
   std::unordered_map<ValueNode *, int> nodes_alive_;
   std::unordered_map<ValueNode *, int> locals_map_;
-  bool missing_value_to_undefine_;
+  bool missing_value_to_undefine_{false};
+  bool vm_mode_{false};
 };
 
 class LoopBodyReCaptureCodeGenerator {
@@ -260,6 +260,10 @@ class CodeBreakGenerator {
  private:
   const CFG *GetCFG() const { return cfg_; }
 
+  // for python3.11+ copy origin instruction need pop the null pointer (which is only consume by call instruction)
+  // from stack. here generate code from node avoid pop null
+  bool IsCopyCapturedInstructions() const { return !IS_PYTHON_3_11_PLUS && no_graph_ && !NeedHandleBreakAtCall(); }
+
   void ExtendCodeInfo(CodeGenerator *cg, bool merge_kw_only) const;
 
   // rebuild parameters of graph, identify parameters that graph only support as constant
@@ -278,7 +282,7 @@ class CodeBreakGenerator {
   void HandleOutputOpt(CodeGenerator *code_gen);
 
   // make function of untracked bytecode, build restore frame operations of untracked bytecode
-  py::object MakeUntrackedCode(int untracked_bci, int untracked_stack_effect) const;
+  py::object MakeUntrackedCode(int untracked_bci, int untracked_stack_effect, int *argc) const;
 
   void ReconstructStack(CodeGenerator *code_gen, int untracked_bci, int untracked_stack_effect) const;
 
@@ -296,6 +300,10 @@ class CodeBreakGenerator {
   // generate specialize code if break point is call
   void BreakAtCall(CodeGenerator *code_gen) const;
   bool NeedHandleBreakAtCall() const;
+  // Similar to BreakAtIf(), but it breaks at conditional-statement of a callee function (or we say a subgraph).
+  void BreakAtCalleeIfCondition(CodeGenerator *code_gen) const;
+  py::object MakeUntrackedCodeForNestedCalls(const std::vector<Graph *> &call_stack, int top_argc, int untracked_bci,
+                                             int stack_effect) const;
 
   void RestoreStack(CodeGenerator *code_gen) const;
 
@@ -334,13 +342,9 @@ class CodeBreakGenerator {
   // used to record the value nodes and the nodes that replaced them
   std::map<ValueNode *, ValueNode *> replaced_nodes_;
 
-  GraphInputInfo graph_inputs_info_;
-
   // break bci alive locals
   std::vector<int> alive_locals_;
 
-  std::shared_ptr<SideEffect> side_effect_handler_;
-
   // break bci
   int break_bci_;
 

mindspore/include/mindspore/ccsrc/pipeline/jit/pi/graph_capture/graph.h

@@ -16,6 +16,7 @@
 #ifndef MINDSPORE_PI_JIT_GRAPH_CAPTURE_GRAPH_H
 #define MINDSPORE_PI_JIT_GRAPH_CAPTURE_GRAPH_H
 
+#include <exception>
 #include <map>
 #include <memory>
 #include <string>
@@ -88,6 +89,12 @@ class FrameStates {
     stack.insert(stack.end() - i, v);
   }
 
+  void Swap(int i) {
+    MS_ASSERT((int)stack.size() - i >= 0);
+    auto top_idx = stack.size() - 1;
+    std::swap(stack[top_idx], stack[top_idx - i]);
+  }
+
   void ResizeLocal(int i) {
     MS_ASSERT((int)locals.size() <= i);
     locals.resize(i, &ValueNode::kUnboundLocal);
@@ -124,6 +131,11 @@ class Graph {
     StopTraceReason reason_;
   };
 
+  struct ExpandParamInfo {
+    ValueNode *node_;
+    std::vector<const AObject *> elements_;
+  };
+
   Graph(PyCodeObject *co, PyObject *globals, const GraphJitConfig &conf);
   virtual ~Graph();
 
@@ -138,10 +150,8 @@ class Graph {
   PyCodeObject *GetCodeObj() const { return reinterpret_cast<PyCodeObject *>(co_.ptr()); }
   const py::object &GetGlobals() const { return f_globals_; }
 
-  void StopTraceAt(int bci, StopTraceReason reason) {
-    break_info_.bci_ = bci;
-    break_info_.reason_ = reason;
-  }
+  /// @throws GraphBreakException if fullgraph=true (graph break is not allowed)
+  void StopTraceAt(int bci, StopTraceReason reason, const std::vector<std::string> &hints = {});
   int GetStopTraceBci() const { return break_info_.bci_; }
   StopTraceReason GetStopTraceReason() const { return break_info_.reason_; }
   const char *GetModuleName() const { return module_name_; }
@@ -167,9 +177,14 @@ class Graph {
   // only func name
   std::string GetCodeName() const {
     PyCodeObject *c = reinterpret_cast<PyCodeObject *>(co_.ptr());
-    return py::str(c->co_name);
+    if (c != nullptr && c->co_name != nullptr) {
+      return py::str(c->co_name);
+    }
+    return "";
   }
 
+  const std::vector<ValueNode *> &GetParameters() const { return params_; }
+
   void GuardParameter(ValueNode *param);
   void GuardGlobal(ValueNode *global_value);
   void GuardAttribute(ValueNode *attr_value);
@@ -181,10 +196,10 @@ class Graph {
 
   TracePtr TraceValueNode(ValueNode *, int max_trace_depth = -1);
   std::vector<TracePtr> TraceValueNodeClosure(ValueNode *, bool *ret, int max_trace_depth = -1);
-  int GetPruneBranchCount() const { return prune_branch_count_; }
-  void SetPruneBranchCount(int count) { prune_branch_count_ = count; }
   const std::shared_ptr<OptCode> &GetGuardManager() const;
   void SetGuard(const std::shared_ptr<OptCode> &guard);
+  void RemoveAllGuardItems() const;
+  std::map<const AObject *, ExpandParamInfo> &GetExpandParamInfo() { return expand_param_info_; }
 
   // (chaiyouheng): restore graph status at loop begin, clear trace values and operations and guards
   bool RestoreLoopStatus() const { return false; }
@@ -198,8 +213,9 @@ class Graph {
   void SetParent(Graph *parent) { parent_ = parent; }
   Graph *GetParent() const { return parent_; }
 
-  const std::shared_ptr<SideEffect> &GetSideEffect() const;
-  void SetSideEffect(const std::shared_ptr<SideEffect> &handler);
+  const std::shared_ptr<SideEffect> &GetSideEffect() const { return side_effect_; }
+  void SetSideEffect(const std::shared_ptr<SideEffect> &handler) { side_effect_ = handler; }
+  const std::shared_ptr<SideEffectHandler> &GetSideEffectHandler() const { return side_effect_handler_; }
 
   std::vector<ValueNode *> CollectAliveNode(int bci, std::vector<int> *ids = nullptr) const;
   // collect alive node, clear the bit if alive local is unbound
@@ -208,6 +224,7 @@ class Graph {
   void FoundInnerClass() { found_inner_class = true; }
 
   const auto &prepare() const { return prepare_; }
+  auto &prepare() { return prepare_; }
   bool PrepareParameter(ValueNode *node);
 
   // return true if has fail guard matched
@@ -227,6 +244,10 @@ class Graph {
   std::map<int, std::unique_ptr<FrameStates>> frame_states_;
   std::vector<ValueNode *> traced_nodes_;
 
+  std::vector<ValueNode *> params_;
+
+  std::map<const AObject *, ExpandParamInfo> expand_param_info_;
+
   // return value
   ValueNode *ret_val_;
 
@@ -247,9 +268,9 @@ class Graph {
 
   const GraphJitConfig &conf_;
 
-  int prune_branch_count_;
   Graph *parent_{nullptr};
   std::shared_ptr<SideEffect> side_effect_;
+  std::shared_ptr<SideEffectHandler> side_effect_handler_;
   bool found_inner_class = false;
 
   struct PrepareInfo {
@@ -261,6 +282,15 @@ class Graph {
   std::shared_ptr<FuncGraphBuilder> func_graph_builder_;
 };
 
+// If using @jit(fullgraph=true), will throw this exception when graph break occurs.
+class GraphBreakException : public std::runtime_error {
+ public:
+  explicit GraphBreakException(const std::string &msg) : std::runtime_error(msg) {}
+  explicit GraphBreakException(const char *msg) : std::runtime_error(msg) {}
+  // Similar to py::builtin_exception::set_error(), call PyErr_SetString() to throw an exception to the Python side.
+  void set_error() const;
+};
+
 // Return the file path of python code.
 std::string GetFileName(const Graph *graph);