[MLIR][Linalg] Refactor transforms to use linalg::getDynOperands helper

[rsuderman]

Summary

getDynOperands behavior is commonly used in a number of passes. Refactored to
use a helper function and avoid code reuse.

[rriddle]

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

372

This looks like something that would go wherever DimOp is defined. Adding here creates an otherwise unnecessary dependency on Linalg in several places.

[mehdi_amini]

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

372

Either we need to add a Utils.h in StandardDialect or we could make it a static method on the DimOp class itself?

[rsuderman]

Moved from Linalg/Utils to StandardOps/Utils

[rriddle]

mlir/include/mlir/Dialect/StandardOps/Utils/Utils.h

19 ↗

(On Diff #315554)

This doesn't look necessary. Can you trim the scope of this include?

mlir/lib/Dialect/StandardOps/Utils/Utils.cpp

24 ↗

(On Diff #315554)

nit: Remove the trivial braces here.

mlir/lib/Transforms/BufferDeallocation.cpp

56

Is this necessary?

mlir/lib/Transforms/PipelineDataTransfer.cpp

20

Is this necessary?

[rsuderman]

River's comments for pairing down deps.

[mehdi_amini]

LGTM, Thanks!

[mehdi_amini]

LGTM, Thanks!

mlir/include/mlir/Dialect/StandardOps/Utils/Utils.h

17 ↗

(On Diff #315557)

The clang-tidy warning is because you're missing a UTILS_ for the directory

[rsuderman]

Utils header guard change.

mlir/include/mlir/Dialect/Linalg/Utils/Utils.h

struct GenerateLoopNest {
 
   static void
   doit(ArrayRef<Range> loopRanges, ValueRange iterArgInitValues,
        ArrayRef<Attribute> iteratorTypes,
        function_ref<scf::ValueVector(ValueRange, ValueRange)> bodyBuilderFn,
        Optional<LinalgLoopDistributionOptions> = None);
 };
 
+/// Given an operation, retrieves the value of each dynamic dimension through
+/// constructing the necessary DimOp operators.
+SmallVector<Value, 4> getDynOperands(Location loc, Value val, OpBuilder &b);
+
 } // namespace linalg
 } // namespace mlir
 
 #endif // MLIR_DIALECT_LINALG_UTILS_H_

mlir/include/mlir/Transforms/Utils.h

 /// In the following cases, `sliceOps` remains empty:
 ///   1. If none of opInst's operands were the result of an affine.apply
 ///      (i.e., there was no affine computation slice to create).
 ///   2. If all the affine.apply op's supplying operands to this opInst did not
 ///      have any uses other than those in this opInst.
 void createAffineComputationSlice(Operation *opInst,
                                   SmallVectorImpl<AffineApplyOp> *sliceOps);
 
+/// Given an operation, retrieves the value of each dynamic dimension through
+/// constructing the necessary DimOp operators.
+SmallVector<Value, 4> getDynOperands(Location loc, Value val, OpBuilder &b);
+
 } // end namespace mlir
 
 #endif // MLIR_TRANSFORMS_UTILS_H

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

 #include "mlir/Dialect/Vector/VectorOps.h"
 #include "mlir/IR/BuiltinDialect.h"
 #include "mlir/IR/Operation.h"
 #include "mlir/Pass/Pass.h"
 
 using namespace ::mlir;
 using namespace ::mlir::linalg;
 
-static SmallVector<Value, 4> getDynOperands(Location loc, Value val,
-                                            OpBuilder &b) {
-  SmallVector<Value, 4> dynOperands;
-  auto shapedType = val.getType().cast<ShapedType>();
-  for (auto dim : llvm::enumerate(shapedType.getShape())) {
-    if (dim.value() == TensorType::kDynamicSize) {
-      dynOperands.push_back(b.create<DimOp>(loc, val, dim.index()));
-    }
-  }
-  return dynOperands;
-}
-
 static Value cloneMemref(Location loc, Value memref, OpBuilder &b) {
   auto memrefType = memref.getType().cast<MemRefType>();
   auto alloc =
       b.create<AllocOp>(loc, memrefType, getDynOperands(loc, memref, b));
   b.create<linalg::CopyOp>(loc, memref, alloc);
   return alloc;
 }
 

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

 //===- ElementwiseToLinalg.cpp - conversion of elementwise to linalg ------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/Linalg/Passes.h"
 
 #include "PassDetail.h"
 #include "mlir/Dialect/Linalg/IR/LinalgOps.h"
+#include "mlir/Dialect/Linalg/Utils/Utils.h"
 #include "mlir/Dialect/StandardOps/IR/Ops.h"
 #include "mlir/Transforms/DialectConversion.h"
 
 using namespace mlir;
 
 static bool isElementwiseMappableOpOnRankedTensors(Operation *op) {
   if (!op->hasTrait<OpTrait::ElementwiseMappable>())
     return false;
for (Value v : operands) {
       }
     }
     if (found)
       continue;
 
     // Extract static / dynamic shape mix from the first operand.
     Value firstOperand = operands.front();
     auto rankedTensorType = t.cast<RankedTensorType>();
