Inline stslice if possible

If stslice store a slice which is made of a cell that has one use and
was entiraly made of stores via a builder, the pass replaces the slice
with correspondin strore instructions. The pass relies on DCE to cleanup
unused values.

llvm/lib/Target/TVM/CMakeLists.txt

@@ -44,6 +44,7 @@ add_llvm_target(TVMCodeGen
   TVMLoadStoreReplace.cpp
   TVMMoveMaterializable.cpp
   TVMIfConversionTerm.cpp
+  TVMInlineSliceStore.cpp
   TVMLowerIntrinsics.cpp
 )
 

llvm/lib/Target/TVM/TVM.h

@@ -44,6 +44,7 @@ FunctionPass *createTVMIfConversionTerm();
 BasicBlockPass *createTVMDefineUndef();
 BasicBlockPass *createTVMLoadStoreReplace();
 BasicBlockPass *createTVMStoreCombine();
+BasicBlockPass *createTVMInlineSliceStore();
 ModulePass *createTVMLowerIntrinsicsPass();
 
 void initializeTVMArgumentMovePass(PassRegistry &);
@@ -63,6 +64,7 @@ void initializeTVMLoadStoreReplacePass(PassRegistry &);
 void initializeTVMMoveMaterializablePass(PassRegistry &);
 void initializeTVMIfConversionTermPass(PassRegistry &);
 void initializeTVMStoreCombinePass(PassRegistry &);
+void initializeTVMInlineSliceStorePass(PassRegistry &);
 void initializeTVMLowerIntrinsicsPass(PassRegistry &);
 
 } // namespace llvm

llvm/lib/Target/TVM/TVMInlineSliceStore.cpp

