doc.h

Go to the documentation of this file.

/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License.  You may obtain a copy of the License at
 *
 *   http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing,
 * software distributed under the License is distributed on an
 * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
 * KIND, either express or implied.  See the License for the
 * specific language governing permissions and limitations
 * under the License.
 */
#ifndef TVM_SCRIPT_PRINTER_DOC_H_
#define TVM_SCRIPT_PRINTER_DOC_H_

#include <tvm/ir/expr.h>
#include <tvm/node/node.h>
#include <tvm/runtime/data_type.h>
#include <tvm/script/printer/traced_object.h>

namespace tvm {
namespace script {
namespace printer {

class DocNode : public Object {
 public:
  mutable Array<ObjectPath> source_paths;

  void VisitAttrs(AttrVisitor* v) { v->Visit("source_paths", &source_paths); }

  static constexpr const char* _type_key = "script.printer.Doc";
  TVM_DECLARE_BASE_OBJECT_INFO(DocNode, Object);

 public:
  virtual ~DocNode() = default;
};

class Doc : public ObjectRef {
 protected:
  Doc() = default;

 public:
  virtual ~Doc() = default;
  TVM_DEFINE_NOTNULLABLE_OBJECT_REF_METHODS(Doc, ObjectRef, DocNode);
};

class ExprDoc;

class ExprDocNode : public DocNode {
 public:
  ExprDoc Attr(String attr) const;

  ExprDoc Attr(TracedObject<String> attr) const;

  ExprDoc operator[](Array<Doc> indices) const;

  ExprDoc Call(Array<ExprDoc, void> args) const;

  ExprDoc Call(Array<ExprDoc, void> args,  //
               Array<String> kwargs_keys,  //
               Array<ExprDoc, void> kwargs_values) const;

  void VisitAttrs(AttrVisitor* v) { DocNode::VisitAttrs(v); }

  static constexpr const char* _type_key = "script.printer.ExprDoc";
  TVM_DECLARE_BASE_OBJECT_INFO(ExprDocNode, DocNode);
};

class ExprDoc : public Doc {
 protected:
  ExprDoc() = default;

 public:
  ExprDoc operator[](Array<Doc> indices) const;

  TVM_DEFINE_NOTNULLABLE_OBJECT_REF_METHODS(ExprDoc, Doc, ExprDocNode);
};

class StmtDocNode : public DocNode {
 public:
  mutable Optional<String> comment{NullOpt};

  void VisitAttrs(AttrVisitor* v) {
    DocNode::VisitAttrs(v);
    v->Visit("comment", &comment);
  }

  static constexpr const char* _type_key = "script.printer.StmtDoc";
  TVM_DECLARE_BASE_OBJECT_INFO(StmtDocNode, DocNode);
};

class StmtDoc : public Doc {
 protected:
  StmtDoc() = default;

 public:
  TVM_DEFINE_NOTNULLABLE_OBJECT_REF_METHODS(StmtDoc, Doc, StmtDocNode);
};

class StmtBlockDocNode : public DocNode {
 public:
  Array<StmtDoc> stmts;

  void VisitAttrs(AttrVisitor* v) {
    DocNode::VisitAttrs(v);
    v->Visit("stmts", &stmts);
  }

  static constexpr const char* _type_key = "script.printer.StmtBlockDoc";
  TVM_DECLARE_FINAL_OBJECT_INFO(StmtBlockDocNode, DocNode);
};

class StmtBlockDoc : public Doc {
 public:
  explicit StmtBlockDoc(Array<StmtDoc> stmts);
  TVM_DEFINE_NOTNULLABLE_OBJECT_REF_METHODS(StmtBlockDoc, Doc, StmtBlockDocNode);
};

class LiteralDocNode : public ExprDocNode {
 public:
  ObjectRef value;

  void VisitAttrs(AttrVisitor* v) {
    ExprDocNode::VisitAttrs(v);
    v->Visit("value", &value);
  }

