//===-- AffineTransforms.h - Matrix based operations ----------------------===//
//
//                     The LLVM Compiler Infrastructure
//
// This file is distributed under the University of Illinois Open Source
// License. See LICENSE.TXT for details.
//
//===----------------------------------------------------------------------===//
//
// Provide an interface to perform affine transformations.
// In particular, the transformation is Unimodular Transformation for
// the affine expressions extracted from loops. After transformation, loops are
// rewritten with new bounds and new access expressions.
//
//===----------------------------------------------------------------------===//

//===----------------------------------------------------------------------===//
//==== Copyright (c) 2017 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_UTILS_AFFINETRANSFORMS_H
#define LLVM_TRANSFORMS_UTILS_AFFINETRANSFORMS_H

#include "llvm/Analysis/ScalarEvolutionExpressions.h"
#include "llvm/Transforms/Utils/LoopUtils.h"
#include "llvm/Transforms/Utils/MatrixOperations.h"
#include "llvm/Transforms/Utils/ScalarEvolutionExpander.h"

namespace llvm {

// TODO: This typedef and other typedefs of LoopNest at other locations are to
// be merged to a single typedef at LoopUtils.
typedef SmallVector<Loop *, 2> LoopNest_t;

typedef SmallVector<Value *, 2> ValueList;
typedef SmallVector<PHINode *, 2> PHIList;
typedef SmallVector<Instruction *, 2> InstructionList;
typedef SmallVector<const SCEV *, 2> SCEVList;

//===----------------------------------------------------------------------===//
// Classes that analyze and manipulate SCEVs.
//===----------------------------------------------------------------------===//

// Convert a SCEV expression to an affine expression.
class SCEVAffinator : public SCEVVisitor<SCEVAffinator, bool> {
protected:
  LoopNest_t LN;
  ScalarEvolution *SE;

private:
  // The number of columns in "A" (and length of "a")
  // will be at least LoopNest.size(), i.e. one for
  // each indvar. Over that, it will have slots for
  // each parameter in "paramList".
  MatOp::RationalMatrix_t A;
  MatOp::RationalVector_t a;

  SCEVList paramList;
  // insert (if not already present) "p" and return index.
  int insertParam(const SCEV *p);

  bool visit(const SCEV *Expr);
  bool visitConstant(const SCEVConstant *Expr);
  bool visitVScale(const SCEVVScale *Expr);
  bool visitTruncateExpr(const SCEVTruncateExpr *Expr);
  bool visitZeroExtendExpr(const SCEVZeroExtendExpr *Expr);
  bool visitSignExtendExpr(const SCEVSignExtendExpr *Expr);
  bool visitAddExpr(const SCEVAddExpr *Expr);
  bool visitMulExpr(const SCEVMulExpr *Expr);
  bool visitUDivExpr(const SCEVUDivExpr *Expr);
  bool visitAddRecExpr(const SCEVAddRecExpr *Expr);
  bool visitSMaxExpr(const SCEVSMaxExpr *Expr);
  bool visitSMinExpr(const SCEVSMinExpr *Expr);
  bool visitUMaxExpr(const SCEVUMaxExpr *Expr);
  bool visitUMinExpr(const SCEVUMinExpr *Expr);
  bool visitSequentialUMinExpr(const SCEVSequentialUMinExpr *Expr);
  bool visitUnknown(const SCEVUnknown *Expr);
  bool visitPtrToIntExpr(const SCEVPtrToIntExpr *Expr);
  bool visitUnhandledSCEV(const SCEV *Expr);

  friend struct SCEVVisitor<SCEVAffinator, bool>;

public:
  SCEVAffinator(LoopNest_t &_LN, ScalarEvolution *_SE) : LN(_LN), SE(_SE) {}

  // return index of "p" if found, else -1.
  int findParam(const SCEV *p);

