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

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

Differential Revision: https://reviews.llvm.org/D94340

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

@@ -0,0 +1,32 @@
+//===- Utils.h - General transformation utilities ---------------*- C++ -*-===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+//
+// This header file defines prototypes for various transformation utilities for
+// the StandardOps dialect. These are not passes by themselves but are used
+// either by passes, optimization sequences, or in turn by other transformation
+// utilities.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_DIALECT_STANDARDOPS_UTILS_UTILS_H
+#define MLIR_DIALECT_STANDARDOPS_UTILS_UTILS_H
+
+#include "mlir/IR/Value.h"
+
+namespace mlir {
+
+class Location;
+class OpBuilder;
+
+/// 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_DIALECT_STANDARDOPS_UTILS_UTILS_H

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

@@ -13,6 +13,7 @@
 #include "mlir/Dialect/Linalg/Transforms/Transforms.h"
 #include "mlir/Dialect/Linalg/Utils/Utils.h"
 #include "mlir/Dialect/StandardOps/Transforms/Passes.h"
+#include "mlir/Dialect/StandardOps/Utils/Utils.h"
 #include "mlir/Dialect/Vector/VectorOps.h"
 #include "mlir/IR/BuiltinDialect.h"
 #include "mlir/IR/Operation.h"
@@ -21,18 +22,6 @@
 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 =

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

@@ -10,7 +10,9 @@
 
 #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/Dialect/StandardOps/Utils/Utils.h"
 #include "mlir/Transforms/DialectConversion.h"
 
 using namespace mlir;
@@ -62,18 +64,9 @@ getOrCreateOperandsMatchingResultTypes(OpBuilder &b, Operation *op) {
     // 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 = getDynOperands(loc, firstOperand, b);
+
     res.push_back(b.create<linalg::InitTensorOp>(
         loc, dynamicShape, staticShape, rankedTensorType.getElementType()));
   }

mlir/lib/Dialect/StandardOps/CMakeLists.txt

@@ -2,6 +2,7 @@ add_mlir_dialect_library(MLIRStandard
   IR/Ops.cpp
   EDSC/Builders.cpp
   EDSC/Intrinsics.cpp
+  Utils/Utils.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/StandardOps

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

@@ -0,0 +1,28 @@
+//===- Utils.cpp - Utilities to support the Linalg dialect ----------------===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+//
+// This file implements utilities for the Linalg dialect.
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Dialect/StandardOps/Utils/Utils.h"
+
+#include "mlir/Dialect/StandardOps/IR/Ops.h"
+
+using namespace mlir;
+
+SmallVector<Value, 4> mlir::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;
+}

mlir/lib/Transforms/BufferDeallocation.cpp

@@ -54,6 +54,7 @@
 #include "PassDetail.h"
 #include "mlir/Dialect/Linalg/IR/LinalgOps.h"
 #include "mlir/Dialect/StandardOps/IR/Ops.h"
+#include "mlir/Dialect/StandardOps/Utils/Utils.h"
 #include "mlir/IR/Operation.h"
 #include "mlir/Interfaces/ControlFlowInterfaces.h"
 #include "mlir/Interfaces/LoopLikeInterface.h"
@@ -394,13 +395,8 @@ private:
 
     // 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 =
+        getDynOperands(terminator->getLoc(), sourceValue, builder);
 
     // TODO: provide a generic interface to create dialect-specific
     // Alloc and CopyOp nodes.

mlir/lib/Transforms/PipelineDataTransfer.cpp

@@ -17,6 +17,7 @@
 #include "mlir/Analysis/LoopAnalysis.h"
 #include "mlir/Analysis/Utils.h"
 #include "mlir/Dialect/Affine/IR/AffineOps.h"
+#include "mlir/Dialect/StandardOps/Utils/Utils.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/Transforms/LoopUtils.h"
 #include "mlir/Transforms/Utils.h"
@@ -83,13 +84,8 @@ static bool doubleBuffer(Value oldMemRef, AffineForOp forOp) {
   // 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 = getDynOperands(forOp.getLoc(), oldMemRef, bOuter);
 
   // Create and place the alloc right before the 'affine.for' operation.
   Value newMemRef =