  static constexpr const char* _type_key = "script.printer.LiteralDoc";
  TVM_DECLARE_FINAL_OBJECT_INFO(LiteralDocNode, ExprDocNode);
};

class LiteralDoc : public ExprDoc {
 protected:
  explicit LiteralDoc(ObjectRef value);
  LiteralDoc(ObjectRef value, ObjectPath object_path);

 public:
  static LiteralDoc None() { return LiteralDoc(ObjectRef(nullptr)); }

  static LiteralDoc None(ObjectPath object_path) {
    return LiteralDoc(ObjectRef(nullptr), object_path);
  }

  static LiteralDoc Int(int v) { return LiteralDoc(IntImm(DataType::Int(64), v)); }

  static LiteralDoc Int(const TracedObject<IntImm>& v) { return LiteralDoc(v.Get(), v.GetPath()); }

  static LiteralDoc Int(const TracedBasicValue<int>& v) {
    return LiteralDoc(IntImm(DataType::Int(64), v.Get()), v.GetPath());
  }
  static LiteralDoc Boolean(bool v) { return LiteralDoc(IntImm(DataType::Bool(), v)); }

  static LiteralDoc Boolean(const TracedBasicValue<bool>& v) {
    return LiteralDoc(IntImm(DataType::Bool(), v.Get()), v.GetPath());
  }

  static LiteralDoc Float(double v) { return LiteralDoc(FloatImm(DataType::Float(64), v)); }

  static LiteralDoc Float(const TracedObject<FloatImm>& v) {
    return LiteralDoc(v.Get(), v.GetPath());
  }

  static LiteralDoc Str(const String& v) { return LiteralDoc(v); }

  static LiteralDoc Str(const TracedObject<String>& v) { return LiteralDoc(v.Get(), v.GetPath()); }

  TVM_DEFINE_NOTNULLABLE_OBJECT_REF_METHODS(LiteralDoc, ExprDoc, LiteralDocNode);
};

class IdDocNode : public ExprDocNode {
 public:
  String name;

  void VisitAttrs(AttrVisitor* v) {
    ExprDocNode::VisitAttrs(v);
    v->Visit("name", &name);
  }

  static constexpr const char* _type_key = "script.printer.IdDoc";
  TVM_DECLARE_FINAL_OBJECT_INFO(IdDocNode, ExprDocNode);
};

class IdDoc : public ExprDoc {
 public:
  explicit IdDoc(String name);
  explicit IdDoc(std::nullptr_t) : ExprDoc(nullptr) {}
  TVM_DEFINE_NOTNULLABLE_OBJECT_REF_METHODS(IdDoc, ExprDoc, IdDocNode);
};

class AttrAccessDocNode : public ExprDocNode {
 public:
  ExprDoc value{nullptr};
  String name;

  void VisitAttrs(AttrVisitor* v) {
    ExprDocNode::VisitAttrs(v);
    v->Visit("value", &value);
    v->Visit("name", &name);
  }

  static constexpr const char* _type_key = "script.printer.AttrAccessDoc";
  TVM_DECLARE_FINAL_OBJECT_INFO(AttrAccessDocNode, ExprDocNode);
};

class AttrAccessDoc : public ExprDoc {
 public:
  explicit AttrAccessDoc(ExprDoc value, String name);
  TVM_DEFINE_NOTNULLABLE_OBJECT_REF_METHODS(AttrAccessDoc, ExprDoc, AttrAccessDocNode);
};

class IndexDocNode : public ExprDocNode {
 public:
  ExprDoc value{nullptr};
  Array<Doc> indices;  // Each element is union of: Slice / ExprDoc

  void VisitAttrs(AttrVisitor* v) {
    ExprDocNode::VisitAttrs(v);
    v->Visit("value", &value);
    v->Visit("indices", &indices);
  }

  static constexpr const char* _type_key = "script.printer.IndexDoc";
  TVM_DECLARE_FINAL_OBJECT_INFO(IndexDocNode, ExprDocNode);
};

class IndexDoc : public ExprDoc {
 public:
  explicit IndexDoc(ExprDoc value, Array<Doc> indices);
  TVM_DEFINE_NOTNULLABLE_OBJECT_REF_METHODS(IndexDoc, ExprDoc, IndexDocNode);
};

class CallDocNode : public ExprDocNode {
 public:
  ExprDoc callee{nullptr};
  Array<ExprDoc> args;
  Array<String> kwargs_keys;
  Array<ExprDoc> kwargs_values;

