[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 ↗

(On Diff #315512)

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 ↗

(On Diff #315512)

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

20

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

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

25

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

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

[rsuderman]

Utils header guard change.

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

+//===- 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_H
+#define MLIR_DIALECT_STANDARDOPS_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_H

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

 //===- Bufferize.cpp - Bufferization of linalg ops ------------------===//
 //
 // 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/Transforms/Bufferize.h"
 #include "PassDetail.h"
 #include "mlir/Dialect/Linalg/IR/LinalgOps.h"
 #include "mlir/Dialect/Linalg/Passes.h"
 #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"
 #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/Dialect/StandardOps/Utils/Utils.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 = 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/StandardOps/CMakeLists.txt

 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
 
   DEPENDS
   MLIRStandardOpsIncGen
 
   LINK_LIBS PUBLIC

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

+//===- 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

 // 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/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"
 #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 =
+        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/StandardOps/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 = 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();