//===- PartialInvarianceAnalysis.h - Analyze loop for partail invariance --===//
//
//===----------------------------------------------------------------------===//
//
//                     The LLVM Compiler Infrastructure
//
// This file is distributed under the University of Illinois Open Source
// License. See LICENSE.TXT for details.
//
//===----------------------------------------------------------------------===//
//
// Copyright (c) 2015 Advanced Micro Devices, Inc.  All rights reserved.
// Developed by: Advanced Micro Devices, Inc.
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// with the Software without restriction, including without limitation the
// rights to use, copy, modify, merge, publish, distribute, sublicense, and/or
// sell copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
// Redistributions of source code must retain the above copyright notice, this
// list of conditions and the following disclaimers.
// Redistributions in binary form must reproduce the above copyright notice,
// this list of conditions and the following disclaimers in the documentation
// and/or other materials provided with the distribution.
// Neither the names of Advanced Micro Devices, Inc., nor the names of its
// contributors may be used to endorse or promote products derived from this
// Software without specific prior written permission.
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// CONTRIBUTORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS WITH
// THE SOFTWARE.
//
//===----------------------------------------------------------------------===//
#include "llvm/ADT/MapVector.h"
#include "llvm/ADT/STLExtras.h"
#include "llvm/ADT/SmallPtrSet.h"
#include "llvm/Analysis/AliasAnalysis.h"
#include "llvm/Analysis/AliasSetTracker.h"
#include "llvm/Analysis/BranchProbabilityInfo.h"
#include "llvm/Analysis/CodeMetrics.h"
#include "llvm/Analysis/GlobalsModRef.h"
#include "llvm/Analysis/InstructionSimplify.h"
#include "llvm/Analysis/LoopAnalysisManager.h"
#include "llvm/Analysis/LoopInfo.h"
#include "llvm/Analysis/LoopPass.h"
#include "llvm/Analysis/ScalarEvolution.h"
#include "llvm/Analysis/TargetTransformInfo.h"
#include "llvm/IR/Constants.h"
#include "llvm/IR/DerivedTypes.h"
#include "llvm/IR/Dominators.h"
#include "llvm/IR/Function.h"
#include "llvm/IR/Instructions.h"
#include "llvm/IR/IntrinsicInst.h"
#include "llvm/IR/MDBuilder.h"
#include "llvm/IR/Module.h"
#include "llvm/InitializePasses.h"
#include "llvm/Support/BranchProbability.h"
#include "llvm/Support/CommandLine.h"
#include "llvm/Support/Debug.h"
#include "llvm/Support/raw_ostream.h"
#include "llvm/Transforms/Scalar.h"
#include "llvm/Transforms/Utils/BasicBlockUtils.h"
#include "llvm/Transforms/Utils/Cloning.h"
#include "llvm/Transforms/Utils/Local.h"
#include "llvm/Transforms/Utils/LoopUtils.h"

using namespace llvm;

namespace llvm {

struct PartialInvarianceInfo {
  // Maintain information about path.
  struct PathInfo {
    unsigned NumberOfInstructions;
    unsigned NumberOfBlocks;
    BranchInst *TBranch;
    BasicBlock *SuccBlock;
    unsigned SuccIndex;
    // Update block in path informaion.
    void Update(BasicBlock *BB) {
      NumberOfBlocks++;
      NumberOfInstructions += BB->size();
    }
    // Clear collected path information.
    void Clear() {
      NumberOfInstructions = 0;
      NumberOfBlocks = 0;
      SuccIndex = 0;
      TBranch = nullptr;
      SuccBlock = nullptr;
    }
    // Initial setter for path information.
    void SetPathInformation(BranchInst *HeaderTI, BasicBlock *SuccBB,
                            unsigned SuccIdx) {
      TBranch = HeaderTI;
      SuccBlock = SuccBB;
      SuccIndex = SuccIdx;
    }
    void dump();
  }; // end: PathInfo

  PartialInvarianceInfo(Loop *L, LoopInfo *LI, AliasAnalysis *AA)
      : L(L), AA(AA), HeaderTI(nullptr), HasPartialInvariance(false),
        IsLoopInvariantCondition(true) {
    PreHeader = L->getLoopPreheader();
    LoopHeader = L->getHeader();
    LatchBlock = L->getLoopLatch();
    // Clear containers.
    MemoryAccessesInCondition.clear();
    LoopExitBlocks.clear();
    LoopExitingBlocks.clear();
    PathExitBlocks.clear();
    PathExitingBlocks.clear();
    InstructionsPartOfCondition.clear();
    InvariantPathInfo.Clear();
    // Analyze for path invariance.
    analyze();
  }

