Commit d1e9c7b6 authored by max's avatar max
Browse files

Revert "[mlir][affine] implement `promoteIfSingleIteration` for `AffineForOp` (#72547)" 

This reverts commit aa6be2f7.
parent 5603bb5f

mlir/include/mlir/Dialect/Affine/Analysis/LoopAnalysis.h

+14 −1
Original line number Diff line number Diff line
@@ -15,7 +15,6 @@ 


#include "mlir/Support/LLVM.h"

#include "llvm/ADT/ArrayRef.h"



#include <optional>



namespace mlir {
@@ -30,6 +29,20 @@ namespace affine { 
class AffineForOp;

class NestedPattern;



/// Returns the trip count of the loop as an affine map with its corresponding

/// operands if the latter is expressible as an affine expression, and nullptr

/// otherwise. This method always succeeds as long as the lower bound is not a

/// multi-result map. The trip count expression is simplified before returning.

/// This method only utilizes map composition to construct lower and upper

/// bounds before computing the trip count expressions

void getTripCountMapAndOperands(AffineForOp forOp, AffineMap *map,

                                SmallVectorImpl<Value> *operands);



/// Returns the trip count of the loop if it's a constant, std::nullopt

/// otherwise. This uses affine expression analysis and is able to determine

/// constant trip count in non-trivial cases.

std::optional<uint64_t> getConstantTripCount(AffineForOp forOp);



/// Returns the greatest known integral divisor of the trip count. Affine

/// expression analysis is used (indirectly through getTripCount), and

/// this method is thus able to determine non-trivial divisors.

mlir/include/mlir/Dialect/Affine/IR/AffineOps.h

+4 −25
Original line number Diff line number Diff line
@@ -117,8 +117,7 @@ public: 
  /// Returns the affine map used to access the source memref.

  AffineMap getSrcMap() { return getSrcMapAttr().getValue(); }

  AffineMapAttr getSrcMapAttr() {

    return cast<AffineMapAttr>(

        *(*this)->getInherentAttr(getSrcMapAttrStrName()));

    return cast<AffineMapAttr>(*(*this)->getInherentAttr(getSrcMapAttrStrName()));

  }



  /// Returns the source memref affine map indices for this DMA operation.
@@ -157,8 +156,7 @@ public: 
  /// Returns the affine map used to access the destination memref.

  AffineMap getDstMap() { return getDstMapAttr().getValue(); }

  AffineMapAttr getDstMapAttr() {

    return cast<AffineMapAttr>(

        *(*this)->getInherentAttr(getDstMapAttrStrName()));

    return cast<AffineMapAttr>(*(*this)->getInherentAttr(getDstMapAttrStrName()));

  }



  /// Returns the destination memref indices for this DMA operation.
@@ -187,8 +185,7 @@ public: 
  /// Returns the affine map used to access the tag memref.

  AffineMap getTagMap() { return getTagMapAttr().getValue(); }

  AffineMapAttr getTagMapAttr() {

    return cast<AffineMapAttr>(

        *(*this)->getInherentAttr(getTagMapAttrStrName()));

    return cast<AffineMapAttr>(*(*this)->getInherentAttr(getTagMapAttrStrName()));

  }



  /// Returns the tag memref indices for this DMA operation.
@@ -310,8 +307,7 @@ public: 
  /// Returns the affine map used to access the tag memref.

  AffineMap getTagMap() { return getTagMapAttr().getValue(); }

  AffineMapAttr getTagMapAttr() {

    return cast<AffineMapAttr>(

        *(*this)->getInherentAttr(getTagMapAttrStrName()));

    return cast<AffineMapAttr>(*(*this)->getInherentAttr(getTagMapAttrStrName()));

  }



  /// Returns the tag memref index for this DMA operation.
@@ -469,23 +465,6 @@ AffineForOp getForInductionVarOwner(Value val); 
/// AffineParallelOp.

AffineParallelOp getAffineParallelInductionVarOwner(Value val);



/// Helper to replace uses of loop carried values (iter_args) and loop

/// yield values while promoting single iteration affine.for ops.

void replaceIterArgsAndYieldResults(AffineForOp forOp);



/// Returns the trip count of the loop as an affine expression if the latter is

/// expressible as an affine expression, and nullptr otherwise. The trip count

/// expression is simplified before returning. This method only utilizes map

/// composition to construct lower and upper bounds before computing the trip

/// count expressions.

void getTripCountMapAndOperands(AffineForOp forOp, AffineMap *tripCountMap,

                                SmallVectorImpl<Value> *tripCountOperands);



/// Returns the trip count of the loop if it's a constant, std::nullopt

/// otherwise. This uses affine expression analysis and is able to determine

/// constant trip count in non-trivial cases.

std::optional<uint64_t> getConstantTripCount(AffineForOp forOp);



/// Extracts the induction variables from a list of AffineForOps and places them

/// in the output argument `ivs`.

void extractForInductionVars(ArrayRef<AffineForOp> forInsts,

mlir/include/mlir/Dialect/Affine/IR/AffineOps.td

+1 −1
Original line number Diff line number Diff line
@@ -121,7 +121,7 @@ def AffineForOp : Affine_Op<"for", 
     ImplicitAffineTerminator, ConditionallySpeculatable,

     RecursiveMemoryEffects, DeclareOpInterfaceMethods<LoopLikeOpInterface,

     ["getSingleInductionVar", "getSingleLowerBound", "getSingleStep",

      "getSingleUpperBound", "getYieldedValuesMutable", "promoteIfSingleIteration",

      "getSingleUpperBound", "getYieldedValuesMutable",

      "replaceWithAdditionalYields"]>,

     DeclareOpInterfaceMethods<RegionBranchOpInterface,

     ["getEntrySuccessorOperands"]>]> {

mlir/include/mlir/Dialect/Affine/LoopUtils.h

+4 −0
Original line number Diff line number Diff line
@@ -83,6 +83,10 @@ LogicalResult loopUnrollJamByFactor(AffineForOp forOp, 
LogicalResult loopUnrollJamUpToFactor(AffineForOp forOp,

                                      uint64_t unrollJamFactor);



/// Promotes the loop body of a AffineForOp to its containing block if the loop

/// was known to have a single iteration.

LogicalResult promoteIfSingleIteration(AffineForOp forOp);



/// Promotes all single iteration AffineForOp's in the Function, i.e., moves

/// their body into the containing Block.

void promoteSingleIterationLoops(func::FuncOp f);

mlir/lib/Dialect/Affine/Analysis/LoopAnalysis.cpp

+79 −1
Original line number Diff line number Diff line
@@ -12,6 +12,7 @@ 


#include "mlir/Dialect/Affine/Analysis/LoopAnalysis.h"



#include "mlir/Analysis/SliceAnalysis.h"

#include "mlir/Dialect/Affine/Analysis/AffineAnalysis.h"

#include "mlir/Dialect/Affine/Analysis/AffineStructures.h"

#include "mlir/Dialect/Affine/Analysis/NestedMatcher.h"
@@ -19,9 +20,9 @@ 
#include "mlir/Dialect/Affine/IR/AffineValueMap.h"

#include "mlir/Support/MathExtras.h"



#include "llvm/ADT/DenseSet.h"

#include "llvm/ADT/SmallPtrSet.h"

#include "llvm/ADT/SmallString.h"



#include <numeric>

#include <optional>

#include <type_traits>
@@ -29,6 +30,83 @@ 
using namespace mlir;

using namespace mlir::affine;



/// Returns the trip count of the loop as an affine expression if the latter is

/// expressible as an affine expression, and nullptr otherwise. The trip count

/// expression is simplified before returning. This method only utilizes map

/// composition to construct lower and upper bounds before computing the trip

/// count expressions.

void mlir::affine::getTripCountMapAndOperands(

    AffineForOp forOp, AffineMap *tripCountMap,

    SmallVectorImpl<Value> *tripCountOperands) {

  MLIRContext *context = forOp.getContext();

  int64_t step = forOp.getStepAsInt();

  int64_t loopSpan;

  if (forOp.hasConstantBounds()) {

    int64_t lb = forOp.getConstantLowerBound();

    int64_t ub = forOp.getConstantUpperBound();

    loopSpan = ub - lb;

    if (loopSpan < 0)

      loopSpan = 0;

    *tripCountMap = AffineMap::getConstantMap(ceilDiv(loopSpan, step), context);

    tripCountOperands->clear();

    return;

  }

  auto lbMap = forOp.getLowerBoundMap();

  auto ubMap = forOp.getUpperBoundMap();

  if (lbMap.getNumResults() != 1) {

    *tripCountMap = AffineMap();

    return;

  }



  // Difference of each upper bound expression from the single lower bound

  // expression (divided by the step) provides the expressions for the trip

  // count map.

  AffineValueMap ubValueMap(ubMap, forOp.getUpperBoundOperands());



  SmallVector<AffineExpr, 4> lbSplatExpr(ubValueMap.getNumResults(),

                                         lbMap.getResult(0));

  auto lbMapSplat = AffineMap::get(lbMap.getNumDims(), lbMap.getNumSymbols(),

                                   lbSplatExpr, context);

  AffineValueMap lbSplatValueMap(lbMapSplat, forOp.getLowerBoundOperands());



  AffineValueMap tripCountValueMap;

  AffineValueMap::difference(ubValueMap, lbSplatValueMap, &tripCountValueMap);

  for (unsigned i = 0, e = tripCountValueMap.getNumResults(); i < e; ++i)

    tripCountValueMap.setResult(i,

                                tripCountValueMap.getResult(i).ceilDiv(step));



  *tripCountMap = tripCountValueMap.getAffineMap();

  tripCountOperands->assign(tripCountValueMap.getOperands().begin(),

                            tripCountValueMap.getOperands().end());

}



/// Returns the trip count of the loop if it's a constant, std::nullopt

/// otherwise. This method uses affine expression analysis (in turn using

/// getTripCount) and is able to determine constant trip count in non-trivial

/// cases.

std::optional<uint64_t> mlir::affine::getConstantTripCount(AffineForOp forOp) {

  SmallVector<Value, 4> operands;

  AffineMap map;

  getTripCountMapAndOperands(forOp, &map, &operands);



  if (!map)

    return std::nullopt;



  // Take the min if all trip counts are constant.

  std::optional<uint64_t> tripCount;

  for (auto resultExpr : map.getResults()) {

    if (auto constExpr = dyn_cast<AffineConstantExpr>(resultExpr)) {

      if (tripCount.has_value())

        tripCount =

            std::min(*tripCount, static_cast<uint64_t>(constExpr.getValue()));

      else

        tripCount = constExpr.getValue();

    } else

      return std::nullopt;

  }

  return tripCount;

}



/// Returns the greatest known integral divisor of the trip count. Affine

/// expression analysis is used (indirectly through getTripCount), and

/// this method is thus able to determine non-trivial divisors.