
//===--------------- ISLValueAnalysis.h - Value Analysis Pass -------------===//
//
//===----------------------------------------------------------------------===//
//                     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.
//
//===----------------------------------------------------------------------===//
//
//===----------------------------------------------------------------------===//

#ifndef LLVM_TRANSFORMS_IPO_ISLVALUEANALYSIS_H
#define LLVM_TRANSFORMS_IPO_ISLVALUEANALYSIS_H

#include "llvm/Pass.h"
#include "llvm/IR/Module.h"
#include "llvm/Analysis/LoopInfo.h"
#include "llvm/IR/BasicBlock.h"

// ISL Headers:
#include "isl/aff.h"
#include "isl/constraint.h"
#include "isl/ctx.h"
#include "isl/id.h"
#include "isl/isl_space_private.h"
#include "isl/map.h"
#include "isl/set.h"
#include "isl/val.h"

#include <set>
#include <map>

namespace llvm {

enum class OpType { Dim, Param, Const };

struct OperandInfo {
  OpType Type;
  union {
    isl_id *PId;
    unsigned DimNum;
    long Const;
  } Opr;
  bool isUndef;
  OperandInfo() : Type(OpType::Dim), isUndef(false) { Opr.Const = 0; }
  OperandInfo(const OperandInfo &O) {
    errs() << "Copy constructor called\n";
    Type = O.Type;
    isUndef = O.isUndef;
    if (Type == OpType::Param)
      Opr.PId = isl_id_copy(O.Opr.PId);
    else
      Opr = O.Opr;
  }
  OperandInfo& operator=(const OperandInfo &O) {
    errs() << "Copy assignment called\n";
    Type = O.Type;
    isUndef = O.isUndef;
    if (Type == OpType::Param)
      Opr.PId = isl_id_copy(O.Opr.PId);
    else
      Opr = O.Opr;
    return *this;
  }
  OperandInfo(const OperandInfo &&O) {
    Type = O.Type;
    isUndef = O.isUndef;
    Opr = O.Opr;    
  }
  OperandInfo& operator=(const OperandInfo &&O) {
    Type = O.Type;
    isUndef = O.isUndef;
    Opr = O.Opr;
    return *this;
  }
};

struct Constraint {
  isl_set *Constr;
  bool IsDead;
  Constraint() {
    Constr = nullptr; IsDead = false;
  };
  Constraint(isl_set *C_, bool IsDead_) : Constr(C_), IsDead(IsDead_) {};
  Constraint copy() {
    return Constraint(isl_set_copy(Constr), IsDead);
  }
  void free() {
    if (Constr) {
      isl_set_free(Constr);
      Constr = nullptr;
    }      
  }
};

struct ValIdx {
  Value *V;
  unsigned Idx;
};

class ISLQueryBuilder;

class ISLFunctionInfo {
public:
  using IslMapVec = SmallVector<isl_map*, 1>;
  using IslIdVec = SmallVector<isl_id*, 1>;
  using ValToIslMapV = DenseMap<Value *, IslMapVec>;
  using ValToIslIdV = DenseMap<Value *, IslIdVec>;
  using LoopToIslId = DenseMap<Loop *, isl_id *>;
  using IslSpaceVec = SmallVector<isl_space*, 1>;  
  using ValSet = std::set<Value *>;
  using BBToConstrTy = DenseMap<BasicBlock *, Constraint>;
  using EdgeToConstrTy =
    DenseMap<std::pair<BasicBlock *, BasicBlock *>, Constraint>;
  using OperandInfoVec = SmallVector<OperandInfo, 1>;
  using PhiDimInfo =
      SmallVector<std::tuple<OperandInfoVec, Constraint, Value *>, 4>;
  friend class ISLQueryBuilder;
  bool isAnalysisValid;
  BBToConstrTy BBConstraintInfo;

private:
  Function *F;
  LoopInfo *LI;
  isl_ctx *Ctx;
  unsigned PseudoPhisCounter;
  unsigned MaxLoopHeightForPow;
  unsigned long MaxNoIslOps;