  void VisitAttrs(AttrVisitor* v) {
    ExprDocNode::VisitAttrs(v);
    v->Visit("callee", &callee);
    v->Visit("args", &args);
    v->Visit("kwargs_keys", &kwargs_keys);
    v->Visit("kwargs_values", &kwargs_values);
  }

  static constexpr const char* _type_key = "script.printer.CallDoc";
  TVM_DECLARE_FINAL_OBJECT_INFO(CallDocNode, ExprDocNode);
};

class CallDoc : public ExprDoc {
 public:
  CallDoc(ExprDoc callee, Array<ExprDoc> args, Array<String> kwargs_keys,
          Array<ExprDoc> kwargs_values);
  TVM_DEFINE_NOTNULLABLE_OBJECT_REF_METHODS(CallDoc, ExprDoc, CallDocNode);
};

class OperationDocNode : public ExprDocNode {
 public:
  enum class Kind : int32_t {
    // Unary operators
    kUnaryStart = 0,
    kUSub = 1,    // -x
    kInvert = 2,  // ~x
    kNot = 3,     // not x
    kUnaryEnd = 4,

    // Binary operators
    kBinaryStart = 5,
    kAdd = 6,        // +
    kSub = 7,        // -
    kMult = 8,       // *
    kDiv = 9,        // /
    kFloorDiv = 10,  // // in Python
    kMod = 11,       // % in Python
    kPow = 12,       // ** in Python
    kLShift = 13,    // <<
    kRShift = 14,    // >>
    kBitAnd = 15,    // &
    kBitOr = 16,     // |
    kBitXor = 17,    // ^
    kLt = 18,        // <
    kLtE = 19,       // <=
    kEq = 20,        // ==
    kNotEq = 21,     // !=
    kGt = 22,        // >
    kGtE = 23,       // >=
    kAnd = 24,       // and
    kOr = 25,        // or
    kBinaryEnd = 26,

    // Special
    kSpecialStart = 27,
    kIfThenElse = 28,  // <operands[1]> if <operands[0]> else <operands[2]>
    kSpecialEnd = 29
  };

  Kind kind;
  Array<ExprDoc> operands;

  void VisitAttrs(AttrVisitor* v) {
    ExprDocNode::VisitAttrs(v);
    v->Visit("kind", &kind);
    v->Visit("operands", &operands);
  }

  static constexpr const char* _type_key = "script.printer.OperationDoc";
  TVM_DECLARE_FINAL_OBJECT_INFO(OperationDocNode, ExprDocNode);
};

class OperationDoc : public ExprDoc {
 public:
  explicit OperationDoc(OperationDocNode::Kind kind, Array<ExprDoc> operands);
  TVM_DEFINE_NOTNULLABLE_OBJECT_REF_METHODS(OperationDoc, ExprDoc, OperationDocNode);
};

class LambdaDocNode : public ExprDocNode {
 public:
  Array<IdDoc> args;
  ExprDoc body{nullptr};

  void VisitAttrs(AttrVisitor* v) {
    ExprDocNode::VisitAttrs(v);
    v->Visit("args", &args);
    v->Visit("body", &body);
  }

  static constexpr const char* _type_key = "script.printer.LambdaDoc";
  TVM_DECLARE_FINAL_OBJECT_INFO(LambdaDocNode, ExprDocNode);
};

class LambdaDoc : public ExprDoc {
 public:
  explicit LambdaDoc(Array<IdDoc> args, ExprDoc body);
  TVM_DEFINE_NOTNULLABLE_OBJECT_REF_METHODS(LambdaDoc, ExprDoc, LambdaDocNode);
};

class TupleDocNode : public ExprDocNode {
 public:
  Array<ExprDoc> elements;

