diff --git a/mlir/include/mlir/Dialect/StandardOps/Utils/Utils.h b/mlir/include/mlir/Dialect/StandardOps/Utils/Utils.h new file mode 100644 --- /dev/null +++ b/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_H +#define MLIR_DIALECT_STANDARDOPS_UTILS_H + +#include "mlir/Dialect/StandardOps/IR/Ops.h" + +namespace mlir { + +class Location; +class OpBuilder; + +/// Given an operation, retrieves the value of each dynamic dimension through +/// constructing the necessary DimOp operators. +SmallVector getDynOperands(Location loc, Value val, OpBuilder &b); + +} // end namespace mlir + +#endif // MLIR_DIALECT_STANDARDOPS_UTILS_H diff --git a/mlir/lib/Dialect/Linalg/Transforms/Bufferize.cpp b/mlir/lib/Dialect/Linalg/Transforms/Bufferize.cpp --- a/mlir/lib/Dialect/Linalg/Transforms/Bufferize.cpp +++ b/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 getDynOperands(Location loc, Value val, - OpBuilder &b) { - SmallVector dynOperands; - auto shapedType = val.getType().cast(); - for (auto dim : llvm::enumerate(shapedType.getShape())) { - if (dim.value() == TensorType::kDynamicSize) { - dynOperands.push_back(b.create(loc, val, dim.index())); - } - } - return dynOperands; -} - static Value cloneMemref(Location loc, Value memref, OpBuilder &b) { auto memrefType = memref.getType().cast(); auto alloc = diff --git a/mlir/lib/Dialect/Linalg/Transforms/ElementwiseToLinalg.cpp b/mlir/lib/Dialect/Linalg/Transforms/ElementwiseToLinalg.cpp --- a/mlir/lib/Dialect/Linalg/Transforms/ElementwiseToLinalg.cpp +++ b/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 @@ // Extract static / dynamic shape mix from the first operand. Value firstOperand = operands.front(); auto rankedTensorType = t.cast(); - SmallVector dynamicShape; - SmallVector 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(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( loc, dynamicShape, staticShape, rankedTensorType.getElementType())); } diff --git a/mlir/lib/Dialect/StandardOps/CMakeLists.txt b/mlir/lib/Dialect/StandardOps/CMakeLists.txt --- a/mlir/lib/Dialect/StandardOps/CMakeLists.txt +++ b/mlir/lib/Dialect/StandardOps/CMakeLists.txt @@ -2,6 +2,7 @@ IR/Ops.cpp EDSC/Builders.cpp EDSC/Intrinsics.cpp + Utils/Utils.cpp ADDITIONAL_HEADER_DIRS ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/StandardOps diff --git a/mlir/lib/Dialect/StandardOps/Utils/Utils.cpp b/mlir/lib/Dialect/StandardOps/Utils/Utils.cpp new file mode 100644 --- /dev/null +++ b/mlir/lib/Dialect/StandardOps/Utils/Utils.cpp @@ -0,0 +1,29 @@ +//===- 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 mlir::getDynOperands(Location loc, Value val, + OpBuilder &b) { + SmallVector dynOperands; + auto shapedType = val.getType().cast(); + for (auto dim : llvm::enumerate(shapedType.getShape())) { + if (dim.value() == TensorType::kDynamicSize) { + dynOperands.push_back(b.create(loc, val, dim.index())); + } + } + return dynOperands; +} diff --git a/mlir/lib/Transforms/BufferDeallocation.cpp b/mlir/lib/Transforms/BufferDeallocation.cpp --- a/mlir/lib/Transforms/BufferDeallocation.cpp +++ b/mlir/lib/Transforms/BufferDeallocation.cpp @@ -53,7 +53,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/IR/Operation.h" #include "mlir/Interfaces/ControlFlowInterfaces.h" #include "mlir/Interfaces/LoopLikeInterface.h" @@ -394,13 +396,8 @@ // Extract information about dynamically shaped types by // extracting their dynamic dimensions. - SmallVector dynamicOperands; - for (auto shapeElement : llvm::enumerate(memRefType.getShape())) { - if (!ShapedType::isDynamic(shapeElement.value())) - continue; - dynamicOperands.push_back(builder.create( - 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. diff --git a/mlir/lib/Transforms/PipelineDataTransfer.cpp b/mlir/lib/Transforms/PipelineDataTransfer.cpp --- a/mlir/lib/Transforms/PipelineDataTransfer.cpp +++ b/mlir/lib/Transforms/PipelineDataTransfer.cpp @@ -17,6 +17,8 @@ #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/Dialect/StandardOps/Utils/Utils.h" #include "mlir/IR/Builders.h" #include "mlir/Transforms/LoopUtils.h" #include "mlir/Transforms/Utils.h" @@ -83,13 +85,8 @@ // The double buffer is allocated right before 'forOp'. OpBuilder bOuter(forOp); // Put together alloc operands for any dynamic dimensions of the memref. - SmallVector allocOperands; - unsigned dynamicDimCount = 0; - for (auto dimSize : oldMemRefType.getShape()) { - if (dimSize == -1) - allocOperands.push_back( - bOuter.create(forOp.getLoc(), oldMemRef, dynamicDimCount++)); - } + + auto allocOperands = getDynOperands(forOp.getLoc(), oldMemRef, bOuter); // Create and place the alloc right before the 'affine.for' operation. Value newMemRef =