  std::set<BasicBlock*> IgnoreList;

  ValToIslMapV FoGMapInfo;
  ValToIslIdV IdInfo;
  LoopToIslId LoopIdInfo;
  ValToIslIdV PseudoPhiIdInfo;

  ValSet FunctionParams;
  ValSet MemoryParams;
  std::set<void *> PseudoPhiParams;

  EdgeToConstrTy EdgeConstraintInfo;

  void allocIdFor(Value *V);
  void preserveIdFor(Value *V, isl_id *Id, unsigned Idx = 0);
  isl_id *getIdFor(Value *V, unsigned Idx = 0);

  void allocPseudoPhiIdFor(Value *V);
  void preservePseudoPhiIdFor(Value *V, isl_id *Id, unsigned Idx);
  isl_id *getPseudoPhiIdFor(Value *V, unsigned Idx = 0);


  void allocIdFor(Loop *L);
  void preserveIdFor(Loop *L, isl_id *Id);
  isl_id *getIdFor(Loop *L);

  void preserveMapFor(Value *V, IslMapVec &M);
  IslMapVec getMapFor(Value *V);
  IslMapVec getFinalMapFor(Value *V);

  void collectFunctionParams();
  bool isLegal(Function *F);
  bool processFunction();
  void processInst(Instruction *I);

  // *******************************************************************
  //                        Instruction Transforms
  // *******************************************************************
  using OpVec = SmallVector<OperandInfo, 1>;
  using OpVecVec = SmallVector<OpVec, 2>;
  OpVec getOperandInfoVecFor(Value *V, ValSet &Params, unsigned &Dim);
  OpVecVec gatherAllOperandInfoFor(Instruction *I, ValSet &Params);
  void freeOperandInfo(OpVecVec &OpsVec);
  IslMapVec getTransformForMemory(Instruction *I, ValSet &Params);
  IslMapVec getTransformForBinaryOp(Instruction *I, OpVecVec &OpsVec);
  IslMapVec getTransformForBinaryOp(unsigned OpCode, OpVecVec &OpsVec,
                                    unsigned NBits);
  IslMapVec getTransformForICMP(CmpInst *CmpI, OpVecVec &OpsVec);
  IslMapVec getTransformForICMP(CmpInst::Predicate Predicate, OpVecVec &OpsVec);
  IslMapVec getTransformForCastOp(Instruction *I, OpVecVec &OpsVec);
  OpVecVec& broadcastOpsVec(OpVecVec &OpsVec);
  IslMapVec getGEPTransform(GetElementPtrInst *Gep, OpVecVec &OpsVec);
  IslSpaceVec getSpace(OpVecVec &OpsVec);
  
  IslMapVec getPureTransformForInst(Instruction *I, std::set<Value *> &ParamsSet);

  // Phis
  PhiDimInfo extractPhiInfofor(PHINode *PN, ValSet &ParamsSet,
                               EdgeToConstrTy &EdgeConstraintInfo);
  IslMapVec getPhiTransform(PhiDimInfo &PhiInfo, ValToIslMapV &MapInfo);

  // Vectors
  IslMapVec getVectorTrans(Instruction *I, ValToIslMapV &FoGMapInfo);
  

