[mlir] NFC: move mergeOffsetsSizesAndStrides into Affine/Utils

So that these utility functions can also be used ViewLikeInterface ops not in the memref dialect. Reviewed By: mravishankar, christopherbate Differential Revision: https://reviews.llvm.org/D134487
2026-02-02 02:00:03 +08:00 · 2022-09-23 13:20:10 -04:00
parent 2c692d891e
commit 465ec4e0b4
7 changed files with 133 additions and 97 deletions
--- a/mlir/include/mlir/Dialect/Affine/ViewLikeInterfaceUtils.h
+++ 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<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
--- 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<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`
 //===----------------------------------------------------------------------===//
--- a/mlir/lib/Dialect/Affine/Utils/CMakeLists.txt
+++ b/mlir/lib/Dialect/Affine/Utils/CMakeLists.txt
@@ -2,6 +2,7 @@ add_mlir_dialect_library(MLIRAffineUtils
  LoopFusionUtils.cpp
  LoopUtils.cpp
  Utils.cpp
+  ViewLikeInterfaceUtils.cpp

  ADDITIONAL_HEADER_DIRS
  ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/Affine
@@ -13,4 +14,5 @@ add_mlir_dialect_library(MLIRAffineUtils
  MLIRArithmeticUtils
  MLIRMemRefDialect
  MLIRTransformUtils
+  MLIRViewLikeInterface
  )
--- a/mlir/lib/Dialect/Affine/Utils/ViewLikeInterfaceUtils.cpp
+++ 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<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);
+}
--- a/mlir/lib/Dialect/Tensor/Transforms/CMakeLists.txt
+++ b/mlir/lib/Dialect/Tensor/Transforms/CMakeLists.txt
@@ -14,6 +14,7 @@ add_mlir_dialect_library(MLIRTensorTransforms

  LINK_LIBS PUBLIC
  MLIRAffineDialect
+  MLIRAffineUtils
  MLIRArithmeticDialect
  MLIRBufferizationDialect
  MLIRBufferizationTransforms
--- 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<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> {
--- a/utils/bazel/llvm-project-overlay/mlir/BUILD.bazel
+++ b/utils/bazel/llvm-project-overlay/mlir/BUILD.bazel
@@ -2440,6 +2440,7 @@ cc_library(
        "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 @@ cc_library(
        ":SCFDialect",
        ":Support",
        ":TransformUtils",
+        ":ViewLikeInterface",
        "//llvm:Support",
    ],
 )
@@ -5067,6 +5069,7 @@ cc_library(
    includes = ["include"],
    deps = [
        ":AffineDialect",
+        ":AffineUtils",
        ":ArithmeticDialect",
        ":ArithmeticUtils",
        ":BufferizationDialect",