@@ -0,0 +1,175 @@
+//===------------- TVMInlineSliceStore.cpp - Inline slice store -----------===//
+//
+//                     The LLVM Compiler Infrastructure
+//
+// This file is distributed under the University of Illinois Open Source
+// License. See LICENSE.TXT for details.
+//
+//===----------------------------------------------------------------------===//
+///
+/// \file
+/// Replace store slice with its content.
+///
+//===----------------------------------------------------------------------===//
+
+#include <vector>
+
+#include "TVM.h"
+#include "llvm/Analysis/LoopInfo.h"
+#include "llvm/Analysis/LoopPass.h"
+#include "llvm/IR/Constants.h"
+#include "llvm/IR/IRBuilder.h"
+#include "llvm/IR/Module.h"
+#include "llvm/IR/Operator.h"
+#include "llvm/Pass.h"
+#include "llvm/Support/Debug.h"
+#include "llvm/Transforms/Utils/Local.h"
+#include "llvm/Transforms/Utils/ModuleUtils.h"
+using namespace llvm;
+
+#define DEBUG_TYPE "tvm-inline-slice-store"
+
+namespace {
+class TVMInlineSliceStore final : public BasicBlockPass {
+  StringRef getPassName() const override {
+    return "Replace store slice with its content";
+  }
+
+  void getAnalysisUsage(AnalysisUsage &AU) const override {
+    AU.setPreservesAll();
+  }
+
+  bool runOnBasicBlock(BasicBlock &BB) override;
+
+public:
+  static char ID;
+  explicit TVMInlineSliceStore() : BasicBlockPass(ID) {}
+};
+} // End anonymous namespace
+
+char TVMInlineSliceStore::ID = 0;
+INITIALIZE_PASS(TVMInlineSliceStore, DEBUG_TYPE, "Inline store slice",
+                false, false)
+
+BasicBlockPass *llvm::createTVMInlineSliceStore() { return new TVMInlineSliceStore(); }
+
+static CallSite storeCallSite(Instruction *I) {
+  auto CS = CallSite(I);
+  if (!CS)
+    return CS;
+  if (CS.getIntrinsicID() == Intrinsic::tvm_stslice)
+    return CS;
+  return {};
+}
+
+static bool isIntStore(unsigned ID) {
+  return ID == Intrinsic::tvm_sti || ID == Intrinsic::tvm_stu;
+}
+
+static bool isVarIntStore(unsigned ID) {
+  return ID == Intrinsic::tvm_stvarint16 || ID == Intrinsic::tvm_stvaruint16
+    ||   ID == Intrinsic::tvm_stvarint32 || ID == Intrinsic::tvm_stvaruint32;
+}
+
+/// For a given store \p CS return its builder argument.
+static Value* argBuilder(CallSite CS) {
+  if (!CS)
+    return nullptr;
+  unsigned ID = CS.getIntrinsicID();
+  if (isIntStore(ID))
+    return CS.getArgument(1);
+  if (isVarIntStore(ID))
+    return CS.getArgument(0);
+  return nullptr;
+}
+
+/// Return stores to a slice in reverse order
+static std::vector<Value *> collectStores(Value *Slice) {
+  assert(Slice);
+  std::vector<Value *> Values;
+  if (!Slice->hasOneUse())
+    return {};
+  auto CS = CallSite(Slice);
+  if (!CS || CS.getIntrinsicID() != Intrinsic::tvm_ctos)
+    return {};
+  Value *Cell = CS.getArgument(0);
+  if (!Cell->hasOneUse())
+    return {};
+  CS = CallSite(Cell);
+  if (!CS || CS.getIntrinsicID() != Intrinsic::tvm_endc)
+    return {};
+  Value *Builder = CS.getArgument(0);
+  CS = CallSite(Builder);
+  while (Builder) {
+    if (!Builder->hasOneUse())
+      break;
+    auto *NextBuilder = argBuilder(CS);
+    if (!NextBuilder)
+      break;
+    Values.push_back(Builder);
+    Builder = NextBuilder;
+    CS = CallSite(Builder);
+  }
+  if (!CS || (CS.getIntrinsicID() != Intrinsic::tvm_newc))
+    return {};
+  return Values;
+}
+
+/// Provide arguments for \p Store replacement.
+/// The argument are the same as for original Store but builder is substituted
+/// with \p Builder.
+static std::vector<Value *> fillArguments(Value *Store, Value *Builder) {
+  auto CS = CallSite(Store);
+  assert(CS);
+  unsigned ID = CS.getIntrinsicID();
+  if (isIntStore(ID))
+    return {CS.getArgument(0), Builder, CS.getArgument(2)};
+  if (isVarIntStore(ID))
+    return {Builder, CS.getArgument(1)};
+  llvm_unreachable("Unexpected Store");
+  return {};
+}
+
+/// Inline \p Stores and return the address of the last builder
+/// to replace the original one.
+static Value *inlineStores(Instruction *I,
+                           const std::vector<Value *> &Stores) {
+  assert(!Stores.empty());
+  IRBuilder<> Builder(I);
+  auto CS = CallSite(I);
+  Value *Prev = CS.getArgument(1);
+  for (Value *Store : make_range(Stores.rbegin(), Stores.rend())) {
+    CS = CallSite(Store);
+    assert(CS);
+    std::vector<Value *> Args = fillArguments(Store, Prev);
+    auto *Fn = Intrinsic::getDeclaration(I->getModule(), CS.getIntrinsicID());
+    Prev = Builder.CreateCall(Fn, Args);
+    llvm::outs() << *Prev << "\n";
+  }
+  return Prev;
+}
+
+bool TVMInlineSliceStore::runOnBasicBlock(BasicBlock &BB) {
+  std::vector<Value *> RemoveInst;
+  auto It = BB.begin(), End = BB.end();
+  while (It != End) {
+    It = std::find_if(It, End, [](Instruction &I) {
+      return static_cast<bool>(I.hasOneUse() && storeCallSite(&I));
+    });
+    if (It != End) {
+      auto CS = storeCallSite(&*It);
+      Value *Slice = CS.getArgument(0);
+      std::vector<Value *> Stores = collectStores(Slice);
+      if (!Stores.empty()) {
+        Value *LastStore = inlineStores(&*It, Stores);
+        llvm::outs() << "repl " << *It << *LastStore << "\n"; 
+        It->replaceAllUsesWith(LastStore);
+        RemoveInst.push_back(&*It);
+      }
+    } else {
+      break;
+    }
+    ++It;
+  }
+  return !RemoveInst.empty();
+}

llvm/lib/Target/TVM/TVMTargetMachine.cpp

@@ -45,6 +45,7 @@ extern "C" void LLVMInitializeTVMTarget() {
   initializeTVMLoadStoreReplacePass(PR);
   initializeTVMMoveMaterializablePass(PR);
   initializeTVMStoreCombinePass(PR);
+  initializeTVMInlineSliceStorePass(PR);
   initializeTVMLowerIntrinsicsPass(PR);
 }
 
@@ -138,6 +139,7 @@ void TVMPassConfig::addIRPasses() {
   addPass(createTVMLoopPrepare());
   addPass(createTVMControlFlowPrepare());
   addPass(createTVMDefineUndef());
+  addPass(createTVMInlineSliceStore());
   addPass(createTVMStoreCombine());
 }
 