  // Transform getter helpers.
  template <typename F>
  IslMapVec getVectorTransForBinOp(OpVecVec &OpsVec, F GetTrans) {
    IslMapVec MpVec;
    IslSpaceVec SpVec = getSpace(OpsVec);
    OpVec &OpV1 = OpsVec[0];
    OpVec &OpV2 = OpsVec[1];
    unsigned VecLen = SpVec.size();
    for (unsigned I = 0; I < VecLen; I++) {
      MpVec.push_back(GetTrans(SpVec[I], OpV1[I], OpV2[I]));
    }
    return MpVec;
  }
  template <typename F>
  IslMapVec getVectorTransForBwOp(OpVecVec &OpsVec, F GetTrans, unsigned NBits) {
    IslMapVec MpVec;
    IslSpaceVec SpVec = getSpace(OpsVec);
    OpVec &OpV1 = OpsVec[0];
    OpVec &OpV2 = OpsVec[1];
    unsigned VecLen = SpVec.size();
    for (unsigned I = 0; I < VecLen; I++) {
      MpVec.push_back(GetTrans(SpVec[I], OpV1[I], OpV2[I], NBits));
    }
    return MpVec;
  }
  template <typename F>
  IslMapVec getVectorTransForUnOp(OpVecVec &OpsVec, F GetTrans) {
    IslMapVec MpVec;
    IslSpaceVec SpVec = getSpace(OpsVec);
    OpVec &OpV1 = OpsVec[0];
    unsigned VecLen = SpVec.size();
    for (unsigned I = 0; I < VecLen; I++) {
      MpVec.push_back(GetTrans(SpVec[I], OpV1[I]));
    }
    return MpVec;
  }
  template <typename F>
  IslMapVec getVectorTransForTerOp(OpVecVec &OpsVec, F GetTrans) {
    IslMapVec MpVec;
    IslSpaceVec SpVec = getSpace(OpsVec);
    OpVec &OpV1 = OpsVec[0];
    OpVec &OpV2 = OpsVec[1];
    OpVec &OpV3 = OpsVec[2];
    unsigned VecLen = SpVec.size();
    for (unsigned I = 0; I < VecLen; I++) {
      MpVec.push_back(GetTrans(SpVec[I], OpV1[I], OpV2[I], OpV3[I]));
    }
    return MpVec;
  }

  // *******************************************************************
  //                        Constraint Propagation
  // *******************************************************************
  std::map<BasicBlock *, SmallVector<isl_set *, 4>> BBPseudoConstraints;
  isl_set *dropMemoryConstraints(isl_set *S);
  void processConstraintsFromBr(BranchInst *BI, ValToIslMapV &MapInfo,
                                std::set<Value *> &Params,
                                BBToConstrTy &BBConstraintInfo,
                                EdgeToConstrTy &EdgeConstraintInfo,
                                bool DropMemoryConstraints);
  void createOutgoingConstraints(BasicBlock *BB, ValToIslMapV &MapInfo,
                                 std::set<Value *> &Params,
                                 BBToConstrTy &BBConstraintInfo,
                                 EdgeToConstrTy &EdgeConstraintInfo,
                                 bool DropMemoryConstraints = true);
  void markOutgoingEdgesAsDead(BasicBlock *BB,
                               EdgeToConstrTy &EdgeConstraintInfo);

  // *******************************************************************
  //                             Loops
  // *******************************************************************
  std::map<Loop *, unsigned> LoopNestingHeight;
  ValSet AllLoopPhis;
  ValToIslMapV ProcessedLoopPhis;
  std::map<PHINode *,
           std::tuple<isl_map *, Loop *, SmallVector<PHINode *, 4>, bool>>
      LoopPowMapInfo;
  int calculateLoopHeight(Loop *L);
  bool processLoop(Loop *L);
  void gatherLoopParams(Loop *L, ValSet &LoopPhisSet, ValSet &LoopParams);
  isl_space *getLoopPhiParamSpace(SmallVector<PHINode *, 4>  &LoopPhisSet);
  isl_map *getLoopPhiTransForBlocks(Loop *L,
                                    SmallVector<PHINode *, 4> &LoopPhiVec,
                                    ValToIslMapV &MapInfo, ValSet &LoopParams,
                                    SmallVector<BasicBlock *, 4> &LatchBlocks,
                                    EdgeToConstrTy &LoopEdgeConstraintInfo);
  isl_map *convertLoopPhiParamsToDims(isl_map *M,
                                      SmallVector<PHINode *, 4> &LoopPhis);
  isl_map *removeInnerLoopsN(isl_map *M, Loop *L);
  isl_map *removeMemoryParams(isl_map *M);
  isl_set *getLoopNSet(Loop *L);
  isl_map *getZeroIterTrans(Loop *L, isl_space *PowMapSpace);
  isl_map *convertParamToDim(isl_map *M, ValSet &ParamsForThisLoop,
                             ValToIslMapV &LoopFoGMapInfo,
                             std::set<void *> &ParamsWhiteList);
  IslMapVec convertParamToDim(IslMapVec MpVec, ValSet &ParamsForThisLoop,
                              ValToIslMapV &LoopFoGMapInfo,
                              std::set<void *> &ParamsWhiteList);
  void processInnerLoopPhiTrans(PHINode *FirstPhi, ValSet &ParamsForThisLoop,
                                ValToIslMapV &LoopFoGMapInfo,
                                EdgeToConstrTy &LoopEdgeConstraintInfo,
                                bool AddPseudoPhis);
  isl_map *addPseudoParamsForPhis(SmallVector<PHINode *, 4> &InnerLoopPhis,
                                  isl_map *CombinedPhiTrans);
  isl_map *removePseudoPhiParams(ValSet &LoopPhis, isl_map *M);
  isl_map *elimPseudoPhis(isl_map* M);
  isl_map *getWeakRelationForPhi(Loop *L, SmallVector<PHINode *, 4> &LoopPhiVec,
                                 ValToIslMapV &MapInfo, ValSet &LoopParams,
                                 SmallVector<BasicBlock *, 4> &LatchBlocks,
                                 EdgeToConstrTy &LoopEdgeConstraintInfo);

