//===- LoopBlock.h - Analyze and tile loops -------------------------------===//
//
//                     The LLVM Compiler Infrastructure
//
// 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 file implements a set of routines that performs strip-mining on loops.
// When done on multiple dimensions, stripmining is also called blocking/tiling.
//
//===----------------------------------------------------------------------===//

#include "llvm/Analysis/CGSCCPassManager.h"
#include "llvm/Analysis/LazyCallGraph.h"
#include "llvm/IR/PassManager.h"
#include "llvm/Pass.h"
#include "llvm/ADT/SmallVector.h"
#include "llvm/Analysis/AliasAnalysis.h"
#include "llvm/Analysis/DependenceAnalysis.h"
#include "llvm/Analysis/Loads.h"
#include "llvm/Analysis/LoopPass.h"
#include "llvm/Analysis/PostDominators.h"
#include "llvm/Analysis/ScalarEvolution.h"
#include "llvm/Analysis/ScalarEvolutionExpressions.h"
#include "llvm/IR/Value.h"
#include "llvm/IR/Instructions.h"
#include "llvm/IR/Dominators.h"
#include "llvm/Support/raw_ostream.h"
#include "llvm/Support/Debug.h"
#include "llvm/Transforms/Scalar/LoopUnrotate.h"
#include "llvm/Transforms/Utils/AffineTransforms.h"
#include "llvm/Transforms/Utils/BasicBlockUtils.h"
#include "llvm/Transforms/Utils/Cloning.h"
#include "llvm/Transforms/Utils/DepUtils.h"
#include "llvm/Transforms/Utils/LoopUtils.h"
#include "llvm/Transforms/Utils/MatrixOperations.h"
#include "llvm/Transforms/Utils/ScalarEvolutionExpander.h"
#include "llvm/Analysis/TargetTransformInfo.h"
#include "llvm/IR/IntrinsicInst.h"
#include "llvm/InitializePasses.h"
#include "llvm/Transforms/Utils/Local.h"
#include "llvm/Transforms/Utils/ValueMapper.h"
#include "llvm/Transforms/Utils/LoopVersioning.h"
#include <map>

using namespace llvm;
using namespace std;
using namespace MatOp;

namespace llvm {

struct LoopBlock {
  AliasAnalysis *AA;
  LoopInfo *LI;
  DependenceInfo *DA;
  DominatorTree *DT;
  PostDominatorTree *PDT;
  ScalarEvolution *SE;
  OptimizationRemarkEmitter *ORE;
  LoopAccessLegacyAnalysis *LAA = nullptr;
  const TargetTransformInfo *TTI;
  const DataLayout *DL;
  AssumptionCache *AC;
  const LoopAccessInfo *LAI = nullptr;
  FunctionAnalysisManager &FAM;
  TargetLibraryInfo *TLI;

  bool IRModified = false;

  const LoopAccessInfo *getLoopInfo(Loop *L) const {
    auto &LAIs = FAM.getResult<LoopAccessAnalysis>(*(L->getHeader()->getParent()));
    return &LAIs.getInfo(*L);
  }

  explicit LoopBlock(TargetTransformInfo *TTI_, AliasAnalysis *AA_,
            DominatorTree *DT_, PostDominatorTree *PDT_, LoopInfo *LI_,
            AssumptionCache *AC_, ScalarEvolution *SE_, DependenceInfo *DA_,
            FunctionAnalysisManager &FAM, OptimizationRemarkEmitter *ORE_,
            TargetLibraryInfo *TLI_): TTI(TTI_), AA(AA_), DT(DT_), PDT(PDT_),
            LI(LI_), AC(AC_), SE(SE_), DA(DA_), FAM(FAM), ORE(ORE_), TLI(TLI_) {}

  // Perfect nest under consideration - outermost to innermost assumed.
  LoopNestType LoopNest;

  // Stripmine candidate loops.
  SmallVector<Loop *, 2> SMCands;

  // Stripmine factor.
  DenseMap<Loop*, unsigned> SMFactor;
  // True if Interchange can be done with runtime checks.
  bool withRtChecks = false;

  bool LegalityChecks();
  bool DependenceChecks();
  bool ProfitabilityChecks();
  bool PerformTransformation();
  // Check if the loops can be fused
  bool checkFusionLegality(Loop *L1, Loop *L2, unsigned L1idx, unsigned L2idx,
                           bool &hasReuse);
  const SCEVAddRecExpr *getARscev(ScalarEvolution *SE, Value *Ptr, Loop *L,
                                  Instruction *LdSt,
                                  SmallVectorImpl<Value *> &MinLBs,
                                  unsigned dim);

  bool HoistInstruction(Instruction * I, Instruction *Ctx,
                      SmallVectorImpl<BasicBlock *> &BBs);

