Refactor common code getting memref access in getMemRefRegion - NFC

- use getAccessMap() instead of repeating it
- fold getMemRefRegion into MemRefRegion ctor (more natural, avoid heap
  allocation and unique_ptr where possible)

- change extractForInductionVars - MutableArrayRef -> ArrayRef for the
  arguments. Since the method is just returning copies of 'Value *', the client
  can't mutate the pointers themselves; it's fine to mutate the 'Value''s
  themselves, but that doesn't mutate the pointers to those.

- change the way extractForInductionVars returns (see b/123437690)

PiperOrigin-RevId: 232359277

4 files changed, 13 insertions, 9 deletions

diff --git a/mlir/lib/Transforms/DmaGeneration.cpp b/mlir/lib/Transforms/DmaGeneration.cpp
index 92ae3767098..40d90c3e27b 100644
--- a/mlir/lib/Transforms/DmaGeneration.cpp
+++ b/mlir/lib/Transforms/DmaGeneration.cpp
@@ -183,7 +183,8 @@ static bool getFullMemRefAsRegion(Instruction *opInst, unsigned numParamLoopIVs,
   SmallVector<OpPointer<AffineForOp>, 4> ivs;
   getLoopIVs(*opInst, &ivs);
   ivs.resize(numParamLoopIVs);
-  SmallVector<Value *, 4> symbols = extractForInductionVars(ivs);
+  SmallVector<Value *, 4> symbols;
+  extractForInductionVars(ivs, &symbols);
   regionCst->reset(rank, numParamLoopIVs, 0);
   regionCst->setIdValues(rank, rank + numParamLoopIVs, symbols);
 
@@ -576,8 +577,8 @@ uint64_t DmaGeneration::runOnBlock(Block::iterator begin, Block::iterator end) {
     }
 
     // Compute the MemRefRegion accessed.
-    auto region = getMemRefRegion(opInst, dmaDepth);
-    if (!region) {
+    auto region = std::make_unique<MemRefRegion>(opInst->getLoc());
+    if (!region->compute(opInst, dmaDepth)) {
       LLVM_DEBUG(llvm::dbgs()
                  << "Error obtaining memory region: semi-affine maps?\n");
       LLVM_DEBUG(llvm::dbgs() << "over-approximating to the entire memref\n");
diff --git a/mlir/lib/Transforms/LoopFusion.cpp b/mlir/lib/Transforms/LoopFusion.cpp
index d7d69e569e5..7a002168528 100644
--- a/mlir/lib/Transforms/LoopFusion.cpp
+++ b/mlir/lib/Transforms/LoopFusion.cpp
@@ -921,7 +921,8 @@ static Value *createPrivateMemRef(OpPointer<AffineForOp> forOp,
   unsigned rank = oldMemRefType.getRank();
 
   // Compute MemRefRegion for 'srcStoreOpInst' at depth 'dstLoopDepth'.
-  auto region = getMemRefRegion(srcStoreOpInst, dstLoopDepth);
+  MemRefRegion region(srcStoreOpInst->getLoc());
+  region.compute(srcStoreOpInst, dstLoopDepth);
   SmallVector<int64_t, 4> newShape;
   std::vector<SmallVector<int64_t, 4>> lbs;
   SmallVector<int64_t, 8> lbDivisors;
@@ -929,11 +930,11 @@ static Value *createPrivateMemRef(OpPointer<AffineForOp> forOp,
   // Query 'region' for 'newShape' and lower bounds of MemRefRegion accessed
   // by 'srcStoreOpInst' at depth 'dstLoopDepth'.
   Optional<int64_t> numElements =
-      region->getConstantBoundingSizeAndShape(&newShape, &lbs, &lbDivisors);
+      region.getConstantBoundingSizeAndShape(&newShape, &lbs, &lbDivisors);
   assert(numElements.hasValue() &&
          "non-constant number of elts in local buffer");
 
-  const FlatAffineConstraints *cst = region->getConstraints();
+  const FlatAffineConstraints *cst = region.getConstraints();
   // 'outerIVs' holds the values that this memory region is symbolic/paramteric
   // on; this would correspond to loop IVs surrounding the level at which the
   // slice is being materialized.
diff --git a/mlir/lib/Transforms/LoopTiling.cpp b/mlir/lib/Transforms/LoopTiling.cpp
index 8b368e5f182..758d434d25e 100644
--- a/mlir/lib/Transforms/LoopTiling.cpp
+++ b/mlir/lib/Transforms/LoopTiling.cpp
@@ -201,7 +201,8 @@ UtilResult mlir::tileCodeGen(MutableArrayRef<OpPointer<AffineForOp>> band,
   // Move the loop body of the original nest to the new one.
   moveLoopBody(origLoops[origLoops.size() - 1], innermostPointLoop);
 
-  SmallVector<Value *, 8> origLoopIVs = extractForInductionVars(band);
+  SmallVector<Value *, 8> origLoopIVs;
+  extractForInductionVars(band, &origLoopIVs);
   SmallVector<Optional<Value *>, 6> ids(origLoopIVs.begin(), origLoopIVs.end());
   FlatAffineConstraints cst;
   getIndexSet(band, &cst);
diff --git a/mlir/lib/Transforms/MemRefDataFlowOpt.cpp b/mlir/lib/Transforms/MemRefDataFlowOpt.cpp
index b2b69dc7b6d..2d06a327315 100644
--- a/mlir/lib/Transforms/MemRefDataFlowOpt.cpp
+++ b/mlir/lib/Transforms/MemRefDataFlowOpt.cpp
@@ -178,8 +178,9 @@ void MemRefDataFlowOpt::visitInstruction(Instruction *opInst) {
       // is trivially loading from a single location at that depth; so there
       // isn't a need to call isRangeOneToOne.
       if (getNestingDepth(*storeOpInst) < loadOpDepth) {
-        auto region = getMemRefRegion(loadOpInst, nsLoops);
-        if (!region->getConstraints()->isRangeOneToOne(
+        MemRefRegion region(loadOpInst->getLoc());
+        region.compute(loadOpInst, nsLoops);
+        if (!region.getConstraints()->isRangeOneToOne(
                 /*start=*/0, /*limit=*/loadOp->getMemRefType().getRank()))
           break;
       }