  bool shouldBailout() {
    if (isl_ctx_get_operations(Ctx) > MaxNoIslOps)
      return true;
    return false;
  }

public:
  explicit ISLFunctionInfo(Function *F, LoopInfo *LI);
  ISLFunctionInfo(ISLFunctionInfo &&FI)
      : BBConstraintInfo(std::move(FI.BBConstraintInfo)),
        F(FI.F), LI(FI.LI), Ctx(FI.Ctx), FoGMapInfo(std::move(FI.FoGMapInfo)),
        IdInfo(std::move(FI.IdInfo)), LoopIdInfo(std::move(FI.LoopIdInfo)),
        PseudoPhiIdInfo(std::move(FI.PseudoPhiIdInfo)),
        FunctionParams(std::move(FI.FunctionParams)),
        MemoryParams(std::move(FI.MemoryParams)),
        EdgeConstraintInfo(std::move(FI.EdgeConstraintInfo)),
        ProcessedLoopPhis(std::move(FI.ProcessedLoopPhis)),
        LoopPowMapInfo(std::move(FI.LoopPowMapInfo)) {
    FI.Ctx = nullptr;
  }

  ISLFunctionInfo(const ISLFunctionInfo &FI) = delete;

  ~ISLFunctionInfo() { 
    if (Ctx) {
      for (auto &It : FoGMapInfo)
        for (auto *M : It.second)
          if (M)
            isl_map_free(M);

      // Free all the constraints.
      for (auto &It : BBConstraintInfo)
        It.second.free();

      for (auto &It : EdgeConstraintInfo)
        It.second.free();

      // Free the params
      for (auto &It : IdInfo)
        for (isl_id *Id : It.second)
          if (Id)
            isl_id_free(Id);

      // Free loop ids
      for (auto &It : LoopIdInfo)
        if (It.second)
          isl_id_free(It.second);

      // Free the pseudo params
      for (auto &It : PseudoPhiIdInfo)
        for (isl_id *Id : It.second)
          if (Id)
            isl_id_free(Id);

      for (auto &It : ProcessedLoopPhis)
        for (auto *M : It.second)
          if(M)
            isl_map_free(M);

      for (auto &It : LoopPowMapInfo) {
        auto &Tuple = It.second;
        isl_map *M = std::get<0>(Tuple);
        if (M)
          isl_map_free(M);
      }

      isl_ctx_free(Ctx);
    }
  }
  bool isBBDead(BasicBlock *BB, ISLFunctionInfo::BBToConstrTy &BBConstraintInfo);
};

} // end namespace llvm

#endif // LLVM_TRANSFORMS_IPO_ISLVALUEANALYSIS_H
