//===------------- MarkMELVLoops.h - Mark candidates for MELV ------------===//
//
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
// See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
//
//===----------------------------------------------------------------------===//
//====  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.
//
//===----------------------------------------------------------------------===//
//
// This pass marks the loops that can be vectorized using the multi exit loop
// vectorizer (MELV). The major objective of this pass is to find a suitable
// vector factor which is used for all these loops. These loops are annotated
// with metadata (along with the chosen vector factor) so that MELV can
// vectorize them directly.
//
//===----------------------------------------------------------------------===//

#ifndef LLVM_TRANSFORMS_SCALAR_MARKMELVLOOPS_H
#define LLVM_TRANSFORMS_SCALAR_MARKMELVLOOPS_H

#include "llvm/IR/PassManager.h"
#include "llvm/Passes/PassBuilder.h"
#include "llvm/Passes/PassPlugin.h"
#include "llvm/ADT/Statistic.h"
#include "llvm/Analysis/AssumptionCache.h"
#include "llvm/Analysis/CFG.h"
#include "llvm/Analysis/DependenceAnalysis.h"
#include "llvm/Analysis/DependenceGraphBuilder.h"
#include "llvm/Analysis/DomTreeUpdater.h"
#include "llvm/Analysis/IVDescriptors.h"
#include "llvm/Analysis/LoopInfo.h"
#include "llvm/Analysis/LoopIterator.h"
#include "llvm/Analysis/LoopPass.h"
#include "llvm/Analysis/MemorySSA.h"
#include "llvm/Analysis/MemorySSAUpdater.h"
#include "llvm/Analysis/ScalarEvolution.h"
#include "llvm/Analysis/ScalarEvolutionExpressions.h"
#include "llvm/Analysis/TargetLibraryInfo.h"
#include "llvm/Analysis/TargetTransformInfo.h"
#include "llvm/IR/BasicBlock.h"
#include "llvm/IR/CFG.h"
#include "llvm/IR/Constants.h"
#include "llvm/IR/DataLayout.h"
#include "llvm/IR/Dominators.h"
#include "llvm/IR/Module.h"
#include "llvm/Pass.h"
#include "llvm/Support/Debug.h"
#include "llvm/Support/raw_ostream.h"
#include "llvm/Transforms/Utils/BasicBlockUtils.h"
#include "llvm/Transforms/Utils/LoopSimplify.h"

#define DEBUG_TYPE "mark-melv-loops"

namespace llvm {

// Class to store the results for this pass.
class MELVLoops {
  // The Loops that can be MELV vectorized.
  SmallVector<Loop*> Loops;
  // The selected vector factor.
  unsigned VF;
  // Map between a particular vector factor and the number of loops that can
  // be vectorized with that vector factor.
  using VFToNumLoopsMap = MapVector<unsigned, unsigned>;
  VFToNumLoopsMap VFToNumLoops;
public:
  MELVLoops() {
    VFToNumLoops.clear();
  }

  SmallVector<Loop *> getLoops() {
    return Loops;
  }
  // Returning false as the results are expected to be unchanged as
  // before LoopVectorize Pass.
  bool invalidate(Module &, const PreservedAnalyses &,
                  ModuleAnalysisManager::Invalidator &) {
    return false;
  }

  // Add the loop that can be MELV vectorized.
  void addLoop(Loop *L) {
    Loops.push_back(L);
  }

  // Increment the number of loops with this vector factor. This information
  // will be used to determine a unified vector factor for all the MELV
  // vectorizable loop in this module
  void incVFToNumLoopsCount(unsigned VF) {
    if (VFToNumLoops.count(VF) == 0) {
      VFToNumLoops[VF] = 1;
      return;
    }
    VFToNumLoops[VF]++;
    return;
  }
  
  // Select the vector factor that maps to highest number of loops that cab be
  // vectorized.
  void selectVF() {
    VF = 0;
    unsigned NumLoops = 0;
    for (auto &I : VFToNumLoops) {
      if (I.second > NumLoops) {
        NumLoops = I.second;
        VF = I.first;
      }
    }
    return;
  }

  // Get the selected vector factor.
  unsigned getSelectedVF() {
    return VF;
  }

  // Check if loop L is found to be vectorized with MELV.
  bool isMELVLoop(Loop *L) {
    for (Loop *L1 : Loops) 
      if (L1 == L)
        return true;
    return false;
  }
};

class MarkMELVLoops {
public:

  MarkMELVLoops(LoopInfo *LI, ScalarEvolution *SE, TargetTransformInfo *TTI,
                DominatorTree *DT, AssumptionCache *AC, MELVLoops *MELV)
                : LI(LI), SE(SE), TTI(TTI), DT(DT), AC(AC), MELV(MELV) { }
  ~MarkMELVLoops() {}

  bool runOnFunction(Function &F);

private:
  // The result of this analysis.
  MELVLoops *MELV;
  ScalarEvolution *SE;
  LoopInfo *LI;
  TargetTransformInfo *TTI;
  DominatorTree *DT;
  AssumptionCache *AC;
  /// The original loop.
  Loop *OrigLoop;
  /// Pointer induction variable if it is found.
  PHINode *PtrInduction = nullptr;
  /// True is Orignal Loop is while loop
  bool isWhile = false;
  /// Trip count of the original loop.
  Value *TripCount;
  /// The vectorization SIMD factor to use. Each vector will have this many
  /// vector elements.
  unsigned VF;
  /// Induction of Original Loop
  PHINode *OrigInduction = nullptr;
  /// The branch instruction used to determine the branches of the vector
  /// compare block.
  BranchInst *OrigLoopBreakBranchInst = nullptr;
  /// Phi incrementing instruction
  Instruction *PhiIncInst = nullptr;

  /// Checks whether the Loop L is vectorizable.
  bool SearchPattern(Loop *L);
  // IsSafe function checks whether the loop with multiple exits is suitable for
  // vectorization or not
  // Currently it checks only for the store conditions.
  bool IsSafe(Loop *L);
  // We allow maximum of two phi nodes. If we have two phi one of them should
  // be of pointer type.
  // FIXME Not Required once MultiExitLoop vectorization is generalised
  bool hasMaxTwoPHIsOneWithPtrType(bool &hasPtrPhi);

  bool isInvariantOrConsecutveAccess(Loop *L, Value *PrevInst);
  // Returns base pointer of memory being accessed within the loop
  // Return null if loop has more than one memory base ptr
  Value *getMemBasePtr(Loop *L);
  // Returns pointer to the compare instruction of BB's terminal branching
  // is based on or else returns null
  CmpInst *getTermCmpInst(BasicBlock *BB);
  // Returns the integer iterator of Loop L
  PHINode *getLoopIterator(Loop *L);
  // Set vector factor according to the type of the loaded value.
  void setVF(LoadInst *L);
  // If Expr is one of the Operands of GEP then returns if the calculated
  // address is accessed consecutively across iterations
  bool isConsecutiveAccess(const SCEV *Expr);
  // Checks whether Loop CmpOperand is loop invariant or load based on the
  // induction variable whose induction constant is 1
  bool SearchRecurr(Loop *L, Value *CmpOperand);
  // CanHandleWhileLoop function returns true only when the induction block
  // has no further uses of incremented induction variable apart from comparing
  // the result.
  bool CanHandleWhileLoop();
  // Collect the loops that are possible candidates for MELV. The checks are
  // exactly similar to that which are performed withing LoopVectorizer pass.
  // Only inner most loops which do not have irreducible control flow is allowed.
  void collectLoops(Loop *L, SmallVectorImpl<Loop *> &Worklist);

  // Data structures and functions extracted from the class
  // melv::LoopVectorizationLegality. For this pass we don't need to have all
  // the functionalities of LoopVectorizationLegality.

  /// InductionList saves induction variables and maps them to the
  /// induction descriptor.
  typedef MapVector<PHINode *, InductionDescriptor> InductionList;
  /// Holds all of the induction variables that we found in the loop.
  /// Notice that inductions don't need to start at zero and that induction
  /// variables can be pointers.
  InductionList Inductions;
  /// Sets Inductions for given Phi and Induction Descriptor ID
  void setInductionVar(PHINode *Phi, InductionDescriptor &ID) {
    Inductions[Phi] = ID;
  }
};

class Function;
class Module;

class MarkMELVLoopsPass : public PassInfoMixin<MarkMELVLoopsPass> {
  void runOnFunction(Function &F, FunctionAnalysisManager &AM);
public:
  MELVLoops MELV;
  PreservedAnalyses run(Module &M, ModuleAnalysisManager &MAM);
};

} // end namespace llvm

#endif // LLVM_TRANSFORMS_SCALAR_MARKMELVLOOPS_H
