lib/RewriteInsertsPass.cpp

// Copyright 2017 The Clspv Authors. All rights reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#include <string>
#include <utility>

#include "llvm/ADT/DenseMap.h"
#include "llvm/ADT/SmallVector.h"
#include "llvm/ADT/UniqueVector.h"
#include "llvm/IR/Attributes.h"
#include "llvm/IR/Constants.h"
#include "llvm/IR/DerivedTypes.h"
#include "llvm/IR/Function.h"
#include "llvm/IR/Instructions.h"
#include "llvm/IR/Module.h"
#include "llvm/Pass.h"
#include "llvm/Support/raw_ostream.h"

#include "clspv/Option.h"

using namespace llvm;
using std::string;

#define DEBUG_TYPE "rewriteinserts"

namespace {

const char *kCompositeConstructFunctionPrefix = "clspv.composite_construct.";

class RewriteInsertsPass : public ModulePass {
public:
  static char ID;
  RewriteInsertsPass() : ModulePass(ID) {}

  bool runOnModule(Module &M) override;

private:
  using InsertionVector = SmallVector<Instruction *, 4>;

  // Replaces chains of insertions that cover the entire value.
  // Such a change always reduces the number of instructions, so
  // we always perform these.  Returns true if the module was modified.
  bool ReplaceCompleteInsertionChains(Module &M);

  // Replaces all InsertValue instructions, even if they aren't part
  // of a complete insetion chain.  Returns true if the module was modified.
  bool ReplacePartialInsertions(Module &M);

  // Load |values| and |chain| with the members of the struct value produced
  // by a chain of InsertValue instructions ending with |iv|, and following
  // the aggregate operand.  Return the start of the chain: the aggregate
  // value which is not an InsertValue instruction, or an InsertValue
  // instruction which inserts a component that is replaced later in the
  // chain.  The |values| vector will match the order of struct members and
  // is initialized to all nullptr members.  The |chain| vector will list
  // the chain of InsertValue instructions, listed in the order we discover
  // them, e.g. begining with |iv|.
  Value *LoadValuesEndingWithInsertion(InsertValueInst *iv,
                                       std::vector<Value *> *values,
                                       InsertionVector *chain) {
    auto *structTy = dyn_cast<StructType>(iv->getType());
    assert(structTy);
    const auto numElems = structTy->getNumElements();

    // Walk backward from the tail to an instruction we don't want to
    // replace.
    Value *frontier = iv;
    while (auto *insertion = dyn_cast<InsertValueInst>(frontier)) {
      chain->push_back(insertion);
      // Only handle single-index insertions.
      if (insertion->getNumIndices() == 1) {
        // Try to replace this one.

        unsigned index = insertion->getIndices()[0];
        assert(index < numElems);
        if ((*values)[index] != nullptr) {
          // We already have a value for this slot.  Stop now.
          break;
        }
        (*values)[index] = insertion->getInsertedValueOperand();
        frontier = insertion->getAggregateOperand();
      } else {
        break;
      }
    }
    return frontier;
  }

  // Returns the number of elements in the struct or array.
  unsigned GetNumElements(Type *type) {
    // CompositeType doesn't implement getNumElements(), but its inheritors
    // do.
    if (auto *struct_ty = dyn_cast<StructType>(type)) {
      return struct_ty->getNumElements();
    } else if (auto *seq_ty = dyn_cast<SequentialType>(type)) {
      return seq_ty->getNumElements();
    }
    return 0;
  }

  // If this is the tail of a chain of InsertValueInst instructions
  // that covers the entire composite, then return a small vector
  // containing the insertion instructions, in member order.
  // Otherwise returns nullptr.
  InsertionVector *CompleteInsertionChain(InsertValueInst *iv) {
    if (iv->getNumIndices() == 1) {
      auto numElems = GetNumElements(iv->getType());
      if (numElems != 0) {
        // Only handle single-index insertions.
        unsigned index = iv->getIndices()[0];
        if (index + 1u != numElems) {
          // Not the last in the chain.
          return nullptr;
        }
        InsertionVector candidates(numElems, nullptr);
        for (unsigned i = index;
             iv->getNumIndices() == 1 && i == iv->getIndices()[0]; --i) {
          // iv inserts the i'th member
          candidates[i] = iv;

          if (i == 0) {
            // We're done!
            return new InsertionVector(candidates);
          }

          if (InsertValueInst *agg =
                  dyn_cast<InsertValueInst>(iv->getAggregateOperand())) {
            iv = agg;
          } else {
            // The chain is broken.
            break;
          }
        }
      }
    }
    return nullptr;
  }