  void VisitAttrs(AttrVisitor* v) {
    ExprDocNode::VisitAttrs(v);
    v->Visit("elements", &elements);
  }

  static constexpr const char* _type_key = "script.printer.TupleDoc";
  TVM_DECLARE_FINAL_OBJECT_INFO(TupleDocNode, ExprDocNode);
};

class TupleDoc : public ExprDoc {
 public:
  TupleDoc() : TupleDoc(runtime::make_object<TupleDocNode>()) {}
  explicit TupleDoc(Array<ExprDoc> elements);
  TVM_DEFINE_NOTNULLABLE_OBJECT_REF_METHODS(TupleDoc, ExprDoc, TupleDocNode);
};

class ListDocNode : public ExprDocNode {
 public:
  Array<ExprDoc> elements;

  void VisitAttrs(AttrVisitor* v) {
    ExprDocNode::VisitAttrs(v);
    v->Visit("elements", &elements);
  }

  static constexpr const char* _type_key = "script.printer.ListDoc";
  TVM_DECLARE_FINAL_OBJECT_INFO(ListDocNode, ExprDocNode);
};

class ListDoc : public ExprDoc {
 public:
  ListDoc() : ListDoc(runtime::make_object<ListDocNode>()) {}
  explicit ListDoc(Array<ExprDoc> elements);
  TVM_DEFINE_NOTNULLABLE_OBJECT_REF_METHODS(ListDoc, ExprDoc, ListDocNode);
};

class DictDocNode : public ExprDocNode {
 public:
  Array<ExprDoc> keys;
  Array<ExprDoc> values;

  void VisitAttrs(AttrVisitor* v) {
    ExprDocNode::VisitAttrs(v);
    v->Visit("keys", &keys);
    v->Visit("values", &values);
  }

  static constexpr const char* _type_key = "script.printer.DictDoc";
  TVM_DECLARE_FINAL_OBJECT_INFO(DictDocNode, ExprDocNode);
};

class DictDoc : public ExprDoc {
 public:
  DictDoc() : DictDoc(runtime::make_object<DictDocNode>()) {}
  explicit DictDoc(Array<ExprDoc> keys, Array<ExprDoc> values);
  TVM_DEFINE_NOTNULLABLE_OBJECT_REF_METHODS(DictDoc, ExprDoc, DictDocNode);
};

class SliceDocNode : public DocNode {
 public:
  Optional<ExprDoc> start;
  Optional<ExprDoc> stop;
  Optional<ExprDoc> step;

  void VisitAttrs(AttrVisitor* v) {
    DocNode::VisitAttrs(v);
    v->Visit("start", &start);
    v->Visit("stop", &stop);
    v->Visit("step", &step);
  }

  static constexpr const char* _type_key = "script.printer.SliceDoc";
  TVM_DECLARE_FINAL_OBJECT_INFO(SliceDocNode, DocNode);
};

class SliceDoc : public Doc {
 public:
  explicit SliceDoc(Optional<ExprDoc> start, Optional<ExprDoc> stop, Optional<ExprDoc> step);
  TVM_DEFINE_NOTNULLABLE_OBJECT_REF_METHODS(SliceDoc, Doc, SliceDocNode);
};

class AssignDocNode : public StmtDocNode {
 public:
  ExprDoc lhs{nullptr};
  Optional<ExprDoc> rhs;
  Optional<ExprDoc> annotation;

  void VisitAttrs(AttrVisitor* v) {
    StmtDocNode::VisitAttrs(v);
    v->Visit("lhs", &lhs);
    v->Visit("rhs", &rhs);
    v->Visit("annotation", &annotation);
  }

  static constexpr const char* _type_key = "script.printer.AssignDoc";
  TVM_DECLARE_FINAL_OBJECT_INFO(AssignDocNode, StmtDocNode);
};

class AssignDoc : public StmtDoc {
 public:
  explicit AssignDoc(ExprDoc lhs, Optional<ExprDoc> rhs, Optional<ExprDoc> annotation);
  TVM_DEFINE_NOTNULLABLE_OBJECT_REF_METHODS(AssignDoc, StmtDoc, AssignDocNode);
};

class IfDocNode : public StmtDocNode {
 public:
  ExprDoc predicate{nullptr};
  Array<StmtDoc> then_branch;
  Array<StmtDoc> else_branch;

