diff --git a/mlir/include/mlir/Dialect/Affine/ViewLikeInterfaceUtils.h b/mlir/include/mlir/Dialect/Affine/ViewLikeInterfaceUtils.h new file mode 100644 --- /dev/null +++ b/mlir/include/mlir/Dialect/Affine/ViewLikeInterfaceUtils.h @@ -0,0 +1,50 @@ +//===- 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 producerOffsets, + ArrayRef producerSizes, + ArrayRef producerStrides, + const llvm::SmallBitVector &droppedProducerDims, + ArrayRef consumerOffsets, + ArrayRef consumerSizes, + ArrayRef consumerStrides, + SmallVector &combinedOffsets, + SmallVector &combinedSizes, + SmallVector &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 &combinedOffsets, + SmallVector &combinedSizes, + SmallVector &combinedStrides); + +} // namespace mlir + +#endif // MLIR_DIALECT_AFFINE_VIEWLIKEINTERFACEUTILS_H diff --git a/mlir/include/mlir/Dialect/Tensor/Transforms/TransformUtils.h b/mlir/include/mlir/Dialect/Tensor/Transforms/TransformUtils.h --- a/mlir/include/mlir/Dialect/Tensor/Transforms/TransformUtils.h +++ b/mlir/include/mlir/Dialect/Tensor/Transforms/TransformUtils.h @@ -14,37 +14,6 @@ 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 producerOffsets, - ArrayRef producerSizes, - ArrayRef producerStrides, - const llvm::SmallBitVector &droppedProducerDims, - ArrayRef consumerOffsets, - ArrayRef consumerSizes, - ArrayRef consumerStrides, - SmallVector &combinedOffsets, - SmallVector &combinedSizes, - SmallVector &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 &combinedOffsets, - SmallVector &combinedSizes, - SmallVector &combinedStrides); - //===----------------------------------------------------------------------===// // Extract slice from `tensor.collapse_shape` //===----------------------------------------------------------------------===// diff --git a/mlir/lib/Dialect/Affine/Utils/CMakeLists.txt b/mlir/lib/Dialect/Affine/Utils/CMakeLists.txt --- a/mlir/lib/Dialect/Affine/Utils/CMakeLists.txt +++ b/mlir/lib/Dialect/Affine/Utils/CMakeLists.txt @@ -2,6 +2,7 @@ LoopFusionUtils.cpp LoopUtils.cpp Utils.cpp + ViewLikeInterfaceUtils.cpp ADDITIONAL_HEADER_DIRS ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/Affine @@ -13,4 +14,5 @@ MLIRArithmeticUtils MLIRMemRefDialect MLIRTransformUtils + MLIRViewLikeInterface ) diff --git a/mlir/lib/Dialect/Affine/Utils/ViewLikeInterfaceUtils.cpp b/mlir/lib/Dialect/Affine/Utils/ViewLikeInterfaceUtils.cpp new file mode 100644 --- /dev/null +++ b/mlir/lib/Dialect/Affine/Utils/ViewLikeInterfaceUtils.cpp @@ -0,0 +1,76 @@ +//===- 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 producerOffsets, + ArrayRef producerSizes, + ArrayRef producerStrides, + const llvm::SmallBitVector &droppedProducerDims, + ArrayRef consumerOffsets, + ArrayRef consumerSizes, + ArrayRef consumerStrides, + SmallVector &combinedOffsets, + SmallVector &combinedSizes, + SmallVector &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(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 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 &combinedOffsets, + SmallVector &combinedSizes, + SmallVector &combinedStrides) { + SmallVector consumerOffsets = consumer.getMixedOffsets(); + SmallVector consumerSizes = consumer.getMixedSizes(); + SmallVector consumerStrides = consumer.getMixedStrides(); + SmallVector producerOffsets = producer.getMixedOffsets(); + SmallVector producerSizes = producer.getMixedSizes(); + SmallVector producerStrides = producer.getMixedStrides(); + return mergeOffsetsSizesAndStrides( + builder, loc, producerOffsets, producerSizes, producerStrides, + droppedProducerDims, consumerOffsets, consumerSizes, consumerStrides, + combinedOffsets, combinedSizes, combinedStrides); +} diff --git a/mlir/lib/Dialect/Tensor/Transforms/CMakeLists.txt b/mlir/lib/Dialect/Tensor/Transforms/CMakeLists.txt --- a/mlir/lib/Dialect/Tensor/Transforms/CMakeLists.txt +++ b/mlir/lib/Dialect/Tensor/Transforms/CMakeLists.txt @@ -14,6 +14,7 @@ LINK_LIBS PUBLIC MLIRAffineDialect + MLIRAffineUtils MLIRArithmeticDialect MLIRBufferizationDialect MLIRBufferizationTransforms diff --git a/mlir/lib/Dialect/Tensor/Transforms/MergeConsecutiveInsertExtractSlicePatterns.cpp b/mlir/lib/Dialect/Tensor/Transforms/MergeConsecutiveInsertExtractSlicePatterns.cpp --- a/mlir/lib/Dialect/Tensor/Transforms/MergeConsecutiveInsertExtractSlicePatterns.cpp +++ b/mlir/lib/Dialect/Tensor/Transforms/MergeConsecutiveInsertExtractSlicePatterns.cpp @@ -6,9 +6,8 @@ // //===----------------------------------------------------------------------===// -#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" @@ -17,70 +16,6 @@ using namespace mlir; using namespace mlir::tensor; -LogicalResult tensor::mergeOffsetsSizesAndStrides( - OpBuilder &builder, Location loc, ArrayRef producerOffsets, - ArrayRef producerSizes, - ArrayRef producerStrides, - const llvm::SmallBitVector &droppedProducerDims, - ArrayRef consumerOffsets, - ArrayRef consumerSizes, - ArrayRef consumerStrides, - SmallVector &combinedOffsets, - SmallVector &combinedSizes, - SmallVector &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(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 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 &combinedOffsets, - SmallVector &combinedSizes, - SmallVector &combinedStrides) { - SmallVector consumerOffsets = consumer.getMixedOffsets(); - SmallVector consumerSizes = consumer.getMixedSizes(); - SmallVector consumerStrides = consumer.getMixedStrides(); - SmallVector producerOffsets = producer.getMixedOffsets(); - SmallVector producerSizes = producer.getMixedSizes(); - SmallVector 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 { diff --git a/utils/bazel/llvm-project-overlay/mlir/BUILD.bazel b/utils/bazel/llvm-project-overlay/mlir/BUILD.bazel --- a/utils/bazel/llvm-project-overlay/mlir/BUILD.bazel +++ b/utils/bazel/llvm-project-overlay/mlir/BUILD.bazel @@ -2440,6 +2440,7 @@ "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 = [ @@ -2453,6 +2454,7 @@ ":SCFDialect", ":Support", ":TransformUtils", + ":ViewLikeInterface", "//llvm:Support", ], ) @@ -5067,6 +5069,7 @@ includes = ["include"], deps = [ ":AffineDialect", + ":AffineUtils", ":ArithmeticDialect", ":ArithmeticUtils", ":BufferizationDialect",