Deprecate linalg.subview in favor of std.subview

This CL uses the now standard std.subview in linalg.
Two shortcuts are currently taken to allow this port:
1. the type resulting from a view is currently degraded to fully dynamic to pass the SubViewOp verifier.
2. indexing into SubViewOp may access out of bounds since lowering to LLVM does not currently enforce it by construction.

These will be fixed in subsequent commits after discussions.

PiperOrigin-RevId: 280250129

4 files changed, 59 insertions, 93 deletions

diff --git a/mlir/lib/Dialect/Linalg/Transforms/Fusion.cpp b/mlir/lib/Dialect/Linalg/Transforms/Fusion.cpp
index c6dffc3ab11..453daba204c 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/Fusion.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/Fusion.cpp
@@ -74,9 +74,8 @@ static llvm::cl::list<unsigned> clTileSizes(
 // a subset of the original loop ranges of `op`.
 // This is achieved by applying the `loopToOperandRangesMaps` permutation maps
 // to the `loopRanges` in order to obtain view ranges.
-static LinalgOp
-cloneWithLoopRanges(OpBuilder &b, Location loc, LinalgOp op,
-                    ArrayRef<mlir::linalg::SubViewOp::Range> loopRanges) {
+static LinalgOp cloneWithLoopRanges(OpBuilder &b, Location loc, LinalgOp op,
+                                    ArrayRef<SubViewOp::Range> loopRanges) {
   auto maps = loopToOperandRangesMaps(op);
   SmallVector<Value *, 8> clonedViews;
   clonedViews.reserve(op.getNumInputsAndOutputs());
@@ -88,8 +87,7 @@ cloneWithLoopRanges(OpBuilder &b, Location loc, LinalgOp op,
     auto map = maps[idx];
     LLVM_DEBUG(dbgs() << "map: " << map << "\n");
     Value *view = en.value();
-    SmallVector<mlir::linalg::SubViewOp::Range, 8> viewRanges(
-        map.getNumResults());
+    SmallVector<SubViewOp::Range, 4> viewRanges(map.getNumResults());
     for (auto en2 : llvm::enumerate(map.getResults())) {
       unsigned d = en2.index();
       // loopToOperandRangesMaps are permutations-only.
@@ -99,16 +97,19 @@ cloneWithLoopRanges(OpBuilder &b, Location loc, LinalgOp op,
                         << "\t"
                         << "loopPos: " << loopPos << "\t" << viewRanges[d]);
     }
-    // TODO(ntv): opportunities for folding/CSE here rather than build new IR.
-    SmallVector<Value *, 12> subViewOperands;
-    subViewOperands.reserve(viewRanges.size() * 3);
+    // Construct a new subview for the tile.
+    unsigned rank = viewRanges.size();
+    SmallVector<Value *, 4> offsets, sizes, strides;
+    offsets.reserve(rank);
+    sizes.reserve(rank);
+    strides.reserve(rank);
     for (auto r : viewRanges) {
-      subViewOperands.push_back(r.min);
-      subViewOperands.push_back(r.max);
-      subViewOperands.push_back(r.step);
+      offsets.push_back(r.offset);
+      sizes.push_back(r.size);
+      strides.push_back(r.stride);
     }
     clonedViews.push_back(
-        b.create<mlir::linalg::SubViewOp>(loc, view, subViewOperands));
+        b.create<SubViewOp>(loc, view, offsets, sizes, strides));
   }
   auto operands = getAssumedNonViewOperands(op);
   clonedViews.append(operands.begin(), operands.end());
@@ -153,7 +154,7 @@ static ViewDimension getViewDefiningLoopRange(LinalgOp op, unsigned loopDepth) {
 static LinalgOp fuse(Value *producedView, LinalgOp producer, LinalgOp consumer,
                      unsigned consumerIdx, unsigned producerIdx,
                      OperationFolder *folder) {
-  auto subView = dyn_cast_or_null<mlir::linalg::SubViewOp>(
+  auto subView = dyn_cast_or_null<SubViewOp>(
       consumer.getInput(consumerIdx)->getDefiningOp());
   auto slice = dyn_cast_or_null<SliceOp>(
       consumer.getInput(consumerIdx)->getDefiningOp());
@@ -172,13 +173,13 @@ static LinalgOp fuse(Value *producedView, LinalgOp producer, LinalgOp consumer,
   unsigned nPar = producer.getNumParallelLoops();
   unsigned nRed = producer.getNumReductionLoops();
   unsigned nWin = producer.getNumWindowLoops();
-  SmallVector<mlir::linalg::SubViewOp::Range, 8> loopRanges(nPar + nRed + nWin);
+  SmallVector<SubViewOp::Range, 8> loopRanges(nPar + nRed + nWin);
 
   // Iterate over dimensions identified by the producer map for `producerIdx`.
   // This defines a subset of the loop ranges that we need to complete later.
   for (auto en : llvm::enumerate(producerMap.getResults())) {
     unsigned posInProducerLoop = en.value().cast<AffineDimExpr>().getPosition();
-    loopRanges[posInProducerLoop] = subView.getRange(en.index());
+    loopRanges[posInProducerLoop] = subView.getRanges()[en.index()];
   }
 
   OpBuilder b(consumer.getOperation());
@@ -187,14 +188,14 @@ static LinalgOp fuse(Value *producedView, LinalgOp producer, LinalgOp consumer,
   // producer map for `producerIdx`, we need to explicitly compute the view that
   // defines the loop ranges using the `producer`.
   for (unsigned i = 0, nLoops = loopRanges.size(); i < nLoops; ++i) {
-    if (loopRanges[i].min)
+    if (loopRanges[i].offset)
       LLVM_DEBUG(llvm::dbgs()
                  << "existing LoopRange: " << loopRanges[i] << "\n");
     else {
       auto viewDim = getViewDefiningLoopRange(producer, i);
-      loopRanges[i] = mlir::linalg::SubViewOp::Range{
-          constant_index(folder, 0), dim(viewDim.view, viewDim.dimension),
-          constant_index(folder, 1)};
+      loopRanges[i] = SubViewOp::Range{constant_index(folder, 0),
+                                       dim(viewDim.view, viewDim.dimension),
+                                       constant_index(folder, 1)};
       LLVM_DEBUG(llvm::dbgs() << "new LoopRange: " << loopRanges[i] << "\n");
     }
   }
@@ -286,8 +287,7 @@ Optional<FusionInfo> mlir::linalg::fuseProducerOf(
 
     // Must be a subview or a slice to guarantee there are loops we can fuse
     // into.
-    auto subView = dyn_cast_or_null<mlir::linalg::SubViewOp>(
-        consumedView->getDefiningOp());
+    auto subView = dyn_cast_or_null<SubViewOp>(consumedView->getDefiningOp());
     auto slice = dyn_cast_or_null<SliceOp>(consumedView->getDefiningOp());
     if (!subView && !slice) {
       LLVM_DEBUG(dbgs() << "\nNot fusable (not a subview or slice)");
diff --git a/mlir/lib/Dialect/Linalg/Transforms/LowerToLLVMDialect.cpp b/mlir/lib/Dialect/Linalg/Transforms/LowerToLLVMDialect.cpp
index 05f4bced063..9d0395346af 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/LowerToLLVMDialect.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/LowerToLLVMDialect.cpp
@@ -503,24 +503,6 @@ public:
   }
 };
 
-/// A non-conversion rewrite pattern kicks in to convert SubViewOp into RangeOps
-/// and SliceOps.
-class SubViewOpConversion : public OpRewritePattern<mlir::linalg::SubViewOp> {
-public:
-  using OpRewritePattern<mlir::linalg::SubViewOp>::OpRewritePattern;
-
-  PatternMatchResult matchAndRewrite(mlir::linalg::SubViewOp op,
-                                     PatternRewriter &rewriter) const override {
-    auto *view = op.getView();
-    SmallVector<Value *, 8> ranges;
-    for (auto sliceRange : op.getRanges())
-      ranges.push_back(rewriter.create<RangeOp>(
-          op.getLoc(), sliceRange.min, sliceRange.max, sliceRange.step));
-    rewriter.replaceOpWithNewOp<SliceOp>(op, view, ranges);
-    return matchSuccess();
-  }
-};
-
 /// Populate the given list with patterns that convert from Linalg to Standard.
 static void
 populateLinalgToStandardConversionPatterns(OwningRewritePatternList &patterns,
@@ -530,8 +512,8 @@ populateLinalgToStandardConversionPatterns(OwningRewritePatternList &patterns,
   patterns.insert<CopyTransposeConversion, LinalgOpConversion<CopyOp>,
                   LinalgOpConversion<DotOp>, LinalgOpConversion<FillOp>,
                   LinalgOpConversion<MatvecOp>, LinalgOpConversion<MatmulOp>,
-                  LinalgOpConversion<ConvOp>, LinalgOpConversion<GenericOp>,
-                  SubViewOpConversion>(ctx);
+                  LinalgOpConversion<ConvOp>, LinalgOpConversion<GenericOp>>(
+      ctx);
 }
 
 /// Populate the given list with patterns that convert from Linalg to LLVM.
diff --git a/mlir/lib/Dialect/Linalg/Transforms/Promotion.cpp b/mlir/lib/Dialect/Linalg/Transforms/Promotion.cpp
index 3afee415405..3c6e1395506 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/Promotion.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/Promotion.cpp
@@ -89,13 +89,13 @@ static Value *allocBuffer(Type elementType, Value *size, bool dynamicBuffers) {
 // boundary tiles. For now this is done with an unconditional `fill` op followed
 // by a partial `copy` op.
 static PromotionInfo promoteFullTileBuffer(OpBuilder &b, Location loc,
-                                           mlir::linalg::SubViewOp subView,
+                                           SubViewOp subView,
                                            bool dynamicBuffers,
                                            OperationFolder *folder) {
   auto zero = constant_index(folder, 0);
   auto one = constant_index(folder, 1);
 
-  auto viewType = subView.getViewType();
+  auto viewType = subView.getType();
   auto rank = viewType.getRank();
   Value *allocSize = one;
   SmallVector<Value *, 8> fullRanges, partialRanges;
@@ -104,12 +104,7 @@ static PromotionInfo promoteFullTileBuffer(OpBuilder &b, Location loc,
   for (auto en : llvm::enumerate(subView.getRanges())) {
     auto rank = en.index();
     auto rangeValue = en.value();
-    Value *d =
-        isa<DimOp>(rangeValue.max->getDefiningOp())
-            ? rangeValue.max
-            : applyMapToValues(b, loc, getAffineDifferenceMap(b.getContext()),
-                               {rangeValue.max, rangeValue.min}, folder)
-                  .front();
+    Value *d = rangeValue.size;
     allocSize = muli(folder, allocSize, d).getValue();
     fullRanges.push_back(d);
     partialRanges.push_back(range(folder, zero, dim(subView, rank), one));
@@ -135,9 +130,8 @@ mlir::linalg::promoteSubViews(OpBuilder &b, Location loc,
   res.reserve(subViews.size());
   DenseMap<Value *, PromotionInfo> promotionInfoMap;
   for (auto *v : subViews) {
-    mlir::linalg::SubViewOp subView =
-        cast<mlir::linalg::SubViewOp>(v->getDefiningOp());
-    auto viewType = subView.getViewType();
+    SubViewOp subView = cast<SubViewOp>(v->getDefiningOp());
+    auto viewType = subView.getType();
     // TODO(ntv): support more cases than just float.
     if (!viewType.getElementType().isa<FloatType>())
       continue;
@@ -148,14 +142,13 @@ mlir::linalg::promoteSubViews(OpBuilder &b, Location loc,
   }
 
   for (auto *v : subViews) {
-    mlir::linalg::SubViewOp subView =
-        cast<mlir::linalg::SubViewOp>(v->getDefiningOp());
+    SubViewOp subView = cast<SubViewOp>(v->getDefiningOp());
     auto info = promotionInfoMap.find(v);
     if (info == promotionInfoMap.end())
       continue;
     // TODO(ntv): value to fill with should be related to the operation.
     // For now, just use APFloat(0.0f).
-    auto t = subView.getViewType().getElementType().cast<FloatType>();
+    auto t = subView.getType().getElementType().cast<FloatType>();
     Value *fillVal = constant_float(folder, APFloat(0.0f), t);
     // TODO(ntv): fill is only necessary if `promotionInfo` has a full local
     // view that is different from the partial local view and we are on the
@@ -167,8 +160,7 @@ mlir::linalg::promoteSubViews(OpBuilder &b, Location loc,
     auto info = promotionInfoMap.find(v);
     if (info == promotionInfoMap.end())
       continue;
-    copy(cast<mlir::linalg::SubViewOp>(v->getDefiningOp()),
-         info->second.partialLocalView);
+    copy(cast<SubViewOp>(v->getDefiningOp()), info->second.partialLocalView);
   }
   return res;
 }
@@ -226,8 +218,7 @@ static void promoteSubViews(FuncOp f, bool dynamicBuffers) {
     // nothing.
     SetVector<Value *> subViews;
     for (auto it : op.getInputsAndOutputs())
-      if (auto sv =
-              dyn_cast_or_null<mlir::linalg::SubViewOp>(it->getDefiningOp()))
+      if (auto sv = dyn_cast_or_null<SubViewOp>(it->getDefiningOp()))
         subViews.insert(sv);
     if (!subViews.empty()) {
       promoteSubViewOperands(op, subViews, dynamicBuffers, &folder);
diff --git a/mlir/lib/Dialect/Linalg/Transforms/Tiling.cpp b/mlir/lib/Dialect/Linalg/Transforms/Tiling.cpp
index b7a5740a387..76b1858b333 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/Tiling.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/Tiling.cpp
@@ -65,7 +65,7 @@ static bool isZero(Value *v) {
 // avoiding affine map manipulations.
 // The returned ranges correspond to the loop ranges, in the proper order, that
 // are tiled and for which new loops will be created.
-static SmallVector<mlir::linalg::SubViewOp::Range, 4>
+static SmallVector<SubViewOp::Range, 4>
 makeTiledLoopRanges(OpBuilder &b, Location loc, AffineMap map,
                     ArrayRef<Value *> allViewSizes,
                     ArrayRef<Value *> allTileSizes, OperationFolder *folder) {
@@ -83,10 +83,10 @@ makeTiledLoopRanges(OpBuilder &b, Location loc, AffineMap map,
   }
 
   // Create a new range with the applied tile sizes.
-  SmallVector<mlir::linalg::SubViewOp::Range, 4> res;
+  SmallVector<SubViewOp::Range, 4> res;
   for (unsigned idx = 0, e = tileSizes.size(); idx < e; ++idx) {
-    res.push_back(mlir::linalg::SubViewOp::Range{
-        constant_index(folder, 0), viewSizes[idx], tileSizes[idx]});
+    res.push_back(SubViewOp::Range{constant_index(folder, 0), viewSizes[idx],
+                                   tileSizes[idx]});
   }
   return res;
 }
@@ -153,16 +153,11 @@ makeTiledViews(OpBuilder &b, Location loc, LinalgOp linalgOp,
 
   // Construct (potentially temporary) mins and maxes on which to apply maps
   // that define tile subviews.
-  SmallVector<Value *, 8> mins, maxes;
+  SmallVector<Value *, 8> lbs, subViewSizes;
   for (unsigned idx = 0, idxIvs = 0, e = tileSizes.size(); idx < e; ++idx) {
-    if (isZero(tileSizes[idx])) {
-      mins.push_back(constant_index(folder, 0));
-      maxes.push_back(viewSizes[idx]);
-    } else {
-      ValueHandle lb(ivs[idxIvs++]), step(tileSizes[idx]);
-      mins.push_back(lb);
-      maxes.push_back(lb + step);
-    }
+    bool isTiled = !isZero(tileSizes[idx]);
+    lbs.push_back(isTiled ? ivs[idxIvs++] : (Value *)constant_index(folder, 0));
+    subViewSizes.push_back(isTiled ? tileSizes[idx] : viewSizes[idx]);
   }
 
   auto *op = linalgOp.getOperation();
@@ -182,41 +177,39 @@ makeTiledViews(OpBuilder &b, Location loc, LinalgOp linalgOp,
     }
 
     // Construct a new subview for the tile.
-    SmallVector<mlir::linalg::SubViewOp::Range, 4> subViewRangeOperands;
-    subViewRangeOperands.reserve(rank * 3);
+    SmallVector<Value *, 4> offsets, sizes, strides;
+    offsets.reserve(rank);
+    sizes.reserve(rank);
+    strides.reserve(rank);
     for (unsigned r = 0; r < rank; ++r) {
       if (!isTiled(map.getSubMap({r}), tileSizes)) {
-        subViewRangeOperands.push_back(mlir::linalg::SubViewOp::Range{
-            constant_index(folder, 0), dim(view, r),
-            constant_index(folder, 1)});
+        offsets.push_back(constant_index(folder, 0));
+        sizes.push_back(dim(view, r));
+        strides.push_back(constant_index(folder, 1));
         continue;
       }
 
-      auto m = map.getSubMap({r});
-      auto *min = applyMapToValues(b, loc, m, mins, folder).front();
-      auto *max = applyMapToValues(b, loc, m, maxes, folder).front();
       // Tiling creates a new slice at the proper index, the slice step is 1
       // (i.e. the slice view does not subsample, stepping occurs in the loop).
-      subViewRangeOperands.push_back(
-          mlir::linalg::SubViewOp::Range{min, max, constant_index(folder, 1)});
-    }
-    SmallVector<Value *, 12> subViewOperands;
-    subViewOperands.reserve(subViewRangeOperands.size() * 3);
-    for (auto r : subViewRangeOperands) {
-      subViewOperands.push_back(r.min);
-      subViewOperands.push_back(r.max);
-      subViewOperands.push_back(r.step);
+      auto m = map.getSubMap({r});
+      auto *offset = applyMapToValues(b, loc, m, lbs, folder).front();
+      offsets.push_back(offset);
+      auto *size = applyMapToValues(b, loc, m, subViewSizes, folder).front();
+      sizes.push_back(size);
+      strides.push_back(constant_index(folder, 1));
     }
-    res.push_back(
-        b.create<mlir::linalg::SubViewOp>(loc, view, subViewOperands));
+    // TODO(b/144419024) Atm std.subview is not guaranteed in-bounds. Depending
+    // on the semantics we attach to it, we may need to use min(size, dim) here
+    // and canonicalize later.
+    res.push_back(b.create<SubViewOp>(loc, view, offsets, sizes, strides));
   }
 
   // Traverse the mins/maxes and erase those that don't have uses left.
   // This is a special type of folding that we only apply when `folder` is
   // defined.
   if (folder) {
-    mins.append(maxes.begin(), maxes.end());
-    for (auto *v : mins)
+    lbs.append(subViewSizes.begin(), subViewSizes.end());
+    for (auto *v : lbs)
       if (v->use_empty())
         v->getDefiningOp()->erase();
   }