Support a large number of BBs

everx-labs · Mar 31, 2020 · 0f62ed2 · 0f62ed2
1 parent c7feb0d
commit 0f62ed2
Show file tree

Hide file tree

Showing 3 changed files with 36 additions and 7 deletions.
diff --git a/llvm/lib/Target/TVM/TVMAsmPrinter.cpp b/llvm/lib/Target/TVM/TVMAsmPrinter.cpp
@@ -176,7 +176,13 @@ void TVMAsmPrinter::EmitFunctionBodyEnd() {
     }
   }
 
-  OutStreamer->EmitRawText("  PUSH s" + Twine(Blocks + Arguments - 1));
+  unsigned Entry = Blocks + Arguments - 1;
+  if (Entry < 256) {
+    OutStreamer->EmitRawText("  PUSH s" + Twine(Entry));
+  } else {
+    OutStreamer->EmitRawText("  PUSHINT " + Twine(Entry));
+    OutStreamer->EmitRawText("  PUSHX");
+  }
   OutStreamer->EmitRawText("  EXECUTE");
 
   if (ReturnValues > 0) {

diff --git a/llvm/lib/Target/TVM/TVMInstrInfo.td b/llvm/lib/Target/TVM/TVMInstrInfo.td
@@ -198,6 +198,7 @@ defm NOP : I<(outs), (ins), (outs), (ins), [(int_tvm_nop)], "NOP", "NOP", 0x00>;
 
 let mayLoad = 1, isAsCheapAsAMove = 1 in {
 defm PUSH : SI<(ins stack_op:$local), "PUSH\t$local", 0x20>;
+defm PUSHX : SI<(ins), "PUSHX", 0x60>;
 defm SWAP : SI<(ins), "SWAP", 0x01>;
 defm XCHG_TOP : SI<(ins stack_op:$src), "XCHG\ts0, $src", 0x01>;
 defm XCHG_TOP_DEEP : SI<(ins stack_op:$src), "XCHG\ts0, $src", 0x11>;

diff --git a/llvm/lib/Target/TVM/TVMStackModel.cpp b/llvm/lib/Target/TVM/TVMStackModel.cpp
@@ -651,19 +651,41 @@ void TVMStackModel::rewriteToSForm(MachineInstr &MI,
         MFI->addStackModelComment(MIB.getInstr(), Then->getName());
         MFI->addStackModelComment(MIB.getInstr(), Else->getName());
       } else {
-        MIB = BuildMI(*MI.getParent(), MI, MI.getDebugLoc(),
-            TII->get(TVM::PUSH)).addImm(ThenID);
+        if (ThenID < 256) {
+          MIB = BuildMI(*MI.getParent(), MI, MI.getDebugLoc(),
+              TII->get(TVM::PUSH)).addImm(ThenID);
+        } else {
+          BuildMI(*MI.getParent(), MI, MI.getDebugLoc(),
+              TII->get(TVM::CONST_I257)).addImm(ThenID);
+          MIB = BuildMI(*MI.getParent(), MI, MI.getDebugLoc(),
+              TII->get(TVM::PUSHX));
+        }
         MFI->addStackModelComment(MIB.getInstr(), Then->getName());
-        MIB = BuildMI(*MI.getParent(), MI, MI.getDebugLoc(),
-            TII->get(TVM::PUSH)).addImm(ElseID + 1);
+        if (ElseID < 255) {
+          MIB = BuildMI(*MI.getParent(), MI, MI.getDebugLoc(),
+              TII->get(TVM::PUSH)).addImm(ElseID + 1);
+        } else {
+          BuildMI(*MI.getParent(), MI, MI.getDebugLoc(),
+              TII->get(TVM::CONST_I257)).addImm(ElseID + 1);
+          MIB = BuildMI(*MI.getParent(), MI, MI.getDebugLoc(),
+              TII->get(TVM::PUSHX));
+        }
         MFI->addStackModelComment(MIB.getInstr(), Else->getName());
       }
     } else if (MI.getOpcode() == TVM::JMPX) {
       auto Dest = MI.getOperand(0).getMBB();
       unsigned DestID = TheStack.size() + BBInfo[Dest].getID();
 
-      auto MIB = BuildMI(*MI.getParent(), MI, MI.getDebugLoc(),
-          TII->get(TVM::PUSH)).addImm(DestID);
+      MachineInstrBuilder MIB;
+      if (DestID < 256) {
+        MIB = BuildMI(*MI.getParent(), MI, MI.getDebugLoc(),
+            TII->get(TVM::PUSH)).addImm(DestID);
+      } else {
+        BuildMI(*MI.getParent(), MI, MI.getDebugLoc(),
+            TII->get(TVM::CONST_I257)).addImm(DestID);
+        MIB = BuildMI(*MI.getParent(), MI, MI.getDebugLoc(),
+            TII->get(TVM::PUSHX));
+      }
       MFI->addStackModelComment(MIB.getInstr(), Dest->getName());
     }