  void VisitAttrs(AttrVisitor* v) {
    StmtDocNode::VisitAttrs(v);
    v->Visit("predicate", &predicate);
    v->Visit("then_branch", &then_branch);
    v->Visit("else_branch", &else_branch);
  }

  static constexpr const char* _type_key = "script.printer.IfDoc";
  TVM_DECLARE_FINAL_OBJECT_INFO(IfDocNode, StmtDocNode);
};

class IfDoc : public StmtDoc {
 public:
  explicit IfDoc(ExprDoc predicate, Array<StmtDoc> then_branch, Array<StmtDoc> else_branch);
  TVM_DEFINE_NOTNULLABLE_OBJECT_REF_METHODS(IfDoc, StmtDoc, IfDocNode);
};

class WhileDocNode : public StmtDocNode {
 public:
  ExprDoc predicate{nullptr};
  Array<StmtDoc> body;

  void VisitAttrs(AttrVisitor* v) {
    StmtDocNode::VisitAttrs(v);
    v->Visit("predicate", &predicate);
    v->Visit("body", &body);
  }

  static constexpr const char* _type_key = "script.printer.WhileDoc";
  TVM_DECLARE_FINAL_OBJECT_INFO(WhileDocNode, StmtDocNode);
};

class WhileDoc : public StmtDoc {
 public:
  explicit WhileDoc(ExprDoc predicate, Array<StmtDoc> body);
  TVM_DEFINE_NOTNULLABLE_OBJECT_REF_METHODS(WhileDoc, StmtDoc, WhileDocNode);
};

class ForDocNode : public StmtDocNode {
 public:
  ExprDoc lhs{nullptr};
  ExprDoc rhs{nullptr};
  Array<StmtDoc> body;

  void VisitAttrs(AttrVisitor* v) {
    StmtDocNode::VisitAttrs(v);
    v->Visit("lhs", &lhs);
    v->Visit("rhs", &rhs);
    v->Visit("body", &body);
  }

  static constexpr const char* _type_key = "script.printer.ForDoc";
  TVM_DECLARE_FINAL_OBJECT_INFO(ForDocNode, StmtDocNode);
};

class ForDoc : public StmtDoc {
 public:
  explicit ForDoc(ExprDoc lhs, ExprDoc rhs, Array<StmtDoc> body);
  TVM_DEFINE_NOTNULLABLE_OBJECT_REF_METHODS(ForDoc, StmtDoc, ForDocNode);
};

class ScopeDocNode : public StmtDocNode {
 public:
  Optional<ExprDoc> lhs{NullOpt};
  ExprDoc rhs{nullptr};
  Array<StmtDoc> body;

  void VisitAttrs(AttrVisitor* v) {
    StmtDocNode::VisitAttrs(v);
    v->Visit("lhs", &lhs);
    v->Visit("rhs", &rhs);
    v->Visit("body", &body);
  }

  static constexpr const char* _type_key = "script.printer.ScopeDoc";
  TVM_DECLARE_FINAL_OBJECT_INFO(ScopeDocNode, StmtDocNode);
};

class ScopeDoc : public StmtDoc {
 public:
  explicit ScopeDoc(Optional<ExprDoc> lhs, ExprDoc rhs, Array<StmtDoc> body);

  explicit ScopeDoc(ExprDoc rhs, Array<StmtDoc> body);

  TVM_DEFINE_NOTNULLABLE_OBJECT_REF_METHODS(ScopeDoc, StmtDoc, ScopeDocNode);
};

class ExprStmtDocNode : public StmtDocNode {
 public:
  ExprDoc expr{nullptr};

  void VisitAttrs(AttrVisitor* v) {
    StmtDocNode::VisitAttrs(v);
    v->Visit("expr", &expr);
  }