  bool LoadPtrFailsDerefCheck(LoadInst *Ld, Instruction *CtxI,
                              AliasSetTracker &AST, bool IsLoadAcessCheckReq = true);
  bool LICMForInstsIn(BasicBlock *BB, LoopInfo *LI, DominatorTree *DT,
                      ScalarEvolution *SE, AAResults *AA);
  void InterchangeLoops(Loop *L1, Loop *L2);
  Loop* StripMineLoop(Loop *L, unsigned SmF);

  bool AnalyseAndTransform();

  bool runOnFunction(Function &F);

  /* Fuse a group of consecutive loops */
  bool fuseTile(std::vector<LoopNestType> & Grp, Function &F);
  bool canInterchange(SmallVectorImpl<Loop *> &Loops);

  void getInnerLoops(SmallVectorImpl<Loop *> &Loops,
                     SmallVectorImpl<Loop *> &InnerLoops, bool &All1DimLoops);

  bool getLoopBounds(SmallVectorImpl<Loop *> &InnerLoops,
                     SmallVectorImpl<const SCEV *> &LBs,
                     SmallVectorImpl<const SCEV *> &UBs,
                     SmallPtrSetImpl<const SCEV *> &MinLBs,
                     SmallPtrSetImpl<const SCEV *> &MinUBs);

  bool checkForInvariantBounds(SmallVectorImpl<Loop *> &Loops,
                               SmallPtrSetImpl<const SCEV *> &MinLBs,
                               SmallPtrSetImpl<const SCEV *> &MinUBs);

  bool fusionLegalityChecks(SmallVectorImpl<Loop *> &InnerLoops,
                            SmallVectorImpl<Loop *> &Loops,
                            unsigned &smallerSet);
  bool TransformLoopNest(LoopNestType &P);
  bool IsProfitableToTile(Loop *L, bool &NotSafeForBlocking, bool CandidateCheck);
  Loop *TileWithFusion(SmallVectorImpl<Loop *> &Loops,
                       SmallVectorImpl<Loop *> &InnerLoops,
                       bool VersioningRequired, int maxDiff);
  bool isUnrotatedCandidateLoop(Loop *L) {
    if (std::find(unRotatedCandidates.begin(), unRotatedCandidates.end(), L) ==
        unRotatedCandidates.end()) {
      return false;
    }
    return true;
  }
  struct LoopBound {
    const SCEV *LB, *UB;
  };

  SmallVector<Value *, 2> LBVals, UBVals;
  SmallVector<Loop *, 2> unRotatedCandidates;
  std::map<Loop *, LoopBound> LoopBoundMap;
  std::map<Value *, unsigned> FirstRMap;
  std::map<Value *, unsigned> FirstWMap;

  void updateAccess(Value *Accs, Value *Arr, unsigned Lidx) {
    if (isa<LoadInst>(Accs)) {
      if (FirstRMap.find(Arr) == FirstRMap.end()) {
        FirstRMap[Arr] = Lidx;
      }
    }
    if (isa<StoreInst>(Accs)) {
      if (FirstWMap.find(Arr) == FirstWMap.end()) {
        FirstWMap[Arr] = Lidx;
      }
    }
  }
  bool wasWrittenBefore(Value *Arr, unsigned idx) {
    if (FirstWMap.find(Arr) != FirstWMap.end()) {
      if (idx >= FirstWMap[Arr])
        return true;
    }
    return false;
  }
  void updateLowerBound(Loop *L, const SCEV *LB, Value *TiledBound,
                        Instruction *Loc);
};

struct LoopBlockLegacyPass : public FunctionPass {
  static char ID;
  LoopBlockLegacyPass() : FunctionPass(ID) {
    initializeLoopBlockLegacyPassPass(*PassRegistry::getPassRegistry());
  }
  bool runOnFunction(Function &F) override;
    void getAnalysisUsage(AnalysisUsage &AU) const override {
    AU.addRequired<AAResultsWrapperPass>();
    AU.addRequired<LoopInfoWrapperPass>();
    AU.addRequired<DependenceAnalysisWrapperPass>();
    AU.addRequired<DominatorTreeWrapperPass>();
    AU.addRequired<PostDominatorTreeWrapperPass>();
    AU.addRequired<ScalarEvolutionWrapperPass>();
    AU.addRequired<LoopAccessLegacyAnalysis>();
    AU.addRequired<OptimizationRemarkEmitterWrapperPass>();
    AU.addRequired<TargetLibraryInfoWrapperPass>();
    AU.addRequired<TargetTransformInfoWrapperPass>();
    AU.addRequired<AssumptionCacheTracker>();
  }
};

struct LoopBlockPass: public PassInfoMixin<LoopBlockPass> {
  PreservedAnalyses run(Module &M, ModuleAnalysisManager &AM);
  bool runOnFunction(Function &F, FunctionAnalysisManager &AM);
};
struct LoopBlockTestPass: public PassInfoMixin<LoopBlockTestPass> {
  PreservedAnalyses run(Function &F, FunctionAnalysisManager &AM);
};

} // namespace