  // Given a bunch of SCEVs, form an affine expression for each of them.
  // upon return, Bi+b forms the affine expression, where "i" will be
  // the vector of canonical indvars + params.
  // 1. All the parameters are returned in "paramList" so that the caller
  //    can know the parameters found and their order assumed.
  // 2. The indvars assumed are implicit, i.e. one canonical indvar for each
  // loop
  //    in loopList, starting from depth 0.
  // All SCEVs that need to be analysed need to be passed together, so that
  // the parameters (and their ordering) returned remains consistent.
  // RETVAL: return true if an affine expression in indvars+params is found,
  // else
  //         return false. "B", "b" and "paramList" are undefined if returning
  //         false.
  // INPUT: Currently, only SCEVs of the below forms are supported (example):
  //   1.  AddRecExpr = {{(3 + (4 * %q)),+,4}<%outer.loop>,+,3}<%inner.loop>
  //     where, "q" is a parameter. There could be more such parameters at
  //     the "start" of the chain of recurrence. The multiplies for each indvar
  //     and parameter must be a compile time constant (else non-affine).
  //     Above expression would transloate to (3 + 4q) + 4i + 3j, assuming
  //     %outer.loop and %inner.loop have canonical indvars i and j
  //     respectively.
  //   2. In case of the AddRecExpr above, we can consider q as a parameter
  //      or (4 * q) as a parameter or the whole of (3+(4*q)) as a parameter.
  //      The current implemention design choice is to go as deep as possible
  //      (i.e., separate out parameters as much as possible).  I'm not sure
  //      if this is the right design. See "The Lambda Loop Transformation
  //      Toolkit - User's reference manual" -- Li and Pingali, Chapter 5.
  bool affinateSCEVs(std::vector<const SCEV *> &ipSCEVs,
         MatOp::RationalMatrix_t &B, MatOp::RationalVector_t &b,
         SCEVList &paramList, SmallVectorImpl<unsigned>& SCEVToRowTracker);
};

// Compute the affine bounds (Ax <= b) for each IndVar in a LoopNest.
class IndVarBounds : public SCEVAffinator {

public:
  IndVarBounds(LoopNest_t &LN, ScalarEvolution *SE) : SCEVAffinator(LN, SE) {}

  // Get loop bounds as Bi <= b where 'i'  is the vector of indvars + params
  // (in that order).
  // All the parameters are returned in "paramList" so that the caller
  // can know the parameters found and their order assumed.
  // RETVAL: return true if all bounds are affine in indvars+params. In this
  // case B,b and paramList are filled accordingly.
  // Otherwise return false and B,b and paramList are undefined.
  bool computeLoopBounds(const PHIList &ivl, MatOp::RationalMatrix_t &B,
                         MatOp::RationalVector_t &b, SCEVList &paramList);
};

// Compute affine access functions "Ai+a" for each array access computation
// (GEP) passed.
class ArrAccFuncs : public SCEVAffinator {

public:
  ArrAccFuncs(LoopNest_t &LN, ScalarEvolution *SE) : SCEVAffinator(LN, SE) {}

  // Try to delinearize the array access if necessary.
  // Subscripts and Sizes are computed if returning true.
  bool tryDelinearize(Instruction *LdSt, SCEVList &Subscripts, SCEVList &Sizes);

  // Get access functions as Ai+a where 'i'  is the vector of indvars + params
  // (in that order).
  // All the parameters are returned in "paramList" so that the caller
  // can know the parameters found and their order assumed.
  // RETVAL: return true if all accesses are affine in indvars+params. In this
  // case
  //         AList, aList and paramList are filled accordingly.
  //         Otherwise return false and AList, aList and paramList are
  //         undefined.
  // AList[i] and aList[i] will be the access function Ai+a for gAccs[i].
  bool computeArrAccs(const InstructionList &Accs,
                      std::vector<MatOp::RationalMatrix_t> &A,
                      std::vector<MatOp::RationalVector_t> &a,
                      SCEVList &paramList);
};

// Affine Transformations.
class AffineTransforms {
private:
  LoopNest_t LN;
  ScalarEvolution *SE;
  LoopInfo *LI;
  const DataLayout *DL;
  LLVMContext *Ctx;
  // The transformation matrix. The caller specifies this matrix.
  MatOp::RationalMatrix_t T;

  PHIList IndVars;    // List of indvars.
  ValueList CmpInsts; // List of loop compare insts.

  MatOp::RationalMatrix_t TInv; // Inverse of T.

  InstructionList Accs;
  // Loop accesses in Ai + a form.
  std::vector<MatOp::RationalMatrix_t> A;
  std::vector<MatOp::RationalVector_t> a;
  SCEVList accParams;

  // Loop bounds in Bi + b form.
  MatOp::RationalMatrix_t B;
  MatOp::RationalVector_t b;
  SCEVList boundParams;

  bool LegalityChecks();

  void getIndVars();
  void getCmpInsts();
  void getAccList();