  static constexpr const char* _type_key = "script.printer.ExprStmtDoc";
  TVM_DECLARE_FINAL_OBJECT_INFO(ExprStmtDocNode, StmtDocNode);
};

class ExprStmtDoc : public StmtDoc {
 public:
  explicit ExprStmtDoc(ExprDoc expr);
  TVM_DEFINE_NOTNULLABLE_OBJECT_REF_METHODS(ExprStmtDoc, StmtDoc, ExprStmtDocNode);
};

class AssertDocNode : public StmtDocNode {
 public:
  ExprDoc test{nullptr};
  Optional<ExprDoc> msg{NullOpt};

  void VisitAttrs(AttrVisitor* v) {
    StmtDocNode::VisitAttrs(v);
    v->Visit("test", &test);
    v->Visit("msg", &msg);
  }

  static constexpr const char* _type_key = "script.printer.AssertDoc";
  TVM_DECLARE_FINAL_OBJECT_INFO(AssertDocNode, StmtDocNode);
};

class AssertDoc : public StmtDoc {
 public:
  explicit AssertDoc(ExprDoc test, Optional<ExprDoc> msg = NullOpt);
  TVM_DEFINE_NOTNULLABLE_OBJECT_REF_METHODS(AssertDoc, StmtDoc, AssertDocNode);
};

class ReturnDocNode : public StmtDocNode {
 public:
  ExprDoc value{nullptr};

  void VisitAttrs(AttrVisitor* v) {
    StmtDocNode::VisitAttrs(v);
    v->Visit("value", &value);
  }

  static constexpr const char* _type_key = "script.printer.ReturnDoc";
  TVM_DECLARE_FINAL_OBJECT_INFO(ReturnDocNode, StmtDocNode);
};

class ReturnDoc : public StmtDoc {
 public:
  explicit ReturnDoc(ExprDoc value);
  TVM_DEFINE_NOTNULLABLE_OBJECT_REF_METHODS(ReturnDoc, StmtDoc, ReturnDocNode);
};

class FunctionDocNode : public StmtDocNode {
 public:
  IdDoc name{nullptr};
  Array<AssignDoc> args;
  Array<ExprDoc> decorators;
  Optional<ExprDoc> return_type{NullOpt};
  Array<StmtDoc> body;

  void VisitAttrs(AttrVisitor* v) {
    StmtDocNode::VisitAttrs(v);
    v->Visit("name", &name);
    v->Visit("args", &args);
    v->Visit("decorators", &decorators);
    v->Visit("return_type", &return_type);
    v->Visit("body", &body);
  }

  static constexpr const char* _type_key = "script.printer.FunctionDoc";
  TVM_DECLARE_FINAL_OBJECT_INFO(FunctionDocNode, StmtDocNode);
};

class FunctionDoc : public StmtDoc {
 public:
  explicit FunctionDoc(IdDoc name, Array<AssignDoc> args, Array<ExprDoc> decorators,
                       Optional<ExprDoc> return_type, Array<StmtDoc> body);
  TVM_DEFINE_NOTNULLABLE_OBJECT_REF_METHODS(FunctionDoc, StmtDoc, FunctionDocNode);
};

class ClassDocNode : public StmtDocNode {
 public:
  IdDoc name{nullptr};
  Array<ExprDoc> decorators;
  Array<StmtDoc> body;

  void VisitAttrs(AttrVisitor* v) {
    StmtDocNode::VisitAttrs(v);
    v->Visit("name", &name);
    v->Visit("decorators", &decorators);
    v->Visit("body", &body);
  }

  static constexpr const char* _type_key = "script.printer.ClassDoc";
  TVM_DECLARE_FINAL_OBJECT_INFO(ClassDocNode, StmtDocNode);
};

class ClassDoc : public StmtDoc {
 public:
  explicit ClassDoc(IdDoc name, Array<ExprDoc> decorators, Array<StmtDoc> body);
  TVM_DEFINE_NOTNULLABLE_OBJECT_REF_METHODS(ClassDoc, StmtDoc, ClassDocNode);
};

}  // namespace printer
}  // namespace script
}  // namespace tvm

#endif  // TVM_SCRIPT_PRINTER_DOC_H_