  // If this is the tail of a chain of InsertElementInst instructions
  // that covers the entire vector, then return a small vector
  // containing the insertion instructions, in member order.
  // Otherwise returns nullptr.  Only handle insertions into vectors.
  InsertionVector *CompleteInsertionChain(InsertElementInst *ie) {
    // Don't handle i8 vectors. Only <4 x i8> is supported and it is
    // translated as i32.  Only handle single-index insertions.
    if (auto vec_ty = dyn_cast<VectorType>(ie->getType())) {
      if (vec_ty->getVectorElementType() == Type::getInt8Ty(ie->getContext())) {
        return nullptr;
      }
    }

    // Only handle single-index insertions.
    if (ie->getNumOperands() == 3) {
      auto numElems = GetNumElements(ie->getType());
      if (numElems != 0) {
        if (auto *const_value = dyn_cast<ConstantInt>(ie->getOperand(2))) {
          uint64_t index = const_value->getZExtValue();
          if (index + 1u != numElems) {
            // Not the last in the chain.
            return nullptr;
          }
          InsertionVector candidates(numElems, nullptr);
          Value *value = ie;
          uint64_t i = index;
          while (auto *insert = dyn_cast<InsertElementInst>(value)) {
            if (insert->getNumOperands() != 3)
              break;
            if (auto *index_const =
                    dyn_cast<ConstantInt>(insert->getOperand(2))) {
              if (i != index_const->getZExtValue())
                break;

              candidates[i] = insert;
              if (i == 0) {
                // We're done!
                return new InsertionVector(candidates);
              }

              value = insert->getOperand(0);
              --i;
            } else {
              break;
            }
          }
        } else {
          return nullptr;
        }
      }
    }
    return nullptr;
  }

  // Return the name for the wrap function for the given type.
  string &WrapFunctionNameForType(Type *type) {
    auto where = function_for_type_.find(type);
    if (where == function_for_type_.end()) {
      // Insert it.
      auto &result = function_for_type_[type] =
          string(kCompositeConstructFunctionPrefix) +
          std::to_string(function_for_type_.size());
      return result;
    } else {
      return where->second;
    }
  }

  // Get or create the composite construct function definition.
  Function *GetConstructFunction(Module &M, CompositeType *constructed_type) {
    // Get or create the composite construct function definition.
    const string &fn_name = WrapFunctionNameForType(constructed_type);
    Function *fn = M.getFunction(fn_name);
    if (!fn) {
      // Make the function.
      SmallVector<Type *, 16> elements;
      unsigned num_elements = GetNumElements(constructed_type);
      for (unsigned i = 0; i != num_elements; ++i)
        elements.push_back(constructed_type->getTypeAtIndex(i));
      FunctionType *fnTy = FunctionType::get(constructed_type, elements, false);
      auto fn_constant = M.getOrInsertFunction(fn_name, fnTy);
      fn = cast<Function>(fn_constant.getCallee());
      fn->addFnAttr(Attribute::ReadOnly);
    }
    return fn;
  }

  // Maps a loaded type to the name of the wrap function for that type.
  DenseMap<Type *, string> function_for_type_;
};
} // namespace

char RewriteInsertsPass::ID = 0;
static RegisterPass<RewriteInsertsPass>
    X("RewriteInserts",
      "Rewrite chains of insertvalue to as composite-construction");

namespace clspv {
llvm::ModulePass *createRewriteInsertsPass() {
  return new RewriteInsertsPass();
}
} // namespace clspv

bool RewriteInsertsPass::runOnModule(Module &M) {
  bool Changed = ReplaceCompleteInsertionChains(M);

  if (clspv::Option::HackInserts()) {
    Changed |= ReplacePartialInsertions(M);
  }

  return Changed;
}

bool RewriteInsertsPass::ReplaceCompleteInsertionChains(Module &M) {
  bool Changed = false;

  SmallVector<InsertionVector *, 16> WorkList;
  for (Function &F : M) {
    for (BasicBlock &BB : F) {
      for (Instruction &I : BB) {
        if (InsertValueInst *iv = dyn_cast<InsertValueInst>(&I)) {
          if (InsertionVector *insertions = CompleteInsertionChain(iv)) {
            WorkList.push_back(insertions);
          }
        } else if (InsertElementInst *ie = dyn_cast<InsertElementInst>(&I)) {
          if (InsertionVector *insertions = CompleteInsertionChain(ie)) {
            WorkList.push_back(insertions);
          }
        }
      }
    }
  }

  if (WorkList.size() == 0) {
    return Changed;
  }

  for (InsertionVector *insertions : WorkList) {
    Changed = true;

    // Gather the member values and types.
    SmallVector<Value *, 4> values;
    for (Instruction *inst : *insertions) {
      if (auto *insert_value = dyn_cast<InsertValueInst>(inst)) {
        values.push_back(insert_value->getInsertedValueOperand());
      } else if (auto *insert_element = dyn_cast<InsertElementInst>(inst)) {
        values.push_back(insert_element->getOperand(1));
      } else {
        llvm_unreachable("Unhandled insertion instruction");
      }
    }

    CompositeType *resultTy =
        cast<CompositeType>(insertions->back()->getType());
    Function *fn = GetConstructFunction(M, resultTy);

    // Replace the chain.
    auto call = CallInst::Create(fn, values);
    call->insertAfter(insertions->back());
    insertions->back()->replaceAllUsesWith(call);

    // Remove the insertions if we can.  Go from the tail back to
    // the head, since the tail uses the previous insertion, etc.
    for (auto iter = insertions->rbegin(), end = insertions->rend();
         iter != end; ++iter) {
      Instruction *insertion = *iter;
      if (!insertion->hasNUsesOrMore(1)) {
        insertion->eraseFromParent();
      }
    }

    delete insertions;
  }

  return Changed;
}