-    SmallVector<Value, 8> dynamicShape;
-    SmallVector<int64_t, 8> staticShape;
-    dynamicShape.reserve(rankedTensorType.getRank());
-    staticShape.reserve(rankedTensorType.getRank());
-    unsigned idx = 0;
-    for (auto shape : rankedTensorType.getShape()) {
-      staticShape.push_back(shape);
-      if (rankedTensorType.isDynamicDim(idx))
-        dynamicShape.push_back(b.create<DimOp>(loc, firstOperand, idx));
-      ++idx;
-    }
-    // Create init tensor.
+    auto staticShape = llvm::to_vector<4>(rankedTensorType.getShape());
+    auto dynamicShape = linalg::getDynOperands(loc, firstOperand, b);
+
     res.push_back(b.create<linalg::InitTensorOp>(
         loc, dynamicShape, staticShape, rankedTensorType.getElementType()));
   }
   return res;
 }
 
 namespace {
 struct ConvertAnyElementwiseMappableOpOnRankedTensors : public RewritePattern {

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

void GenerateLoopNest<scf::ParallelOp>::doit(
   };
   generateParallelLoopNest(lbs, ubs, steps, iteratorTypes,
                            bodyBuilderWithoutIterArgsFn, ivs,
                            distributionMethod);
 
   assert(ivs.size() == iteratorTypes.size() && "did not generate enough loops");
 }
 
+SmallVector<Value, 4> getDynOperands(Location loc, Value val, OpBuilder &b) {
+  SmallVector<Value, 4> dynOperands;
+  auto shapedType = val.getType().cast<ShapedType>();
+  for (auto dim : llvm::enumerate(shapedType.getShape())) {
+    if (dim.value() == TensorType::kDynamicSize) {
+      dynOperands.push_back(b.create<DimOp>(loc, val, dim.index()));
+    }
+  }
+  return dynOperands;
+}
+
 } // namespace linalg
 } // namespace mlir

mlir/lib/Transforms/BufferDeallocation.cpp

 // the resulting code will be invalid with respect to program semantics.
 // However, structured control-flow loops are fully supported. Furthermore, it
 // doesn't accept functions which return buffers already.
 //
 //===----------------------------------------------------------------------===//
 
 #include "PassDetail.h"
 #include "mlir/Dialect/Linalg/IR/LinalgOps.h"
+#include "mlir/Dialect/Linalg/Utils/Utils.h"
 #include "mlir/Dialect/StandardOps/IR/Ops.h"
 #include "mlir/IR/Operation.h"
 #include "mlir/Interfaces/ControlFlowInterfaces.h"
 #include "mlir/Interfaces/LoopLikeInterface.h"
 #include "mlir/Pass/Pass.h"
 #include "mlir/Transforms/BufferUtils.h"
 #include "mlir/Transforms/Passes.h"
 #include "llvm/ADT/SetOperations.h"
Value introduceBufferCopy(Value sourceValue, Operation *terminator) {
     if (copiedValues.contains(sourceValue))
       return sourceValue;
     // Create a new alloc at the current location of the terminator.
     auto memRefType = sourceValue.getType().cast<MemRefType>();
     OpBuilder builder(terminator);
 
     // Extract information about dynamically shaped types by
     // extracting their dynamic dimensions.
-    SmallVector<Value, 4> dynamicOperands;
-    for (auto shapeElement : llvm::enumerate(memRefType.getShape())) {
-      if (!ShapedType::isDynamic(shapeElement.value()))
-        continue;
-      dynamicOperands.push_back(builder.create<DimOp>(
-          terminator->getLoc(), sourceValue, shapeElement.index()));
-    }
+    auto dynamicOperands =
+        linalg::getDynOperands(terminator->getLoc(), sourceValue, builder);
 
     // TODO: provide a generic interface to create dialect-specific
     // Alloc and CopyOp nodes.
     auto alloc = builder.create<AllocOp>(terminator->getLoc(), memRefType,
                                          dynamicOperands);
 
     // Create a new copy operation that copies to contents of the old
     // allocation to the new one.

mlir/lib/Transforms/PipelineDataTransfer.cpp

 
 #include "PassDetail.h"
 #include "mlir/Transforms/Passes.h"
 
 #include "mlir/Analysis/AffineAnalysis.h"
 #include "mlir/Analysis/LoopAnalysis.h"
 #include "mlir/Analysis/Utils.h"
 #include "mlir/Dialect/Affine/IR/AffineOps.h"
+#include "mlir/Dialect/Linalg/Utils/Utils.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/Transforms/LoopUtils.h"
 #include "mlir/Transforms/Utils.h"
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/Support/Debug.h"
 
 #define DEBUG_TYPE "affine-pipeline-data-transfer"
 
static bool doubleBuffer(Value oldMemRef, AffineForOp forOp) {
   };
 
   auto oldMemRefType = oldMemRef.getType().cast<MemRefType>();
   auto newMemRefType = doubleShape(oldMemRefType);
 
   // The double buffer is allocated right before 'forOp'.
   OpBuilder bOuter(forOp);
   // Put together alloc operands for any dynamic dimensions of the memref.
-  SmallVector<Value, 4> allocOperands;
-  unsigned dynamicDimCount = 0;
-  for (auto dimSize : oldMemRefType.getShape()) {
-    if (dimSize == -1)
-      allocOperands.push_back(
-          bOuter.create<DimOp>(forOp.getLoc(), oldMemRef, dynamicDimCount++));
-  }
+
+  auto allocOperands =
+      linalg::getDynOperands(forOp.getLoc(), oldMemRef, bOuter);
 
   // Create and place the alloc right before the 'affine.for' operation.
   Value newMemRef =
       bOuter.create<AllocOp>(forOp.getLoc(), newMemRefType, allocOperands);
 
   // Create 'iv mod 2' value to index the leading dimension.
   auto d0 = bInner.getAffineDimExpr(0);
   int64_t step = forOp.getStep();