--- /dev/null
+//===- ViewLikeInterfaceUtils.h ---------------------------------*- 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
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_DIALECT_AFFINE_VIEWLIKEINTERFACEUTILS_H
+#define MLIR_DIALECT_AFFINE_VIEWLIKEINTERFACEUTILS_H
+
+#include "mlir/IR/OpDefinition.h"
+#include "mlir/Interfaces/ViewLikeInterface.h"
+
+namespace mlir {
+
+/// Fills the `combinedOffsets`, `combinedSizes` and `combinedStrides` to use
+/// when combining a producer slice **into** a consumer slice.
+///
+/// This function performs the following computation:
+/// - Combined offsets = producer_offsets * consumer_strides + consumer_offsets
+/// - Combined sizes = consumer_sizes
+/// - Combined strides = producer_strides * consumer_strides
+LogicalResult
+mergeOffsetsSizesAndStrides(OpBuilder &builder, Location loc,
+ ArrayRef<OpFoldResult> producerOffsets,
+ ArrayRef<OpFoldResult> producerSizes,
+ ArrayRef<OpFoldResult> producerStrides,
+ const llvm::SmallBitVector &droppedProducerDims,
+ ArrayRef<OpFoldResult> consumerOffsets,
+ ArrayRef<OpFoldResult> consumerSizes,
+ ArrayRef<OpFoldResult> consumerStrides,
+ SmallVector<OpFoldResult> &combinedOffsets,
+ SmallVector<OpFoldResult> &combinedSizes,
+ SmallVector<OpFoldResult> &combinedStrides);
+
+/// Fills the `combinedOffsets`, `combinedSizes` and `combinedStrides` to use
+/// when combining a `producer` slice op **into** a `consumer` slice op.
+LogicalResult
+mergeOffsetsSizesAndStrides(OpBuilder &builder, Location loc,
+ OffsetSizeAndStrideOpInterface producer,
+ OffsetSizeAndStrideOpInterface consumer,
+ const llvm::SmallBitVector &droppedProducerDims,
+ SmallVector<OpFoldResult> &combinedOffsets,
+ SmallVector<OpFoldResult> &combinedSizes,
+ SmallVector<OpFoldResult> &combinedStrides);
+
+} // namespace mlir
+
+#endif // MLIR_DIALECT_AFFINE_VIEWLIKEINTERFACEUTILS_H
namespace mlir {
namespace tensor {
-/// Fills the `combinedOffsets`, `combinedSizes` and `combinedStrides` to use
-/// when combining a producer slice **into** a consumer slice.
-///
-/// This function performs the following computation:
-/// - Combined offsets = producer_offsets * consumer_strides + consumer_offsets
-/// - Combined sizes = consumer_sizes
-/// - Combined strides = producer_strides * consumer_strides
-LogicalResult
-mergeOffsetsSizesAndStrides(OpBuilder &builder, Location loc,
- ArrayRef<OpFoldResult> producerOffsets,
- ArrayRef<OpFoldResult> producerSizes,
- ArrayRef<OpFoldResult> producerStrides,
- const llvm::SmallBitVector &droppedProducerDims,
- ArrayRef<OpFoldResult> consumerOffsets,
- ArrayRef<OpFoldResult> consumerSizes,
- ArrayRef<OpFoldResult> consumerStrides,
- SmallVector<OpFoldResult> &combinedOffsets,
- SmallVector<OpFoldResult> &combinedSizes,
- SmallVector<OpFoldResult> &combinedStrides);
-
-/// Fills the `combinedOffsets`, `combinedSizes` and `combinedStrides` to use
-/// when combining a `producer` slice op **into** a `consumer` slice op.
-LogicalResult
-mergeOffsetsSizesAndStrides(OpBuilder &builder, Location loc,
- OffsetSizeAndStrideOpInterface producer,
- OffsetSizeAndStrideOpInterface consumer,
- const llvm::SmallBitVector &droppedProducerDims,
- SmallVector<OpFoldResult> &combinedOffsets,
- SmallVector<OpFoldResult> &combinedSizes,
- SmallVector<OpFoldResult> &combinedStrides);
-
//===----------------------------------------------------------------------===//
// Extract slice from `tensor.collapse_shape`
//===----------------------------------------------------------------------===//
LoopFusionUtils.cpp
LoopUtils.cpp
Utils.cpp
+ ViewLikeInterfaceUtils.cpp
ADDITIONAL_HEADER_DIRS
${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/Affine
MLIRArithmeticUtils
MLIRMemRefDialect
MLIRTransformUtils
+ MLIRViewLikeInterface
)
--- /dev/null
+//===- ViewLikeInterfaceUtils.cpp -----------------------------------------===//
+//
+// 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/Affine/ViewLikeInterfaceUtils.h"
+#include "mlir/Dialect/Affine/IR/AffineOps.h"
+
+using namespace mlir;
+
+LogicalResult mlir::mergeOffsetsSizesAndStrides(
+ OpBuilder &builder, Location loc, ArrayRef<OpFoldResult> producerOffsets,
+ ArrayRef<OpFoldResult> producerSizes,
+ ArrayRef<OpFoldResult> producerStrides,
+ const llvm::SmallBitVector &droppedProducerDims,
+ ArrayRef<OpFoldResult> consumerOffsets,
+ ArrayRef<OpFoldResult> consumerSizes,
+ ArrayRef<OpFoldResult> consumerStrides,
+ SmallVector<OpFoldResult> &combinedOffsets,
+ SmallVector<OpFoldResult> &combinedSizes,
+ SmallVector<OpFoldResult> &combinedStrides) {
+ combinedOffsets.resize(producerOffsets.size());
+ combinedSizes.resize(producerOffsets.size());
+ combinedStrides.resize(producerOffsets.size());
+
+ AffineExpr s0, s1, s2;
+ bindSymbols(builder.getContext(), s0, s1, s2);
+
+ unsigned consumerPos = 0;
+ for (auto i : llvm::seq<unsigned>(0, producerOffsets.size())) {
+ if (droppedProducerDims.test(i)) {
+ // For dropped dims, get the values from the producer.
+ combinedOffsets[i] = producerOffsets[i];
+ combinedSizes[i] = producerSizes[i];
+ combinedStrides[i] = producerStrides[i];
+ continue;
+ }
+ SmallVector<OpFoldResult> offsetSymbols, strideSymbols;
+ // The combined offset is computed as
+ // producer_offset + consumer_offset * producer_strides.
+ combinedOffsets[i] = makeComposedFoldedAffineApply(
+ builder, loc, s0 * s1 + s2,
+ {consumerOffsets[consumerPos], producerStrides[i], producerOffsets[i]});
+ combinedSizes[i] = consumerSizes[consumerPos];
+ // The combined stride is computed as
+ // consumer_stride * producer_stride.
+ combinedStrides[i] = makeComposedFoldedAffineApply(
+ builder, loc, s0 * s1,
+ {consumerStrides[consumerPos], producerStrides[i]});
+
+ consumerPos++;
+ }
+ return success();
+}
+
+LogicalResult mlir::mergeOffsetsSizesAndStrides(
+ OpBuilder &builder, Location loc, OffsetSizeAndStrideOpInterface producer,
+ OffsetSizeAndStrideOpInterface consumer,
+ const llvm::SmallBitVector &droppedProducerDims,
+ SmallVector<OpFoldResult> &combinedOffsets,
+ SmallVector<OpFoldResult> &combinedSizes,
+ SmallVector<OpFoldResult> &combinedStrides) {
+ SmallVector<OpFoldResult> consumerOffsets = consumer.getMixedOffsets();
+ SmallVector<OpFoldResult> consumerSizes = consumer.getMixedSizes();
+ SmallVector<OpFoldResult> consumerStrides = consumer.getMixedStrides();
+ SmallVector<OpFoldResult> producerOffsets = producer.getMixedOffsets();
+ SmallVector<OpFoldResult> producerSizes = producer.getMixedSizes();
+ SmallVector<OpFoldResult> producerStrides = producer.getMixedStrides();
+ return mergeOffsetsSizesAndStrides(
+ builder, loc, producerOffsets, producerSizes, producerStrides,
+ droppedProducerDims, consumerOffsets, consumerSizes, consumerStrides,
+ combinedOffsets, combinedSizes, combinedStrides);
+}
LINK_LIBS PUBLIC
MLIRAffineDialect
+ MLIRAffineUtils
MLIRArithmeticDialect
MLIRBufferizationDialect
MLIRBufferizationTransforms
//
//===----------------------------------------------------------------------===//
-#include "mlir/Dialect/Affine/IR/AffineOps.h"
+#include "mlir/Dialect/Affine/ViewLikeInterfaceUtils.h"
#include "mlir/Dialect/Tensor/IR/Tensor.h"
-#include "mlir/Dialect/Tensor/Transforms/TransformUtils.h"
#include "mlir/Dialect/Tensor/Transforms/Transforms.h"
#include "mlir/IR/BuiltinTypes.h"
#include "mlir/IR/OpDefinition.h"
using namespace mlir;
using namespace mlir::tensor;
-LogicalResult tensor::mergeOffsetsSizesAndStrides(
- OpBuilder &builder, Location loc, ArrayRef<OpFoldResult> producerOffsets,
- ArrayRef<OpFoldResult> producerSizes,
- ArrayRef<OpFoldResult> producerStrides,
- const llvm::SmallBitVector &droppedProducerDims,
- ArrayRef<OpFoldResult> consumerOffsets,
- ArrayRef<OpFoldResult> consumerSizes,
- ArrayRef<OpFoldResult> consumerStrides,
- SmallVector<OpFoldResult> &combinedOffsets,
- SmallVector<OpFoldResult> &combinedSizes,
- SmallVector<OpFoldResult> &combinedStrides) {
- combinedOffsets.resize(producerOffsets.size());
- combinedSizes.resize(producerOffsets.size());
- combinedStrides.resize(producerOffsets.size());
-
- AffineExpr s0, s1, s2;
- bindSymbols(builder.getContext(), s0, s1, s2);
-
- unsigned consumerPos = 0;
- for (auto i : llvm::seq<unsigned>(0, producerOffsets.size())) {
- if (droppedProducerDims.test(i)) {
- // For dropped dims, get the values from the producer.
- combinedOffsets[i] = producerOffsets[i];
- combinedSizes[i] = producerSizes[i];
- combinedStrides[i] = producerStrides[i];
- continue;
- }
- SmallVector<OpFoldResult> offsetSymbols, strideSymbols;
- // The combined offset is computed as
- // producer_offset + consumer_offset * producer_strides.
- combinedOffsets[i] = makeComposedFoldedAffineApply(
- builder, loc, s0 * s1 + s2,
- {consumerOffsets[consumerPos], producerStrides[i], producerOffsets[i]});
- combinedSizes[i] = consumerSizes[consumerPos];
- // The combined stride is computed as
- // consumer_stride * producer_stride.
- combinedStrides[i] = makeComposedFoldedAffineApply(
- builder, loc, s0 * s1,
- {consumerStrides[consumerPos], producerStrides[i]});
-
- consumerPos++;
- }
- return success();
-}
-
-LogicalResult tensor::mergeOffsetsSizesAndStrides(
- OpBuilder &builder, Location loc, OffsetSizeAndStrideOpInterface producer,
- OffsetSizeAndStrideOpInterface consumer,
- const llvm::SmallBitVector &droppedProducerDims,
- SmallVector<OpFoldResult> &combinedOffsets,
- SmallVector<OpFoldResult> &combinedSizes,
- SmallVector<OpFoldResult> &combinedStrides) {
- SmallVector<OpFoldResult> consumerOffsets = consumer.getMixedOffsets();
- SmallVector<OpFoldResult> consumerSizes = consumer.getMixedSizes();
- SmallVector<OpFoldResult> consumerStrides = consumer.getMixedStrides();
- SmallVector<OpFoldResult> producerOffsets = producer.getMixedOffsets();
- SmallVector<OpFoldResult> producerSizes = producer.getMixedSizes();
- SmallVector<OpFoldResult> producerStrides = producer.getMixedStrides();
- return tensor::mergeOffsetsSizesAndStrides(
- builder, loc, producerOffsets, producerSizes, producerStrides,
- droppedProducerDims, consumerOffsets, consumerSizes, consumerStrides,
- combinedOffsets, combinedSizes, combinedStrides);
-}
-
namespace {
/// Merges consecutive tensor.extract_slice ops into one.
struct MergeConsecutiveExtractSlice : public OpRewritePattern<ExtractSliceOp> {
"include/mlir/Dialect/Affine/LoopFusionUtils.h",
"include/mlir/Dialect/Affine/LoopUtils.h",
"include/mlir/Dialect/Affine/Utils.h",
+ "include/mlir/Dialect/Affine/ViewLikeInterfaceUtils.h",
],
includes = ["include"],
deps = [
":SCFDialect",
":Support",
":TransformUtils",
+ ":ViewLikeInterface",
"//llvm:Support",
],
)
includes = ["include"],
deps = [
":AffineDialect",
+ ":AffineUtils",
":ArithmeticDialect",
":ArithmeticUtils",
":BufferizationDialect",