bool RewriteInsertsPass::ReplacePartialInsertions(Module &M) {
  bool Changed = false;

  // First find candidates.  Collect all InsertValue instructions
  // into struct type, but track their interdependencies.  To minimize
  // the number of new instructions, generate a construction for each
  // tail of an insertion chain.

  UniqueVector<InsertValueInst *> insertions;
  for (Function &F : M) {
    for (BasicBlock &BB : F) {
      for (Instruction &I : BB) {
        if (InsertValueInst *iv = dyn_cast<InsertValueInst>(&I)) {
          if (iv->getType()->isStructTy()) {
            insertions.insert(iv);
          }
        }
      }
    }
  }

  // Now count how many times each InsertValue is used by another InsertValue.
  // The |num_uses| vector is indexed by the unique id that |insertions|
  // assigns to it.
  std::vector<unsigned> num_uses(insertions.size() + 1);
  // Count from the user's perspective.
  for (InsertValueInst *insertion : insertions) {
    if (auto *agg =
            dyn_cast<InsertValueInst>(insertion->getAggregateOperand())) {
      ++(num_uses[insertions.idFor(agg)]);
    }
  }

  // Proceed in rounds.  Each round rewrites any chains ending with an
  // insertion that is not used by another insertion.

  // Get the first list of insertion tails.
  InsertionVector WorkList;
  for (InsertValueInst *insertion : insertions) {
    if (num_uses[insertions.idFor(insertion)] == 0) {
      WorkList.push_back(insertion);
    }
  }

  // This records insertions in the order they should be removed.
  // In this list, an insertion preceds any insertions it uses.
  // (This is post-dominance order.)
  InsertionVector ordered_candidates_for_removal;

  // Proceed in rounds.
  while (WorkList.size()) {
    Changed = true;

    // Record the list of tails for the next round.
    InsertionVector NextRoundWorkList;

    for (Instruction *inst : WorkList) {
      InsertValueInst *insertion = cast<InsertValueInst>(inst);
      // Rewrite |insertion|.

      StructType *resultTy = cast<StructType>(insertion->getType());

      const unsigned num_members = resultTy->getNumElements();
      std::vector<Value *> members(num_members, nullptr);
      InsertionVector chain;
      // Gather the member values.  Walk backward from the insertion.
      Value *base = LoadValuesEndingWithInsertion(insertion, &members, &chain);

      // Populate remaining entries in |values| by extracting elements
      // from |base|.  Only make a new extractvalue instruction if we can't
      // make a constant or undef.  New instructions are inserted before
      // the insertion we plan to remove.
      for (unsigned i = 0; i < num_members; ++i) {
        if (!members[i]) {
          Type *memberTy = resultTy->getTypeAtIndex(i);
          if (isa<UndefValue>(base)) {
            members[i] = UndefValue::get(memberTy);
          } else if (const auto *caz = dyn_cast<ConstantAggregateZero>(base)) {
            members[i] = caz->getElementValue(i);
          } else if (const auto *ca = dyn_cast<ConstantAggregate>(base)) {
            members[i] = ca->getOperand(i);
          } else {
            members[i] = ExtractValueInst::Create(base, {i}, "", insertion);
          }
        }
      }

      // Create the call.  It's dominated by any extractions we've just
      // created.
      Function *construct_fn = GetConstructFunction(M, resultTy);
      auto *call = CallInst::Create(construct_fn, members, "", insertion);

      // Disconnect this insertion.  We'll remove it later.
      insertion->replaceAllUsesWith(call);

      // Trace backward through the chain, removing uses and deleting where
      // we can.  Stop at the first element that has a remaining use.
      for (auto *chainElem : chain) {
        if (chainElem->hasNUsesOrMore(1)) {
          unsigned &use_count =
              num_uses[insertions.idFor(cast<InsertValueInst>(chainElem))];
          assert(use_count > 0);
          --use_count;
          if (use_count == 0) {
            NextRoundWorkList.push_back(chainElem);
          }
          break;
        } else {
          chainElem->eraseFromParent();
        }
      }
    }
    WorkList = std::move(NextRoundWorkList);
  }

  return Changed;
}