  // Return true if loop has invariant in any of the path.
  bool hasPartialInvariance() const {
    return HasPartialInvariance && !hasLoopInvariantCondition();
  }
  // Return true if condition in header is loop invariant.
  bool hasLoopInvariantCondition() const { return IsLoopInvariantCondition; }
  // Return true if invariant path has no instructions.
  bool hasEmptyInvariantPath() const {
    return !InvariantPathInfo.NumberOfInstructions;
  }
  // Return the sucessor number of the header condition where header condition
  // remains invariant.
  unsigned getInvariancePathIndex() const {
    return InvariantPathInfo.SuccIndex;
  }
  // Return exit block list in path.
  SmallVector<BasicBlock *, 32> &getPathExitBlocks() { return PathExitBlocks; }
  // Return exiting block list in path.
  SmallVector<BasicBlock *, 32> &getPathExitingBlocks() {
    return PathExitingBlocks;
  }
  // Return if path has a single exit block, else nullptr.
  BasicBlock *getUniqueExitBlockInPath() const {
    if (PathExitBlocks.size() == 1)
      return PathExitBlocks[0];
    else
      return nullptr;
  }
  // Return true if path has single exiting block, else nullptr.
  BasicBlock *getUniqueExitingBlockInPath() const {
    if (PathExitingBlocks.size() == 1)
      return PathExitingBlocks[0];
    else
      return nullptr;
  }
  // Return true if input instrucion is part of condition.
  bool isInstructionsPartOfCondition(Instruction *I) const {
    return InstructionsPartOfCondition.find(I) !=
           InstructionsPartOfCondition.end();
  }

private:
  Loop *L;
  AliasAnalysis *AA;
  BranchInst *HeaderTI;
  BasicBlock *PreHeader;
  BasicBlock *LoopHeader;
  BasicBlock *LatchBlock;
  bool HasPartialInvariance;
  bool IsLoopInvariantCondition;
  SmallVector<Instruction *, 32> MemoryAccessesInCondition;
  SmallVector<BasicBlock *, 32> LoopExitBlocks;
  SmallVector<BasicBlock *, 32> LoopExitingBlocks;
  SmallVector<BasicBlock *, 32> PathExitBlocks;
  SmallVector<BasicBlock *, 32> PathExitingBlocks;
  MapVector<Instruction *, bool> InstructionsPartOfCondition;
  PathInfo InvariantPathInfo;

  void analyze();
  bool hasLegalLoopStructure();

  bool hasInvariantInAnyBranchPath();
  bool analyzePathForInvariance(BasicBlock *BB, BasicBlock *LatchBlock,
                                MapVector<BasicBlock *, bool> &BlockList);

  bool isSafeCall(CallInst *CI);
  bool isInstructionSafeForPath(Instruction &I);

  bool hasSimpleConditionInHeader();
  bool analyzeConditionDependencies(Instruction *Inst, BasicBlock *CurrBlock);

  // Return true if input basic block is a loop exit block.
  bool isLoopExitBlock(const BasicBlock *BB) const {
    return std::find(LoopExitBlocks.begin(), LoopExitBlocks.end(), BB) !=
           LoopExitBlocks.end();
  }
  // Return true if input basic block is a loop exiting block.
  bool isLoopExitingBlock(const BasicBlock *BB) const {
    return std::find(LoopExitingBlocks.begin(), LoopExitingBlocks.end(), BB) !=
           LoopExitingBlocks.end();
  }
  // Add if not already in exit block list.
  void addBlockToPathExitBlocks(BasicBlock *BB) {
    if (std::find(PathExitBlocks.begin(), PathExitBlocks.end(), BB) ==
        PathExitBlocks.end()) {
      PathExitBlocks.push_back(BB);
    }
  }
  // Add if not already in exiting block list.
  void addBlockToPathExitingBlocks(BasicBlock *BB) {
    if (std::find(PathExitingBlocks.begin(), PathExitingBlocks.end(), BB) ==
        PathExitingBlocks.end()) {
      PathExitingBlocks.push_back(BB);
    }
  }
  // Reset path information.
  void ResetPathInfo() {
    InvariantPathInfo.Clear();
    PathExitBlocks.clear();
    PathExitingBlocks.clear();
  }
};

struct LoopPartialInvarianceAnalysis {
  LoopInfo *LI;
  AliasAnalysis *AA;
  DenseMap<Loop *, std::unique_ptr<PartialInvarianceInfo>>
     PartialInvarianceInfoMap;

  LoopPartialInvarianceAnalysis(LoopInfo *LI, AliasAnalysis *AA)
         : LI(LI), AA(AA) {
    PartialInvarianceInfoMap.clear();
  }

  void releaseMemory() {
    // Invalidate the cache when the pass is freed.
    PartialInvarianceInfoMap.clear();
  }

  void processLoop(Loop *L) {
    PartialInvarianceInfoMap[L] =
       std::make_unique<PartialInvarianceInfo>(L, LI, AA);
  }

  // Will return PartialInvarianceInfo.
  const PartialInvarianceInfo &getInfo(Loop *L);

  // Update the analysis
  void updateAnalysis(LoopInfo *NewLI, AliasAnalysis *NewAA) {
    LI = NewLI;
    AA = NewAA;
  }
};

struct LoopPartialInvarianceAnalysisLegacy : public LoopPass {
  static char ID;
  LoopPartialInvarianceAnalysis *LPIA;
  explicit LoopPartialInvarianceAnalysisLegacy()
    : LoopPass(ID), LPIA(nullptr) {
    initializeLoopPartialInvarianceAnalysisLegacyPass(
        *PassRegistry::getPassRegistry());
  }
  bool runOnLoop(Loop *L, LPPassManager &LPM) override;
  void getAnalysisUsage(AnalysisUsage &AU) const override {
    AU.addRequired<LoopInfoWrapperPass>();
    AU.addRequired<AAResultsWrapperPass>();
    AU.setPreservesAll();
  }
  // Will return PartialInvarianceInfo.
  const PartialInvarianceInfo &getInfo(Loop *L) {
    return LPIA->getInfo(L);
  }
  ~LoopPartialInvarianceAnalysisLegacy() {
    if (LPIA)
      delete LPIA;
  }
};

class LoopPartialInvarianceAnalysisPass :
      public AnalysisInfoMixin<LoopPartialInvarianceAnalysisPass> {
  friend AnalysisInfoMixin<LoopPartialInvarianceAnalysisPass>;
  static AnalysisKey Key;
public:
  typedef PartialInvarianceInfo Result;
  Result run(Loop &L, LoopAnalysisManager &AM, LoopStandardAnalysisResults &AR);
};

} // end llvm namespace
