junrushao1994 commented on code in PR #12148:
URL: https://github.com/apache/tvm/pull/12148#discussion_r934086940
##########
src/script/printer/python_doc_printer.cc:
##########
@@ -31,6 +31,114 @@ namespace tvm {
namespace script {
namespace printer {
+/*!
+ * \brief Operator precedence
+ *
+ * This is based on
+ * https://docs.python.org/3/reference/expressions.html#operator-precedence
+ */
+enum class ExprPrecedence : int32_t {
+ /*! \brief Unknown precedence */
+ kUnkown = 0,
+ /*! \brief Lambda Expression */
+ kLambda = 1,
+ /*! \brief Conditional Expression */
+ kIfThenElse = 2,
+ /*! \brief Boolean OR */
+ kBooleanOr = 3,
+ /*! \brief Boolean AND */
+ kBooleanAnd = 4,
+ /*! \brief Boolean NOT */
+ kBooleanNot = 5,
+ /*! \brief Comparisons */
+ kComparison = 6,
+ /*! \brief Bitwise OR */
+ kBitwiseOr = 7,
+ /*! \brief Bitwise XOR */
+ kBitwiseXor = 8,
+ /*! \brief Bitwise AND */
+ kBitwiseAnd = 9,
+ /*! \brief Shift Operators */
+ kShift = 10,
+ /*! \brief Addition and subtraction */
+ kAdd = 11,
+ /*! \brief Multiplication, division, floor division, remainder */
+ kMult = 12,
+ /*! \brief Positive negative and bitwise NOT */
+ kUnary = 13,
+ /*! \brief Exponentiation */
+ kExp = 14,
+ /*! \brief Index access, attribute access, call and atom expression */
+ kIdentity = 15,
+};
+
+#define DOC_PRECEDENCE_ENTRY(RefType, Precedence) \
+ { RefType::ContainerType::RuntimeTypeIndex(), ExprPrecedence::Precedence }
+
+ExprPrecedence GetExprPrecedence(const ExprDoc& doc) {
+ // Key is the value of OperationDocNode::Kind
+ static const std::vector<ExprPrecedence> op_kind_precedence = []() {
+ using OpKind = OperationDocNode::Kind;
+ std::map<OpKind, ExprPrecedence> raw_table = {
+ {OpKind::kUSub, ExprPrecedence::kUnary}, //
+ {OpKind::kInvert, ExprPrecedence::kUnary}, //
+ {OpKind::kAdd, ExprPrecedence::kAdd}, //
+ {OpKind::kSub, ExprPrecedence::kAdd}, //
+ {OpKind::kMult, ExprPrecedence::kMult}, //
+ {OpKind::kDiv, ExprPrecedence::kMult}, //
+ {OpKind::kFloorDiv, ExprPrecedence::kMult}, //
+ {OpKind::kMod, ExprPrecedence::kMult}, //
+ {OpKind::kPow, ExprPrecedence::kExp}, //
+ {OpKind::kLShift, ExprPrecedence::kShift}, //
+ {OpKind::kRShift, ExprPrecedence::kShift}, //
+ {OpKind::kBitAnd, ExprPrecedence::kBitwiseAnd}, //
+ {OpKind::kBitOr, ExprPrecedence::kBitwiseOr}, //
+ {OpKind::kBitXor, ExprPrecedence::kBitwiseXor}, //
+ {OpKind::kLt, ExprPrecedence::kComparison}, //
+ {OpKind::kLtE, ExprPrecedence::kComparison}, //
+ {OpKind::kEq, ExprPrecedence::kComparison}, //
+ {OpKind::kNotEq, ExprPrecedence::kComparison}, //
+ {OpKind::kGt, ExprPrecedence::kComparison}, //
+ {OpKind::kGtE, ExprPrecedence::kComparison}, //
+ {OpKind::kIfThenElse, ExprPrecedence::kIfThenElse}, //
+ };
+
+ std::vector<ExprPrecedence> table;
+ table.resize(static_cast<int>(OperationDocNode::Kind::kSpecialEnd) + 1);
+
+ for (const auto& kv : raw_table) {
+ table[static_cast<int>(kv.first)] = kv.second;
+ }
+
+ return table;
+ }();
+
+ // Key is the type index of Doc
+ static const std::unordered_map<uint32_t, ExprPrecedence>
doc_type_precedence = {
+ DOC_PRECEDENCE_ENTRY(LiteralDoc, kIdentity), //
+ DOC_PRECEDENCE_ENTRY(IdDoc, kIdentity), //
+ DOC_PRECEDENCE_ENTRY(AttrAccessDoc, kIdentity), //
+ DOC_PRECEDENCE_ENTRY(IndexDoc, kIdentity), //
+ DOC_PRECEDENCE_ENTRY(CallDoc, kIdentity), //
+ DOC_PRECEDENCE_ENTRY(LambdaDoc, kLambda), //
+ DOC_PRECEDENCE_ENTRY(TupleDoc, kIdentity), //
+ DOC_PRECEDENCE_ENTRY(ListDoc, kIdentity), //
+ DOC_PRECEDENCE_ENTRY(DictDoc, kIdentity), //
+ };
+
+ if (const auto* op_doc = doc.as<OperationDocNode>()) {
+ ExprPrecedence precedence =
op_kind_precedence[static_cast<int>(op_doc->kind)];
+ ICHECK(precedence != ExprPrecedence::kUnkown)
+ << "Precedence for operator " << static_cast<int>(op_doc->kind) << "
is unknown";
+ return precedence;
+ } else if (doc_type_precedence.find(doc->type_index()) !=
doc_type_precedence.end()) {
+ return doc_type_precedence.at(doc->type_index());
Review Comment:
nit: we may save one look-up overhead
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscr...@tvm.apache.org
For queries about this service, please contact Infrastructure at:
us...@infra.apache.org
