From b46c56a9c78292b973eeb9fb8622649f07cb5ef2 Mon Sep 17 00:00:00 2001
From: Dmitry Borisenkov <dmitry.borisenkov89@gmail.com>
Date: Fri, 20 Nov 2020 01:16:09 +0200
Subject: [PATCH] Add PUSHPOW2, PUSHPOW2DEC, PUSHNEGPOW2

---
 llvm/lib/Target/TVM/TVMInstrInfo.td         | 12 ++++++
 llvm/lib/Target/TVM/TVMPeephole.cpp         | 41 +++++++++++++++++++++
 llvm/test/CodeGen/TVM/optimizations/pow2.ll | 23 ++++++++++++
 3 files changed, 76 insertions(+)
 create mode 100644 llvm/test/CodeGen/TVM/optimizations/pow2.ll

diff --git a/llvm/lib/Target/TVM/TVMInstrInfo.td b/llvm/lib/Target/TVM/TVMInstrInfo.td
index 0a65e09213d4..7421a16c6513 100644
--- a/llvm/lib/Target/TVM/TVMInstrInfo.td
+++ b/llvm/lib/Target/TVM/TVMInstrInfo.td
@@ -262,6 +262,18 @@ defm CONST_U257 : I<(outs I257:$res), (ins UImm257:$uimm),
                     (outs), (ins UImm257:$uimm),
                     [(set I257:$res, (TVMconst_u257 timm:$uimm))],
                     "PUSHINT\t$res, $uimm", "PUSHINT\t$uimm", 0x82>;
+defm CONST_POW2 : I<(outs I257:$res), (ins i257imm:$imm),
+                    (outs), (ins i257imm:$imm),
+                    [(set I257:$res, imm:$imm)],
+                    "PUSHPOW2\t$res, $imm", "PUSHPOW2\t$imm", 0x83>;
+defm CONST_POW2DEC : I<(outs I257:$res), (ins i257imm:$imm),
+                      (outs), (ins i257imm:$imm),
+                      [(set I257:$res, imm:$imm)],
+                       "PUSHPOW2DEC\t$res, $imm", "PUSHPOW2DEC\t$imm", 0x84>;
+defm CONST_NEGPOW2 : I<(outs I257:$res), (ins i257imm:$imm),
+                       (outs), (ins i257imm:$imm),
+                       [(set I257:$res, imm:$imm)],
+                       "PUSHNEGPOW2\t$res, $imm", "PUSHNEGPOW2\t$imm", 0x85>;
 }
 
 
diff --git a/llvm/lib/Target/TVM/TVMPeephole.cpp b/llvm/lib/Target/TVM/TVMPeephole.cpp
index 1522aa59c5fc..258318996a4b 100644
--- a/llvm/lib/Target/TVM/TVMPeephole.cpp
+++ b/llvm/lib/Target/TVM/TVMPeephole.cpp
@@ -41,6 +41,11 @@ static cl::opt<bool>
 static cl::opt<bool> DisableTVMBlkPushOpt(
     "disable-tvm-push-opt", cl::Hidden,
     cl::desc("TVM: Disable BLKPUSH peephole optimizations."), cl::init(false));
+static cl::opt<bool> DisableTVMPushintOpt(
+    "disable-tvm-pushint-opt", cl::Hidden,
+    cl::desc("TVM: Disable power of 2 optimizations for constants."),
+    cl::init(false));
+
 static unsigned BlkPushLimit = 15;
 namespace {
 class TVMPeephole final : public MachineFunctionPass {
@@ -58,6 +63,7 @@ class TVMPeephole final : public MachineFunctionPass {
                                      const TargetInstrInfo &TII);
   bool runBlkDropCombine(MachineBasicBlock &MBB, const TargetInstrInfo &TII);
   bool runBlkPushCombine(MachineBasicBlock &MBB, const TargetInstrInfo &TII);
+  bool runPushintOpt(MachineBasicBlock &MBB, const TargetInstrInfo &TII);
   bool runMbbInlineOptimization(MachineBasicBlock &MBB,
                                 const TargetInstrInfo &TII);
   bool runIfElseOptimization(MachineBasicBlock &MBB,
@@ -183,6 +189,39 @@ bool TVMPeephole::runBlkPushCombine(MachineBasicBlock &MBB,
   return Changed;
 }
 
+bool TVMPeephole::runPushintOpt(MachineBasicBlock &MBB,
+                                const TargetInstrInfo &TII) {
+  std::vector<MachineInstr *> MIRemove;
+  for (auto &I : MBB) {
+    if (I.getOpcode() != TVM::CONST_I257_S && I.getOpcode() != TVM::CONST_U257_S)
+      continue;
+    APInt Val = I.getOperand(0).getCImm()->getValue();
+    if (Val.isPowerOf2() && Val.ugt(64)) {
+      MIRemove.push_back(&I);
+      BuildMI(&I, TII.get(TVM::CONST_POW2_S))
+        .addImm(Val.getBitWidth() - Val.countLeadingZeros() - 1);
+      continue;
+    }
+    APInt AbsVal = Val.abs();
+    if (AbsVal.isPowerOf2() && AbsVal.ugt(64)) {
+      MIRemove.push_back(&I);
+      BuildMI(&I, TII.get(TVM::CONST_NEGPOW2_S))
+        .addImm(AbsVal.getBitWidth() - AbsVal.countLeadingZeros() - 1);
+      continue;
+    }
+    Val += 1;
+    if (Val.isPowerOf2() && Val.ugt(64)) {
+      MIRemove.push_back(&I);
+      BuildMI(&I, TII.get(TVM::CONST_POW2DEC_S))
+        .addImm(Val.getBitWidth() - Val.countLeadingZeros() - 1);
+      continue;
+    }
+  }
+  for (auto *I : MIRemove)
+    I->eraseFromParent();
+  return !MIRemove.empty();
+}
+
 bool TVMPeephole::runImplicitReturnOptimization(MachineBasicBlock &MBB,
                                                 const TargetInstrInfo &TII) {
   if (MBB.empty())
@@ -563,6 +602,8 @@ bool TVMPeephole::runOnMachineBasicBlock(MachineBasicBlock &MBB,
     Changed |= runBlkDropCombine(MBB, TII);
   if (!DisableTVMBlkPushOpt)
     Changed |= runBlkPushCombine(MBB, TII);
+  if (!DisableTVMPushintOpt)
+    Changed |= runPushintOpt(MBB, TII);
 
   return Changed;
 }
diff --git a/llvm/test/CodeGen/TVM/optimizations/pow2.ll b/llvm/test/CodeGen/TVM/optimizations/pow2.ll
new file mode 100644
index 000000000000..c4774e5b7c49
--- /dev/null
+++ b/llvm/test/CodeGen/TVM/optimizations/pow2.ll
@@ -0,0 +1,23 @@
+; RUN: llc -O3 < %s -march=tvm -filetype=asm | FileCheck %s
+target datalayout = "E-S257-i1:257:257-i8:257:257-i16:257:257-i32:257:257-i64:257:257-i257:257:257-p:257:257-a:257:257"
+target triple = "tvm"
+
+; CHECK-LABEL: ret1
+define i257 @ret1() {
+  ret i257 128
+}
+
+; CHECK-LABEL: ret2
+define i257 @ret2() {
+  ret i257 -256
+}
+
+; CHECK-LABEL: ret3
+define i257 @ret3() {
+  ret i257 255
+}
+
+; CHECK-LABEL: retneg
+define i257 @retneg() {
+  ret i257 0
+}