Change the ForInst induction variable to be a block argument of the body instead of the ForInst itself. This is a necessary step in converting ForInst into an operation.

PiperOrigin-RevId: 231064139

7 files changed, 25 insertions, 25 deletions

diff --git a/mlir/lib/Analysis/AffineAnalysis.cpp b/mlir/lib/Analysis/AffineAnalysis.cpp
index 1ecad8d4e90..a4d969bc203 100644
--- a/mlir/lib/Analysis/AffineAnalysis.cpp
+++ b/mlir/lib/Analysis/AffineAnalysis.cpp
@@ -555,7 +555,7 @@ void mlir::getReachableAffineApplyOps(
 // setExprStride(ArrayRef<int64_t> expr, int64_t stride)
 bool mlir::getIndexSet(ArrayRef<ForInst *> forInsts,
                        FlatAffineConstraints *domain) {
-  SmallVector<Value *, 4> indices(forInsts.begin(), forInsts.end());
+  auto indices = extractForInductionVars(forInsts);
   // Reset while associated Values in 'indices' to the domain.
   domain->reset(forInsts.size(), /*numSymbols=*/0, /*numLocals=*/0, indices);
   for (auto *forInst : forInsts) {
@@ -677,7 +677,7 @@ static void buildDimAndSymbolPositionMaps(
   auto updateValuePosMap = [&](ArrayRef<Value *> values, bool isSrc) {
     for (unsigned i = 0, e = values.size(); i < e; ++i) {
       auto *value = values[i];
-      if (!isa<ForInst>(values[i])) {
+      if (!isForInductionVar(values[i])) {
         assert(values[i]->isValidSymbol() &&
                "access operand has to be either a loop IV or a symbol");
         valuePosMap->addSymbolValue(value);
@@ -739,7 +739,7 @@ void initDependenceConstraints(const FlatAffineConstraints &srcDomain,
   // Set values for the symbolic identifier dimensions.
   auto setSymbolIds = [&](ArrayRef<Value *> values) {
     for (auto *value : values) {
-      if (!isa<ForInst>(value)) {
+      if (!isForInductionVar(value)) {
         assert(value->isValidSymbol() && "expected symbol");
         dependenceConstraints->setIdValue(valuePosMap.getSymPos(value), value);
       }
@@ -907,7 +907,7 @@ addMemRefAccessConstraints(const AffineValueMap &srcAccessMap,
   // Add equality constraints for any operands that are defined by constant ops.
   auto addEqForConstOperands = [&](ArrayRef<const Value *> operands) {
     for (unsigned i = 0, e = operands.size(); i < e; ++i) {
-      if (isa<ForInst>(operands[i]))
+      if (isForInductionVar(operands[i]))
         continue;
       auto *symbol = operands[i];
       assert(symbol->isValidSymbol());
@@ -976,8 +976,8 @@ static unsigned getNumCommonLoops(const FlatAffineConstraints &srcDomain,
       std::min(srcDomain.getNumDimIds(), dstDomain.getNumDimIds());
   unsigned numCommonLoops = 0;
   for (unsigned i = 0; i < minNumLoops; ++i) {
-    if (!isa<ForInst>(srcDomain.getIdValue(i)) ||
-        !isa<ForInst>(dstDomain.getIdValue(i)) ||
+    if (!isForInductionVar(srcDomain.getIdValue(i)) ||
+        !isForInductionVar(dstDomain.getIdValue(i)) ||
         srcDomain.getIdValue(i) != dstDomain.getIdValue(i))
       break;
     ++numCommonLoops;
@@ -998,8 +998,9 @@ static const Block *getCommonBlock(const MemRefAccess &srcAccess,
     return block;
   }
   auto *commonForValue = srcDomain.getIdValue(numCommonLoops - 1);
-  assert(isa<ForInst>(commonForValue));
-  return cast<ForInst>(commonForValue)->getBody();
+  auto *forInst = getForInductionVarOwner(commonForValue);
+  assert(forInst && "commonForValue was not an induction variable");
+  return forInst->getBody();
 }
 
 // Returns true if the ancestor operation instruction of 'srcAccess' appears
diff --git a/mlir/lib/Analysis/AffineStructures.cpp b/mlir/lib/Analysis/AffineStructures.cpp
index 268fbe0c9c6..7aa23bbe480 100644
--- a/mlir/lib/Analysis/AffineStructures.cpp
+++ b/mlir/lib/Analysis/AffineStructures.cpp
@@ -1251,7 +1251,7 @@ void FlatAffineConstraints::setDimSymbolSeparation(unsigned newSymbolCount) {
 bool FlatAffineConstraints::addForInstDomain(const ForInst &forInst) {
   unsigned pos;
   // Pre-condition for this method.
-  if (!findId(forInst, &pos)) {
+  if (!findId(*forInst.getInductionVar(), &pos)) {
     assert(0 && "Value not found");
     return false;
   }
diff --git a/mlir/lib/Analysis/Dominance.cpp b/mlir/lib/Analysis/Dominance.cpp
index b98efa73e54..5cdeebbdf4a 100644
--- a/mlir/lib/Analysis/Dominance.cpp
+++ b/mlir/lib/Analysis/Dominance.cpp
@@ -53,9 +53,9 @@ bool DominanceInfo::properlyDominates(const Block *a, const Block *b) {
   if (blockListA == blockListB)
     return DominatorTreeBase::properlyDominates(a, b);
 
-  // Otherwise, 'a' properly dominates 'b' if 'b' is defined in an
-  // IfInst/ForInst that (recursively) ends up being dominated by 'a'. Walk up
-  // the list of containers enclosing B.
+  // Otherwise, 'a' properly dominates 'b' if 'b' is defined in an instruction
+  // region that (recursively) ends up being dominated by 'a'. Walk up the list
+  // of containers enclosing B.
   Instruction *bAncestor;
   do {
     bAncestor = blockListB->getContainingInst();
@@ -106,11 +106,6 @@ bool DominanceInfo::properlyDominates(const Value *a, const Instruction *b) {
   if (auto *aInst = a->getDefiningInst())
     return properlyDominates(aInst, b);
 
-  // The induction variable of a ForInst properly dominantes its body, so we
-  // can just do a simple block dominance check.
-  if (auto *forInst = dyn_cast<ForInst>(a))
-    return dominates(forInst->getBody(), b->getBlock());
-
   // block arguments properly dominate all instructions in their own block, so
   // we use a dominates check here, not a properlyDominates check.
   return dominates(cast<BlockArgument>(a)->getOwner(), b->getBlock());
diff --git a/mlir/lib/Analysis/LoopAnalysis.cpp b/mlir/lib/Analysis/LoopAnalysis.cpp
index b154ebab105..640984bf866 100644
--- a/mlir/lib/Analysis/LoopAnalysis.cpp
+++ b/mlir/lib/Analysis/LoopAnalysis.cpp
@@ -125,7 +125,7 @@ uint64_t mlir::getLargestDivisorOfTripCount(const ForInst &forInst) {
 }
 
 bool mlir::isAccessInvariant(const Value &iv, const Value &index) {
-  assert(isa<ForInst>(iv) && "iv must be a ForInst");
+  assert(isForInductionVar(&iv) && "iv must be a ForInst");
   assert(index.getType().isa<IndexType>() && "index must be of IndexType");
   SmallVector<OperationInst *, 4> affineApplyOps;
   getReachableAffineApplyOps({const_cast<Value *>(&index)}, affineApplyOps);
@@ -288,8 +288,10 @@ bool mlir::isVectorizableLoopAlongFastestVaryingMemRefDim(
       [fastestVaryingDim](const ForInst &loop, const OperationInst &op) {
         auto load = op.dyn_cast<LoadOp>();
         auto store = op.dyn_cast<StoreOp>();
-        return load ? isContiguousAccess(loop, *load, fastestVaryingDim)
-                    : isContiguousAccess(loop, *store, fastestVaryingDim);
+        return load ? isContiguousAccess(*loop.getInductionVar(), *load,
+                                         fastestVaryingDim)
+                    : isContiguousAccess(*loop.getInductionVar(), *store,
+                                         fastestVaryingDim);
       });
   return isVectorizableLoopWithCond(loop, fun);
 }
diff --git a/mlir/lib/Analysis/SliceAnalysis.cpp b/mlir/lib/Analysis/SliceAnalysis.cpp
index a8cec771f0d..d16a7fcb1b3 100644
--- a/mlir/lib/Analysis/SliceAnalysis.cpp
+++ b/mlir/lib/Analysis/SliceAnalysis.cpp
@@ -64,7 +64,7 @@ void mlir::getForwardSlice(Instruction *inst,
       }
     }
   } else if (auto *forInst = dyn_cast<ForInst>(inst)) {
-    for (auto &u : forInst->getUses()) {
+    for (auto &u : forInst->getInductionVar()->getUses()) {
       auto *ownerInst = u.getOwner();
       if (forwardSlice->count(ownerInst) == 0) {
         getForwardSlice(ownerInst, forwardSlice, filter,
diff --git a/mlir/lib/Analysis/Utils.cpp b/mlir/lib/Analysis/Utils.cpp
index 39e58e8983c..939a2ede618 100644
--- a/mlir/lib/Analysis/Utils.cpp
+++ b/mlir/lib/Analysis/Utils.cpp
@@ -149,7 +149,8 @@ bool mlir::getMemRefRegion(OperationInst *opInst, unsigned loopDepth,
     // A rank 0 memref has a 0-d region.
     SmallVector<ForInst *, 4> ivs;
     getLoopIVs(*opInst, &ivs);
-    SmallVector<Value *, 4> regionSymbols(ivs.begin(), ivs.end());
+
+    SmallVector<Value *, 8> regionSymbols = extractForInductionVars(ivs);
     regionCst->reset(0, loopDepth, 0, regionSymbols);
     return true;
   }
@@ -172,7 +173,7 @@ bool mlir::getMemRefRegion(OperationInst *opInst, unsigned loopDepth,
   unsigned numSymbols = accessMap.getNumSymbols();
   // Add inequalties for loop lower/upper bounds.
   for (unsigned i = 0; i < numDims + numSymbols; ++i) {
-    if (auto *loop = dyn_cast<ForInst>(accessValueMap.getOperand(i))) {
+    if (auto *loop = getForInductionVarOwner(accessValueMap.getOperand(i))) {
       // Note that regionCst can now have more dimensions than accessMap if the
       // bounds expressions involve outer loops or other symbols.
       // TODO(bondhugula): rewrite this to use getInstIndexSet; this way
@@ -207,7 +208,7 @@ bool mlir::getMemRefRegion(OperationInst *opInst, unsigned loopDepth,
   outerIVs.resize(loopDepth);
   for (auto *operand : accessValueMap.getOperands()) {
     ForInst *iv;
-    if ((iv = dyn_cast<ForInst>(operand)) &&
+    if ((iv = getForInductionVarOwner(operand)) &&
         std::find(outerIVs.begin(), outerIVs.end(), iv) == outerIVs.end()) {
       regionCst->projectOut(operand);
     }
diff --git a/mlir/lib/Analysis/VectorAnalysis.cpp b/mlir/lib/Analysis/VectorAnalysis.cpp
index 37eed71508f..125020e92a3 100644
--- a/mlir/lib/Analysis/VectorAnalysis.cpp
+++ b/mlir/lib/Analysis/VectorAnalysis.cpp
@@ -113,7 +113,8 @@ static AffineMap makePermutationMap(
                                   getAffineConstantExpr(0, context));
   for (auto kvp : enclosingLoopToVectorDim) {
     assert(kvp.second < perm.size());
-    auto invariants = getInvariantAccesses(*kvp.first, unwrappedIndices);
+    auto invariants =
+        getInvariantAccesses(*kvp.first->getInductionVar(), unwrappedIndices);
     unsigned numIndices = unwrappedIndices.size();
     unsigned countInvariantIndices = 0;
     for (unsigned dim = 0; dim < numIndices; ++dim) {