[Heavy] Add multiple arguments to nbdl_gen::VisitOp

Author: ricejasonf

heavy/include/nbdl_gen/Nbdl.td

@@ -141,13 +141,12 @@ def Nbdl_GetOp : Nbdl_Op<"get", []> {
   let results   = (outs Nbdl_Type:$result);
 }
 
-def VisitOp : Nbdl_Op<"visit", [Terminator]> {
+def Nbdl_VisitOp : Nbdl_Op<"visit", [Terminator]> {
   let description = [{
-    Visit a single result with a function object.
-    VisitOp is a simplification of MatchOp for leaf nodes.
+    Visit matched results with a function object.
   }];
 
-  let arguments = (ins Nbdl_Type:$fn, Nbdl_Type:$arg);
+  let arguments = (ins Nbdl_Type:$fn, Variadic<Nbdl_Type>:$args);
   let results   = (outs);
 }
 

heavy/lib/Nbdl/NbdlWriter.cpp

@@ -351,7 +351,10 @@ class FuncWriter : public NbdlWriter<FuncWriter> {
   void Visit(VisitOp Op) {
     WriteForwardedExpr(Op.getFn());
     OS << '(';
-    WriteForwardedExpr(Op.getArg());
+    llvm::interleave(Op.getArgs(), OS,
+        [&](mlir::Value V) {
+          WriteForwardedExpr(V);
+        }, ",\n");
     OS << ");\n";
   }