AndrewZhaoLuo commented on code in PR #12580:
URL: https://github.com/apache/tvm/pull/12580#discussion_r955387959


##########
src/relay/transforms/simplify_expr.cc:
##########
@@ -708,6 +708,50 @@ class SimplifyRSqrt : public DFPatternRewrite {
   DFPattern numerator_;
 };
 
+class SimplifyDQArgFunc : public DFPatternRewrite {
+  /*! Base class for simplifying dequantize followed by arg ops */
+ public:
+  explicit SimplifyDQArgFunc(std::string op) : op_(op) {
+    x_ = IsWildcard();
+    dq_ = IsOp("qnn.dequantize")({x_, IsWildcard(), IsWildcard()});
+    pattern_ = IsOp(op_)({dq_});
+  }
+
+  Expr Callback(const Expr& pre, const Expr& post,
+                const Map<DFPattern, Array<Expr>>& node_map) const override {
+    const CallNode* call = pre.as<CallNode>();
+    ICHECK(call);
+    auto x = node_map[x_][0];
+    return Call(Op::Get(op_), {x}, call->attrs);
+  }
+
+ protected:
+  /*! \brief Pattern input */
+  DFPattern x_;
+  /*! \brief dequantize op */
+  DFPattern dq_;
+  /*! \brief Name of op to simplify */
+  String op_;
+};
+
+class SimplifyDQArgMax : public SimplifyDQArgFunc {
+  /*! Simplify dequantize follwed by argmax */

Review Comment:
   Probably want to follow format used in rest of file:
   
   ```
   
   /*! \brief Make two consecutive add able to be constant_folded.
    * This pattern matching supports commutative property for addition.
    */
   class SimplifyConsecutiveAdd : public DFPatternRewrite {
   ```



-- 
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: [email protected]

For queries about this service, please contact Infrastructure at:
[email protected]

Reply via email to