@@ -171,6 +173,7 @@ void TVMPassConfig::addPreEmitPass() {
   addPass(createTVMRegStackify());
   addPass(createTVMLoopInstructions());
   addPass(createTVMMoveMaterializable());
+  addPass(createGlobalDCEPass());
   addPass(createTVMStackModel());
 
   // Perform the very last peephole optimizations on the code.

llvm/test/CodeGen/TVM/optimizations/stslice_inline.ll

@@ -0,0 +1,74 @@
+; RUN: llc < %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: stslice
+define builder @stslice() {
+; CHECK: PUSHINT 5
+; CHECK: STUR
+; CHECK-NOT: STSLICE
+  %b1 = tail call builder @llvm.tvm.newc()
+  %b2 = tail call builder @llvm.tvm.stu(i257 0, builder %b1, i257 2)
+  %b3 = tail call builder @llvm.tvm.sti(i257 -1, builder %b2, i257 2)
+  %c1 = tail call cell @llvm.tvm.endc(builder %b3)
+  %s1 = tail call slice @llvm.tvm.ctos(cell %c1)
+  %b4 = tail call builder @llvm.tvm.newc()
+  %b5 = tail call builder @llvm.tvm.stslice(slice %s1, builder %b4)
+  ret builder %b5
+}
+
+; CHECK-LABEL: stslice_varint
+define builder @stslice_varint() {
+; CHECK-NOT: STSLICE
+  %b1 = tail call builder @llvm.tvm.newc()
+  %b2 = tail call builder @llvm.tvm.stvarint16(builder %b1, i257 2)
+  %b3 = tail call builder @llvm.tvm.stvaruint32(builder %b2, i257 2)
+  %b4 = tail call builder @llvm.tvm.sti(i257 -1, builder %b3, i257 2)
+  %c1 = tail call cell @llvm.tvm.endc(builder %b4)
+  %s1 = tail call slice @llvm.tvm.ctos(cell %c1)
+  %b5 = tail call builder @llvm.tvm.newc()
+  %b6 = tail call builder @llvm.tvm.stslice(slice %s1, builder %b5)
+  ret builder %b6
+}
+
+; CHECK-LABEL: stslice_neg
+define builder @stslice_neg() {
+; CHECK: STSLICE
+; CHECK: STSLICE
+  %b1 = tail call builder @llvm.tvm.newc()
+  %b2 = tail call builder @llvm.tvm.stu(i257 0, builder %b1, i257 2)
+  %b3 = tail call builder @llvm.tvm.sti(i257 -1, builder %b2, i257 2)
+  %c1 = tail call cell @llvm.tvm.endc(builder %b3)
+  %s1 = tail call slice @llvm.tvm.ctos(cell %c1)
+  %b4 = tail call builder @llvm.tvm.newc()
+  %b5 = tail call builder @llvm.tvm.stslice(slice %s1, builder %b4)
+  %b6 = tail call builder @llvm.tvm.stslice(slice %s1, builder %b5)
+  ret builder %b6
+}
+
+; CHECK-LABEL: stslice_rec
+define builder @stslice_rec() {
+; CHECK: PUSHINT -1
+; CHECK: STIR
+; CHECK-NOT: STSLICE
+  %b1 = tail call builder @llvm.tvm.newc()
+  %b2 = tail call builder @llvm.tvm.stu(i257 0, builder %b1, i257 2)
+  %b3 = tail call builder @llvm.tvm.sti(i257 -1, builder %b2, i257 2)
+  %c1 = tail call cell @llvm.tvm.endc(builder %b3)
+  %s1 = tail call slice @llvm.tvm.ctos(cell %c1)
+  %b4 = tail call builder @llvm.tvm.newc()
+  %b5 = tail call builder @llvm.tvm.stslice(slice %s1, builder %b4)
+  %c2 = tail call cell @llvm.tvm.endc(builder %b5)
+  %s2 = tail call slice @llvm.tvm.ctos(cell %c2)
+  %b7 = tail call builder @llvm.tvm.newc()
+  %b8 = tail call builder @llvm.tvm.stslice(slice %s2, builder %b7)
+  ret builder %b8
+}
+declare builder @llvm.tvm.stu(i257, builder, i257)
+declare builder @llvm.tvm.sti(i257, builder, i257)
+declare builder @llvm.tvm.stvarint16(builder, i257)
+declare builder @llvm.tvm.stvaruint32(builder, i257)
+declare builder @llvm.tvm.newc()
+declare cell @llvm.tvm.endc(builder)
+declare builder @llvm.tvm.stslice(slice, builder)
+declare slice @llvm.tvm.ctos(cell)