  // This function performs fourier-motzkin elimination of variables,
  // starting from "firstProject" till 0.
  // assert (height(B) == height(b)).
  // The assumed canonical form is Bi <= b, where 'i' is the vector of
  // loop indices followed by loop invariants (but not constants). 'i'
  // here is implicit, and would be as long as width(B).
  // assert(firstProject <= width(B)). firstProject is considered to
  // be the last loop index variable. The rest are loop invariants, such
  // as 'n' that we commonly have as loop bound.
  // RETURN VALUE: return true if there is a solution to the set of bounds,
  //               i.e. there is at least one point in the polyhedra.
  // Computed value: The computed bounds (projected bounds) are returned
  // in pB and pb, where two consecutive slots are assigned for each index
  // variable, first for lower bounds and second for upper bounds. i.e.
  // upon returning, pB.size() == pb.size() == (2*firstProject). For example,
  // if i0 is the first index variable, then
  //     i0 >= pB[0](i) + pb[0] && i0 <= pB[1](i) + pb[1]
  // where 'i' is the set of index and invariant variables in the same order
  // as was implicit in the input B and b.
  bool doFourierMotzkin(const MatOp::RationalMatrix_t &B,
                        const MatOp::RationalVector_t &b, int firstProject,
                        std::vector<MatOp::RationalMatrix_t> &pB,
                        std::vector<MatOp::RationalVector_t> &pb);

  // SCEVAffinator computes parameters as SCEVs. Some of these are
  // trivially Values, while others need to be expanded into instructions.
  // Its likely that instructions to compute these SCEVs already exist
  // (such as indvar of an outer loop variable), but we have no way of
  // knowing them. So expand unconditionally and hope gvn fixes it.
  void expandParams(const SCEVList &boundParams,
                    ValueList &boundParamsExpanded);

  // It is possible that the bound looks like this:
  //   (1/2)i + (1/4)j +(1/2)
  // i.e. the denominators aren't common. normalizeBound() will
  // return a common factor out, and replace the const multiplies
  // in-place. i.e., for above example it will return "4" and
  // set the inputs (in-place) to:
  //   2i + j + 2
  long normalizeBound(MatOp::RationalVector_t &B, MatOp::Rational &b);

  // Generate a sum of products expressions, with each rational in V being the
  // multiplier for each indvar in newivl and each parameter in paramList.
  // NOTE: It is ok for newivl.size() < LoopNest.size(), but in that case, the
  //       corresponding multipliers in V must be 0.
  // REQ: Each element in V must be an integer.
  Value *generateSOP(const MatOp::RationalVector_t &V, const PHIList &newivl,
                     const ValueList &paramList, Instruction *insertBefore,
                     const char *nameBase);

  // Generate ceil(X / Y).
  Value *generateCeilingXbyY(Value *X, Value *Y, Instruction *insertBefore);

  // Generate floor(X / Y).
  Value *generateFloorXbyY(Value *X, Value *Y, Instruction *insertBefore);

  // Generate IR for the bounds specified.
  Value *generateBound(const MatOp::RationalMatrix_t &B,
                       const MatOp::RationalVector_t &b, const PHIList &newivl,
                       const ValueList &paramList, Instruction *insertBefore,
                       bool lowerOrUpper);

  // Assuming pB, pb contain the bounds for the transformed indices, and
  // paramList contains the parameter list assumed, generate LLVM IR for the new
  // indices and their bounds (including setting the branch condition).
  // LoopSimplify form is assumed.
  int generateNewBounds(const std::vector<MatOp::RationalMatrix_t> &pB,
                        const std::vector<MatOp::RationalVector_t> &pb,
                        const ValueList &paramList, PHIList &newivl,
                        InstructionList &newIncrList);

  // Every use of old indvars will now be replaced with TInv * new ones, for the
  // bounds computation. We may accidentally replace old bounds computation too,
  // but thats okay, its anyway already dead.
  // NOTE: Instead of doing (A*TInv)i', we are doing A(TInv * i'). This may lead
  // to two sets of multiplication+addition being performed.
  // TODO: Get A explicitly and do the faster one.
  void modifyAccessFunctions(const MatOp::RationalMatrix_t &TInv,
                             const PHIList &ivl, const PHIList &newivl);

  // Delete old indvars as they are no longer used.
  void DeleteOldIndVars(PHIList &ivl);

  // Print B, b and boundParams.
  void printBounds();

public:
  AffineTransforms(LoopNest_t LN, ScalarEvolution *SE, LoopInfo *LI,
                   const DataLayout *DL, LLVMContext *Ctx,
                   MatOp::RationalMatrix_t T)
                 : LN(LN), SE(SE), LI(LI), DL(DL), Ctx(Ctx), T(T) {}

  bool PerformTransformation();

  bool AnalyseAndTransform();
};
}
#endif
