[mlir] NFC: move mergeOffsetsSizesAndStrides into Affine/Utils

[antiagainst]

Summary

So that these utility functions can also be used ViewLikeInterface
ops not in the memref dialect.

[mravishankar]

mlir/lib/Dialect/Affine/Utils/ViewLikeInterfaceUtils.cpp

59

Nit: you could just inline this into the header file...

[christopherbate]

Thanks!

[antiagainst]

mlir/lib/Dialect/Affine/Utils/ViewLikeInterfaceUtils.cpp

59

I typically inline for oneliners or similar. I'll leave this as-is given this is quite a few lines.

mlir/include/mlir/Dialect/Affine/ViewLikeInterfaceUtils.h

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

mlir/include/mlir/Dialect/Tensor/Transforms/TransformUtils.h

 //===- TransformsUtils.h - Tensor 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
 //
 //===----------------------------------------------------------------------===//
 #ifndef MLIR_DIALECT_TENSOR_TRANSFORMS_TRANSFORMUTILS_H
 #define MLIR_DIALECT_TENSOR_TRANSFORMS_TRANSFORMUTILS_H
 
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/IR/PatternMatch.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`
 //===----------------------------------------------------------------------===//
 
 /// This class assists with generating IR required to materialize an
 /// arbitrary-sized slice from the result of a CollapseShapeOp. In order to
 /// accomplish this, a loop nest or similar operation must be created by the
 /// caller. The purpose of the loop nest is to generate a "tiling by 1" of all

mlir/lib/Dialect/Affine/Utils/CMakeLists.txt

 add_mlir_dialect_library(MLIRAffineUtils
   LoopFusionUtils.cpp
   LoopUtils.cpp
   Utils.cpp
+  ViewLikeInterfaceUtils.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/Affine
 
   LINK_LIBS PUBLIC
   MLIRAffineDialect
   MLIRAffineAnalysis
   MLIRAnalysis
   MLIRArithmeticUtils
   MLIRMemRefDialect
   MLIRTransformUtils
+  MLIRViewLikeInterface
   )

mlir/lib/Dialect/Affine/Utils/ViewLikeInterfaceUtils.cpp

+//===- 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);
+}

mlir/lib/Dialect/Tensor/Transforms/CMakeLists.txt

 add_mlir_dialect_library(MLIRTensorTransforms
   BufferizableOpInterfaceImpl.cpp
   Bufferize.cpp
   ExtractSliceFromReshapeUtils.cpp
   MergeConsecutiveInsertExtractSlicePatterns.cpp
   SplitPaddingPatterns.cpp
   SwapExtractSliceWithProducerPatterns.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/Tensor/Transforms
 
   DEPENDS
   MLIRTensorTransformsIncGen
 
   LINK_LIBS PUBLIC
   MLIRAffineDialect
+  MLIRAffineUtils
   MLIRArithmeticDialect
   MLIRBufferizationDialect
   MLIRBufferizationTransforms
   MLIRIR
   MLIRMemRefDialect
   MLIRPass
   MLIRSCFDialect
   MLIRTensorDialect
   MLIRTilingInterface
   MLIRTransforms
   )

mlir/lib/Dialect/Tensor/Transforms/MergeConsecutiveInsertExtractSlicePatterns.cpp

 //===- MergeConsecutiveInsertExtractSlicePatterns.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/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"
 #include "mlir/IR/PatternMatch.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> {
   using OpRewritePattern::OpRewritePattern;
 
   LogicalResult matchAndRewrite(ExtractSliceOp nextOp,
                                 PatternRewriter &rewriter) const override {
     auto prevOp = nextOp.getSource().getDefiningOp<ExtractSliceOp>();

utils/bazel/llvm-project-overlay/mlir/BUILD.bazel

srcs = glob(
             "lib/Dialect/Affine/Utils/*.cpp",
             "lib/Dialect/Affine/Utils/*.h",
         ],
     ),
     hdrs = [
         "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 = [
         ":AffineAnalysis",
         ":AffineDialect",
         ":Analysis",
         ":ArithmeticUtils",
         ":FuncDialect",
         ":IR",
         ":MemRefDialect",
         ":SCFDialect",
         ":Support",
         ":TransformUtils",
+        ":ViewLikeInterface",
         "//llvm:Support",
     ],
 )
 
 gentbl_cc_library(
     name = "AffinePassIncGen",
     strip_include_prefix = "include",
     tbl_outs = [
hdrs = [
         "include/mlir/Dialect/Tensor/Transforms/BufferizableOpInterfaceImpl.h",
         "include/mlir/Dialect/Tensor/Transforms/Passes.h",
         "include/mlir/Dialect/Tensor/Transforms/TransformUtils.h",
         "include/mlir/Dialect/Tensor/Transforms/Transforms.h",
     ],
     includes = ["include"],
     deps = [
         ":AffineDialect",
+        ":AffineUtils",
         ":ArithmeticDialect",
         ":ArithmeticUtils",
         ":BufferizationDialect",
         ":BufferizationTransforms",
         ":DialectUtils",
         ":FuncDialect",
         ":IR",
         ":MemRefDialect",