[mlir][linalg] Cleanup LinalgOp usage in tiling (NFC).

[gysit]

Summary

Replace the uses of deprecated Structured Op Interface methods in Tiling.cpp and Utils.cpp. This patch is based on https://reviews.llvm.org/D103394.

mlir/lib/Dialect/Linalg/Transforms/Tiling.cpp

auto tiledLoopBodyBuilder = [&](OpBuilder &b, Location loc,
     // loop ranges and the iterator types. Apply its inverse to the
     // resulting loop `ivs` to match the op definition.
     SmallVector<Value, 4> interchangedIvs;
     if (!options.interchangeVector.empty())
       interchangedIvs = applyMapToValues(b, loc, invPermutationMap, ivs);
     else
       interchangedIvs.assign(ivs.begin(), ivs.end());
 
-    assert(op.getNumOutputTensors() == iterArgs.size() &&
+    assert(op.getOutputTensorOperands().size() == iterArgs.size() &&
            "num output tensors must match number of loop iter arguments");
 
-    auto operands = llvm::to_vector<4>(op.getInputs());
-    SmallVector<Value, 4> outputBuffers = op.getOutputBuffers();
+    SmallVector<Value> operands = op.getInputOperands();
+    SmallVector<Value> outputBuffers = op.getOutputBufferOperands();
     // TODO: thanks to simplifying assumption we do not need to worry about
     // order of output buffers and tensors: there is only ever one kind.
     assert(outputBuffers.empty() || iterArgs.empty());
     operands.append(outputBuffers.begin(), outputBuffers.end());
     operands.append(iterArgs.begin(), iterArgs.end());
     auto sizeBounds =
         applyMapToValues(b, loc, shapeSizesToLoopsMap, allShapeSizes);
     SmallVector<Value, 4> tiledOperands = makeTiledShapes(
         b, loc, op, operands, interchangedIvs, tileSizes, sizeBounds);
     auto nonShapedOperands = op.getAssumedNonShapedOperands();
     tiledOperands.append(nonShapedOperands.begin(), nonShapedOperands.end());
 
     // TODO: use an interface/adaptor to avoid leaking position in
     // `tiledOperands`.
     SmallVector<Type, 4> resultTensorTypes;
-    for (OpOperand *opOperand : op.getOutputTensorsOpOperands())
+    for (OpOperand *opOperand : op.getOutputTensorOperands())
       resultTensorTypes.push_back(
           tiledOperands[opOperand->getOperandNumber()].getType());
 
     res = op.clone(b, loc, resultTensorTypes, tiledOperands);
 
     // Insert a subtensor_insert for each output tensor.
     unsigned resultIdx = 0;
-    for (OpOperand *opOperand : op.getOutputTensorsOpOperands()) {
+    for (OpOperand *opOperand : op.getOutputTensorOperands()) {
       // TODO: use an interface/adaptor to avoid leaking position in
       // `tiledOperands`.
       Value outputTensor = tiledOperands[opOperand->getOperandNumber()];
       if (auto subtensor = outputTensor.getDefiningOp<SubTensorOp>()) {
         tensorResults.push_back(b.create<SubTensorInsertOp>(
             loc, subtensor.source().getType(), res->getResult(resultIdx),
             subtensor.source(), subtensor.offsets(), subtensor.sizes(),
             subtensor.strides(), subtensor.static_offsets(),

mlir/lib/Dialect/Linalg/Utils/Utils.cpp

 void GenerateLoopNest<scf::ForOp>::doit(
     OpBuilder &b, Location loc, ArrayRef<Range> loopRanges, LinalgOp linalgOp,
     ArrayRef<Attribute> iteratorTypes,
     function_ref<scf::ValueVector(OpBuilder &, Location, ValueRange,
                                   ValueRange)>
         bodyBuilderFn,
     Optional<LinalgLoopDistributionOptions> distributionOptions,
     ArrayRef<StringRef> distributionTypes) {
-  auto iterArgInitValues = linalgOp.getOutputTensors();
+  SmallVector<Value> iterArgInitValues = linalgOp.getOutputTensorOperands();
   // Create procInfo so it dominates loops, if appropriate.
   SmallVector<ProcInfo, 4> procInfo;
   SmallVector<DistributionMethod, 0> distributionMethod;
   if (distributionOptions.hasValue()) {
     // Collect loop ranges for parallel dimensions.
     SmallVector<Range, 2> parallelLoopRanges;
     for (auto iteratorType : enumerate(iteratorTypes))
       if (isParallelIteratorType(iteratorType.value()))

 template <>
 void GenerateLoopNest<AffineForOp>::doit(
     OpBuilder &b, Location loc, ArrayRef<Range> loopRanges, LinalgOp linalgOp,
     ArrayRef<Attribute> iteratorTypes,
     function_ref<scf::ValueVector(OpBuilder &, Location, ValueRange,
                                   ValueRange)>
         bodyBuilderFn,
     Optional<LinalgLoopDistributionOptions>, ArrayRef<StringRef>) {
-  auto iterArgInitValues = linalgOp.getOutputTensors();
+  SmallVector<Value> iterArgInitValues = linalgOp.getOutputTensorOperands();
   assert(iterArgInitValues.empty() && "unexpected AffineForOp init values");
   SmallVector<Value, 4> lbs, ubs, steps;
   unpackRanges(loopRanges, lbs, ubs, steps);
 
   // Affine loops require constant steps.
   SmallVector<int64_t, 4> constantSteps;
   constantSteps.reserve(steps.size());
   for (Value v : steps) {
void GenerateLoopNest<TiledLoopOp>::doit(
     ArrayRef<StringRef> distributionTypes) {
   SmallVector<ProcInfo, 2> procInfo;
   SmallVector<Value, 4> lbs, ubs, steps;
   unpackRanges(loopRanges, lbs, ubs, steps);
 
   auto wrappedBuilderFn = [&](OpBuilder &nestedBuilder, Location nestedLoc,
                               ValueRange ivs, ValueRange inputs,
                               ValueRange outputs) {
-    scf::ValueVector results = bodyBuilderFn(nestedBuilder, nestedLoc, ivs,
-                                             linalgOp.getOutputTensors());
+    SmallVector<Value> outputTensors = linalgOp.getOutputTensorOperands();
+    scf::ValueVector results =
+        bodyBuilderFn(nestedBuilder, nestedLoc, ivs, outputTensors);
     nestedBuilder.create<linalg::YieldOp>(nestedLoc, results);
   };
 
-  auto tiledLoop = b.create<TiledLoopOp>(
-      loc, lbs, ubs, steps, linalgOp.getInputs(), linalgOp.getOutputs(),
+  SmallVector<Value> inputOperands = linalgOp.getInputOperands();
+  SmallVector<Value> outputOperands = linalgOp.getOutputOperands();
+  auto tiledLoop =
+      b.create<TiledLoopOp>(loc, lbs, ubs, steps, inputOperands, outputOperands,
                             b.getArrayAttr(iteratorTypes), wrappedBuilderFn);
   if (!distributionTypes.empty())
     tiledLoop.setDistributionTypes(b, distributionTypes);
 
   // Replace inputs/outputs with the corresponding region args.
   auto isInsideTiledLoop = [&](OpOperand &operand) {
     return operand.getOwner()->getBlock() == tiledLoop.getBody();
   };
-  for (auto it :
-       llvm::zip(linalgOp.getInputs(), tiledLoop.getRegionInputArgs()))
+  for (auto it : llvm::zip(inputOperands, tiledLoop.getRegionInputArgs()))
     std::get<0>(it).replaceUsesWithIf(std::get<1>(it), isInsideTiledLoop);
-  for (auto it :
-       llvm::zip(linalgOp.getOutputs(), tiledLoop.getRegionOutputArgs()))
+  for (auto it : llvm::zip(outputOperands, tiledLoop.getRegionOutputArgs()))
     std::get<0>(it).replaceUsesWithIf(std::get<1>(it), isInsideTiledLoop);
 }
 
 /// Update the `lb`, `ub` and `step` to get per processor `lb`, `ub` and `step`.
 void updateBoundsForCyclicDistribution(OpBuilder &b, Location loc, Value procId,
                                        Value nprocs, Value &lb, Value &ub,
                                        Value &step) {
   AffineExpr d0, d1;

 void GenerateLoopNest<scf::ParallelOp>::doit(
     OpBuilder &b, Location loc, ArrayRef<Range> loopRanges, LinalgOp linalgOp,
     ArrayRef<Attribute> iteratorTypes,
     function_ref<scf::ValueVector(OpBuilder &, Location, ValueRange,
                                   ValueRange)>
         bodyBuilderFn,
     Optional<LinalgLoopDistributionOptions> distributionOptions,
     ArrayRef<StringRef> distributionTypes) {
-  auto iterArgInitValues = linalgOp.getOutputTensors();
+  SmallVector<Value> iterArgInitValues = linalgOp.getOutputTensorOperands();
   assert(iterArgInitValues.empty() && "unexpected ParallelOp init values");
   // This function may be passed more iterator types than ranges.
   assert(iteratorTypes.size() >= loopRanges.size() &&
          "expected iterator type for all ranges");
   iteratorTypes = iteratorTypes.take_front(loopRanges.size());
   SmallVector<Value, 8> lbsStorage, ubsStorage, stepsStorage, ivs;
   unsigned numLoops = iteratorTypes.size();
   ivs.reserve(numLoops);
generateParallelLoopNest(
       },
       ivs, distributionMethod);
 
   assert(ivs.size() == iteratorTypes.size() && "did not generate enough loops");
 }
 
 SmallVector<Value, 4> makeTiledShapes(OpBuilder &b, Location loc,
                                       LinalgOp linalgOp,
-                                      ArrayRef<Value> tiledOperands,
+                                      ArrayRef<Value> valuesToTile,
                                       ValueRange ivs, ValueRange tileSizes,
                                       ArrayRef<Value> sizeBounds) {
   assert(ivs.size() == static_cast<size_t>(llvm::count_if(
                            llvm::make_range(tileSizes.begin(), tileSizes.end()),
                            [](Value v) { return !isZero(v); })) &&
          "expected as many ivs as non-zero sizes");
 
   // Construct (potentially temporary) mins and maxes on which to apply maps
   // that define tile subshapes.
   SmallVector<Value, 8> lbs, subShapeSizes;
   for (unsigned idx = 0, idxIvs = 0, e = tileSizes.size(); idx < e; ++idx) {
     LLVM_DEBUG(llvm::dbgs() << "makeTiledShapes: for loop#" << idx << "\n");
     bool isTiled = !isZero(tileSizes[idx]);
     lbs.push_back(isTiled ? ivs[idxIvs++]
                           : (Value)b.create<ConstantIndexOp>(loc, 0));
     // Before composing, we need to make range a closed interval.
     Value size = isTiled ? tileSizes[idx] : sizeBounds[idx];
     AffineExpr d0 = getAffineDimExpr(0, b.getContext());
     subShapeSizes.push_back(makeComposedAffineApply(b, loc, d0 - 1, size));
     LLVM_DEBUG(llvm::dbgs() << "lb: " << lbs.back() << "\n");
     LLVM_DEBUG(llvm::dbgs() << "size: " << subShapeSizes.back() << "\n");
   }
 
+  assert(valuesToTile.size() == linalgOp.getNumInputsAndOutputs() &&
+         "expected one value to tile for every operand");
   MLIRContext *context = b.getContext();
   SmallVector<Value, 4> tiledShapes;
-  tiledShapes.reserve(tiledOperands.size());
-  for (auto en : llvm::enumerate(tiledOperands)) {
-    Value shapedOp = en.value();
+  tiledShapes.reserve(valuesToTile.size());
+  for (OpOperand *opOperand : linalgOp.getInputAndOutputOperands()) {
+    Value shapedOp = valuesToTile[opOperand->getOperandNumber()];
     LLVM_DEBUG(llvm::dbgs() << "makeTiledShapes: for operand " << shapedOp);
-    ShapedType shapedType = shapedOp.getType().cast<ShapedType>();
-    unsigned rank = shapedType.getRank();
-    AffineMap map = linalgOp.getIndexingMap(en.index());
+    int64_t rank = linalgOp.getRank(opOperand);
+    ArrayRef<int64_t> shape = linalgOp.getShape(opOperand);
+    AffineMap map = linalgOp.getTiedIndexingMap(opOperand);
     // If the shape is not tiled, we can use it as is.
     if (!isTiled(map, tileSizes)) {
       tiledShapes.push_back(shapedOp);
-      LLVM_DEBUG(llvm::dbgs()
-                 << ": not tiled: use shape: " << shapedType << "\n");
+      LLVM_DEBUG(llvm::dbgs() << ": not tiled: use shape: "
+                              << opOperand->get().getType() << "\n");
       continue;
     }
     LLVM_DEBUG(llvm::dbgs() << ": tiled: figure out subshape...\n");
 
     // Construct a new subview / subtensor for the tile.
     SmallVector<OpFoldResult, 4> offsets, sizes, strides;
     offsets.reserve(rank);
     sizes.reserve(rank);
for (unsigned r = 0; r < rank; ++r) {
       // Resulting size needs to be made half open interval again.
       AffineExpr s0 = getAffineSymbolExpr(0, b.getContext());
       Value size = makeComposedAffineApply(b, loc, s0 + 1, closedIntSize);
       LLVM_DEBUG(llvm::dbgs() << "makeTiledShapes: raw size: " << size << "\n");
 
       // The size of the subview / subtensor should be trimmed to avoid
       // out-of-bounds accesses, unless we statically know the subshape size
       // divides the shape size evenly.
-      int64_t shapeSize = shapedType.getDimSize(r);
+      int64_t shapeSize = shape[r];
       auto sizeCst = size.getDefiningOp<ConstantIndexOp>();
       if (ShapedType::isDynamic(shapeSize) || !sizeCst ||
           (shapeSize % sizeCst.getValue()) != 0) {
         LLVM_DEBUG(llvm::dbgs() << "makeTiledShapes: shapeSize=" << shapeSize
                                 << ", size: " << size
                                 << ": make sure in bound with affine.min\n");
         AffineExpr dim0, dim1, dim2;
         bindDims(context, dim0, dim1, dim2);
for (unsigned r = 0; r < rank; ++r) {
 
       sizes.push_back(size);
       LLVM_DEBUG(llvm::dbgs()
                  << "makeTiledShapes: new offset: " << offset << "\n");
       LLVM_DEBUG(llvm::dbgs() << "makeTiledShapes: new size: " << size << "\n");
       strides.push_back(b.getIndexAttr(1));
     }
 
-    if (shapedType.isa<MemRefType>())
+    if (opOperand->get().getType().isa<MemRefType>())
       tiledShapes.push_back(
           b.create<memref::SubViewOp>(loc, shapedOp, offsets, sizes, strides));
     else
       tiledShapes.push_back(
           b.create<SubTensorOp>(loc, shapedOp, offsets, sizes, strides));
   }
 
   return tiledShapes;
 }
 
 } // namespace linalg
 } // namespace mlir