[mlir][linalg] Remove redundant dimension size helper functions

mlir/include/mlir/Dialect/Linalg/IR/Linalg.h

 /// symbol-less identity map of `rank`.
 AffineMap extractOrIdentityMap(std::optional<AffineMap> maybeMap, unsigned rank,
                                MLIRContext *context);
 
 /// Return the vector that is the concatenation of `a` and `b`.
 SmallVector<AffineExpr, 4> concat(ArrayRef<AffineExpr> a,
                                   ArrayRef<AffineExpr> b);
 
+/// Create one memref::DimOp or tensor::DimOp depending on the type of `val`.
+/// This is a polymorphic convenience function to abstract away the rank and
+/// concrete type of `val`.
+/// Asserts that `val` is a memref or tensor type.
+Value createOrFoldDimOp(OpBuilder &b, Location loc, Value val, int64_t dim);
+
+/// Create one memref::DimOp or tensor::DimOp depending on the type of `val`.
+/// This is a polymorphic convenience function to abstract away the rank and
+/// concrete type of `val`.
+/// Asserts that `val` is a memref or tensor type.
+OpFoldResult createFoldedDimOp(OpBuilder &b, Location loc, Value val,
+                               int64_t dim);
+
 } // namespace linalg
 } // namespace mlir
 
 //===----------------------------------------------------------------------===//
 // Linalg Dialect
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/Linalg/IR/LinalgOpsDialect.h.inc"

mlir/include/mlir/Dialect/Linalg/Utils/IndexingUtils.h

-//===- IndexingUtils.h - Indexing utilities supporting Linalg ---*- 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_LINALG_UTILS_INDEXINGUTILS_H
-#define MLIR_DIALECT_LINALG_UTILS_INDEXINGUTILS_H
-
-#include "mlir/Dialect/Linalg/IR/Linalg.h"
-#include "mlir/Dialect/SCF/IR/SCF.h"
-#include "mlir/Dialect/Utils/StructuredOpsUtils.h"
-#include "llvm/ADT/StringSet.h"
-#include <optional>
-
-namespace mlir {
-namespace linalg {
-
-/// Create one memref::DimOp or tensor::DimOp depending on the type of `val`.
-/// This is a polymorphic convenience function to abstract away the rank and
-/// concrete type of `val`.
-/// Asserts that `val` is a memref or tensor type.
-Value createOrFoldDimOp(OpBuilder &b, Location loc, Value val, int64_t dim);
-
-/// Create one memref::DimOp or tensor::DimOp depending on the type of `val`.
-/// This is a polymorphic convenience function to abstract away the rank and
-/// concrete type of `val`.
-/// Asserts that `val` is a memref or tensor type.
-OpFoldResult createFoldedDimOp(OpBuilder &b, Location loc, Value val,
-                               int64_t dim);
-
-/// Build the list of DimOp for the dynamic dimensions of `val`.
-/// Asserts that `val` is a ranked shaped type.
-SmallVector<Value> createDynamicDimensions(OpBuilder &b, Location loc,
-                                           Value val);
-
-/// Build the list of all dimensions for `val`, mixing static attributes and
-/// dynamic values where appropriate.
-/// Asserts that `val` is a ranked shaped type.
-SmallVector<OpFoldResult> getMixedDimensions(OpBuilder &b, Location loc,
-                                             Value val);
-
-} // namespace linalg
-} // namespace mlir
-#endif // MLIR_DIALECT_LINALG_UTILS_INDEXINGUTILS_H

mlir/include/mlir/Dialect/Linalg/Utils/Utils.h

 //===- Utils.h - Utilities to support the Linalg dialect --------*- 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_LINALG_UTILS_UTILS_H
 #define MLIR_DIALECT_LINALG_UTILS_UTILS_H
 
 #include "mlir/Dialect/Linalg/IR/Linalg.h"
-#include "mlir/Dialect/Linalg/Utils/IndexingUtils.h"
 #include "mlir/Dialect/SCF/IR/SCF.h"
 #include "mlir/Dialect/Utils/StructuredOpsUtils.h"
 #include "llvm/ADT/StringSet.h"
 #include <optional>
 
 namespace mlir {
 class AffineExpr;
 class AffineMap;

mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp

 //===- TosaToLinalg.cpp - Lowering Tosa to 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
 //
 //===----------------------------------------------------------------------===//
 //
 // These rewriters lower from the Tosa to the Linalg dialect.
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Conversion/TosaToLinalg/TosaToLinalg.h"
 #include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Dialect/Arith/Utils/Utils.h"
 #include "mlir/Dialect/Linalg/IR/Linalg.h"
-#include "mlir/Dialect/Linalg/Utils/IndexingUtils.h"
 #include "mlir/Dialect/Math/IR/Math.h"
 #include "mlir/Dialect/SCF/IR/SCF.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/Dialect/Tensor/Utils/Utils.h"
 #include "mlir/Dialect/Tosa/IR/TosaOps.h"
 #include "mlir/Dialect/Tosa/Utils/ConversionUtils.h"
 #include "mlir/Dialect/Utils/ReshapeOpsUtils.h"
 #include "mlir/Dialect/Utils/StaticValueUtils.h"
static OpFoldResult halfPlusOne(OpBuilder &builder, Location loc,
     auto plusOne = builder.createOrFold<arith::AddIOp>(loc, divBy2, one);
     return getAsOpFoldResult(plusOne);
   }
 
   static RankedTensorType
   computeOutputShape(OpBuilder &builder, Location loc, Value input,
                      llvm::SmallVectorImpl<Value> &dynamicSizes) {
     // Get [N, H, W]
-    auto dims = linalg::getMixedDimensions(builder, loc, input);
+    auto dims = tensor::getMixedSizes(builder, loc, input);
 
     // Set W = (W / 2) + 1 to account for the half-sized W dimension of the
     // output tensors.
     dims[2] = halfPlusOne(builder, loc, dims[2]);
 
     llvm::SmallVector<int64_t, 3> staticSizes;
     dispatchIndexOpFoldResults(dims, dynamicSizes, staticSizes);
 

mlir/lib/Dialect/Linalg/IR/LinalgInterfaces.cpp

 //===- LinalgInterfaces.cpp - Linalg interfaces implementation ------------===//
 //
 // 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/Linalg/IR/LinalgInterfaces.h"
 
 #include "mlir/Dialect/Affine/IR/AffineOps.h"
 #include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Dialect/Arith/Utils/Utils.h"
 #include "mlir/Dialect/Complex/IR/Complex.h"
+#include "mlir/Dialect/Linalg/IR/Linalg.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/IR/AffineExprVisitor.h"
 #include "mlir/IR/AffineMap.h"
 #include "mlir/IR/TypeUtilities.h"
 #include "llvm/ADT/SetOperations.h"
 #include "llvm/ADT/SmallBitVector.h"
 #include "llvm/ADT/SmallVector.h"
LogicalResult mlir::linalg::detail::verifyFillInterface(Operation *op) {
 
   return success();
 }
 
 //===----------------------------------------------------------------------===//
 // StructuredOpInterface implementation
 //===----------------------------------------------------------------------===//
 
-/// Helper function that creates a memref::DimOp or tensor::DimOp depending on
-/// the type of `source`.
-static Value createOrFoldDimOp(OpBuilder &b, Location loc, Value source,
-                               int64_t dim) {
-  if (llvm::isa<UnrankedMemRefType, MemRefType>(source.getType()))
-    return b.createOrFold<memref::DimOp>(loc, source, dim);
-  if (llvm::isa<UnrankedTensorType, RankedTensorType>(source.getType()))
-    return b.createOrFold<tensor::DimOp>(loc, source, dim);
-  llvm_unreachable("Expected MemRefType or TensorType");
-}
-static OpFoldResult createFoldedDimOp(OpBuilder &b, Location loc, Value source,
-                                      int64_t dim) {
-  auto shapedType = llvm::cast<ShapedType>(source.getType());
-  if (!shapedType.hasRank() || shapedType.isDynamicDim(dim))
-    return createOrFoldDimOp(b, loc, source, dim);
-  return b.getIndexAttr(shapedType.getDimSize(dim));
-}
-
 SmallVector<OpFoldResult> LinalgOp::createFlatListOfOperandDims(OpBuilder &b,
                                                                 Location loc) {
   SmallVector<OpFoldResult> res;
   for (OpOperand &opOperand : getOperation()->getOpOperands()) {
     for (int64_t i = 0, e = getRank(&opOperand); i < e; ++i)
       res.push_back(createFoldedDimOp(b, loc, opOperand.get(), i));
   }
   return res;

mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp

 #include "llvm/Support/MathExtras.h"
 #include "llvm/Support/raw_ostream.h"
 #include <optional>
 
 using namespace mlir;
 using namespace mlir::linalg;
 
 //===----------------------------------------------------------------------===//
+// Helper functions
+//===----------------------------------------------------------------------===//
+
+Value linalg::createOrFoldDimOp(OpBuilder &b, Location loc, Value source,
+                                int64_t dim) {
+  if (llvm::isa<UnrankedMemRefType, MemRefType>(source.getType()))
+    return b.createOrFold<memref::DimOp>(loc, source, dim);
+  if (llvm::isa<UnrankedTensorType, RankedTensorType>(source.getType()))
+    return b.createOrFold<tensor::DimOp>(loc, source, dim);
+  llvm_unreachable("Expected MemRefType or TensorType");
+}
+
+OpFoldResult linalg::createFoldedDimOp(OpBuilder &b, Location loc, Value source,
+                                       int64_t dim) {
+  auto shapedType = llvm::cast<ShapedType>(source.getType());
+  if (!shapedType.hasRank() || shapedType.isDynamicDim(dim))
+    return createOrFoldDimOp(b, loc, source, dim);
+  return b.getIndexAttr(shapedType.getDimSize(dim));
+}
+
+//===----------------------------------------------------------------------===//
 // Support for named Linalg ops defined in ods-gen.
 //===----------------------------------------------------------------------===//
 
 using RegionBuilderFn = llvm::function_ref<void(ImplicitLocOpBuilder &, Block &,
                                                 ArrayRef<NamedAttribute>)>;
 
 /// Fills the region of a structured operation using the provided
 /// `regionBuilder`. The method is used by both named structured ops created by

mlir/lib/Dialect/Linalg/Transforms/ElementwiseToLinalg.cpp

for (Value v : operands) {
         res.push_back(v);
         break;
       }
     }
     if (found)
       continue;
 
     // Extract static / dynamic shape mix from the first operand.
-    Value firstOperand = operands.front();
-    auto rankedTensorType = cast<RankedTensorType>(t);
-    auto staticShape = llvm::to_vector<4>(rankedTensorType.getShape());
-    auto dynamicShape = linalg::createDynamicDimensions(b, loc, firstOperand);
-
     res.push_back(b.create<tensor::EmptyOp>(
-        loc, staticShape, rankedTensorType.getElementType(), dynamicShape));
+        loc, tensor::getMixedSizes(b, loc, operands.front()),
+        cast<RankedTensorType>(t).getElementType()));
   }
   return res;
 }
 
 namespace {
 struct ConvertAnyElementwiseMappableOpOnRankedTensors : public RewritePattern {
   ConvertAnyElementwiseMappableOpOnRankedTensors(MLIRContext *context)
       : RewritePattern(MatchAnyOpTypeTag(), /*benefit=*/1, context) {}

mlir/lib/Dialect/Linalg/Transforms/Transforms.cpp

if (packOp.isLikePad()) {
     // Just insert the pad in the higher ranked tensor.
     auto emptyOp =
         rewriter.create<tensor::EmptyOp>(loc, packedTensorType, ValueRange{});
     // Offsets.
     SmallVector<OpFoldResult> zeros(packedRank, rewriter.getIndexAttr(0));
     // Strides.
     SmallVector<OpFoldResult> ones(packedRank, rewriter.getIndexAttr(1));
     SmallVector<OpFoldResult> sizes =
-        getMixedDimensions(rewriter, loc, packOp.getDest());
+        tensor::getMixedSizes(rewriter, loc, packOp.getDest());
 
     auto insertSliceOp = rewriter.create<tensor::InsertSliceOp>(
         loc, /*source=*/padOp, /*dest=*/emptyOp,
         /*offsets=*/zeros, sizes,
         /*strides=*/ones);
 
     LLVM_DEBUG(DBGS() << "insert_slice op: " << insertSliceOp; DBGSNL(););
 
FailureOr<LowerUnPackOpResult> linalg::lowerUnPack(RewriterBase &rewriter,
   auto destTensorType = cast<RankedTensorType>(unPackOp.getDest().getType());
   if (unPackOp.isLikeUnPad()) {
     // This unpack is just a plain unpad.
     // Just extract the slice from the higher ranked tensor.
     ArrayRef<int64_t> destShape = destTensorType.getShape();
     // The inner dimensions stay the same as the destination tensor, but the
     // outer ones are additional 1s.
     SmallVector<OpFoldResult> sizes(packedRank - destShape.size(), one);
-    sizes.append(getMixedDimensions(rewriter, loc, unPackOp.getDest()));
+    sizes.append(tensor::getMixedSizes(rewriter, loc, unPackOp.getDest()));
 
     auto extractSliceOp = rewriter.create<tensor::ExtractSliceOp>(
         loc, destTensorType, unPackOp.getSource(),
         SmallVector<OpFoldResult>(packedRank, zero), sizes,
         SmallVector<OpFoldResult>(packedRank, one));
 
     rewriter.replaceOp(unPackOp, extractSliceOp->getResults());
 

mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp

vectorizeAsTensorPadOp(RewriterBase &rewriter, tensor::PadOp padOp,
   LogicalResult status =
       cast<ReifyRankedShapedTypeOpInterface>(padOp.getOperation())
           .reifyResultShapes(rewriter, reifiedReturnShapes);
   (void)status; // prevent unused variable warning on non-assert builds
   assert(succeeded(status) && "failed to reify result shapes");
   auto emptyOp = rewriter.create<tensor::EmptyOp>(loc, reifiedReturnShapes[0],
                                                   padValue.getType());
   SmallVector<OpFoldResult> mixedSourceDims =
-      getMixedDimensions(rewriter, loc, padOp.getSource());
+      tensor::getMixedSizes(rewriter, loc, padOp.getSource());
   Value mask =
       rewriter.create<vector::CreateMaskOp>(loc, maskType, mixedSourceDims);
   auto zero = rewriter.create<arith::ConstantIndexOp>(loc, 0);
   auto transferReadOp = rewriter.create<vector::TransferReadOp>(
       loc,
       /*vectorType=*/vectorType,
       /*source=*/padOp.getSource(),
       /*indices=*/SmallVector<Value>(rank, zero),

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

 add_mlir_dialect_library(MLIRLinalgUtils
   Utils.cpp
-  IndexingUtils.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/Linalg
 
   LINK_LIBS PUBLIC
   MLIRAffineDialect
   MLIRAffineAnalysis
   MLIRAffineUtils

mlir/lib/Dialect/Linalg/Utils/IndexingUtils.cpp

-//===- IndexingUtils.cpp - Indexing utilities supporting Linalg -----------===//
-//
-// 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 indexing utilities for the Linalg dialect.
-//
-//===----------------------------------------------------------------------===//
-
-#include "mlir/Dialect/Linalg/Utils/Utils.h"
-
-#include "mlir/Analysis/SliceAnalysis.h"
-#include "mlir/Dialect/Affine/Analysis/AffineStructures.h"
-#include "mlir/Dialect/Affine/IR/AffineOps.h"
-#include "mlir/Dialect/Affine/IR/AffineValueMap.h"
-#include "mlir/Dialect/Affine/LoopUtils.h"
-#include "mlir/Dialect/Arith/IR/Arith.h"
-#include "mlir/Dialect/Arith/Utils/Utils.h"
-#include "mlir/Dialect/Func/IR/FuncOps.h"
-#include "mlir/Dialect/Linalg/IR/Linalg.h"
-#include "mlir/Dialect/MemRef/IR/MemRef.h"
-#include "mlir/Dialect/SCF/IR/SCF.h"
-#include "mlir/Dialect/Tensor/IR/Tensor.h"
-#include "mlir/Dialect/Tensor/Utils/Utils.h"
-#include "mlir/Dialect/Utils/IndexingUtils.h"
-#include "mlir/Dialect/Utils/StaticValueUtils.h"
-#include "mlir/IR/AffineExpr.h"
-#include "mlir/IR/AffineExprVisitor.h"
-#include "mlir/IR/AffineMap.h"
-#include "mlir/IR/Matchers.h"
-#include "mlir/IR/OpImplementation.h"
-#include "mlir/Pass/Pass.h"
-#include "llvm/ADT/SetOperations.h"
-#include "llvm/ADT/TypeSwitch.h"
-#include "llvm/Support/Debug.h"
-#include <optional>
-
-#define DEBUG_TYPE "linalg-utils"
-
-namespace mlir {
-namespace linalg {
-Value createOrFoldDimOp(OpBuilder &b, Location loc, Value val, int64_t dim) {
-  if (isa<UnrankedMemRefType, MemRefType>(val.getType()))
-    return b.createOrFold<memref::DimOp>(loc, val, dim);
-  if (isa<UnrankedTensorType, RankedTensorType>(val.getType()))
-    return b.createOrFold<tensor::DimOp>(loc, val, dim);
-  llvm_unreachable("Expected MemRefType or TensorType");
-}
-
-OpFoldResult createFoldedDimOp(OpBuilder &b, Location loc, Value val,
-                               int64_t dim) {
-  auto shapedType = cast<ShapedType>(val.getType());
-  if (!shapedType.hasRank() || shapedType.isDynamicDim(dim))
-    return createOrFoldDimOp(b, loc, val, dim);
-  return b.getIndexAttr(shapedType.getDimSize(dim));
-}
-
-SmallVector<Value> createDynamicDimensions(OpBuilder &b, Location loc,
-                                           Value val) {
-  auto shapedType = cast<ShapedType>(val.getType());
-  assert(shapedType.hasRank() && "`val` must have a static rank");
-  SmallVector<Value> res;
-  res.reserve(shapedType.getRank());
-  for (const auto &dim : llvm::enumerate(shapedType.getShape())) {
-    if (dim.value() == ShapedType::kDynamic)
-      res.push_back(createOrFoldDimOp(b, loc, val, dim.index()));
-  }
-  return res;
-}
-
-SmallVector<OpFoldResult> getMixedDimensions(OpBuilder &b, Location loc,
-                                             Value val) {
-  auto shapedType = cast<ShapedType>(val.getType());
-  assert(shapedType.hasRank() && "`val` must have a static rank");
-  SmallVector<Value> dynamicDims = createDynamicDimensions(b, loc, val);
-  return getMixedValues(shapedType.getShape(), dynamicDims, b);
-}
-} // namespace linalg
-} // namespace mlir

mlir/test/Dialect/Linalg/generalize-tensor-pack.mlir

 func.func @simple_pad_and_pack(%input: tensor<5x1xf32>, %output: tensor<1x1x8x2xf32>, %pad: f32) -> tensor<1x1x8x2xf32> {
   %0 = tensor.pack %input padding_value(%pad : f32) inner_dims_pos = [0, 1] inner_tiles = [8, 2] into %output : tensor<5x1xf32> -> tensor<1x1x8x2xf32>
   return %0 : tensor<1x1x8x2xf32>
 }
 // CHECK-LABEL: func.func @simple_pad_and_pack
 // CHECK-SAME:    %[[SRC:[a-zA-Z0-9]+]]
 // CHECK-SAME:    %[[DEST:[a-zA-Z0-9]+]]
 // CHECK-SAME:    %[[PAD_VAL:[a-zA-Z0-9]+]]
-// CHECK-DAG:     %[[C0:.+]] = arith.constant 0 : index
-// CHECK:         %[[PAD:.+]] = tensor.pad %[[SRC]] low[%[[C0]], %[[C0]]] high[3, 1]
+// CHECK:         %[[PAD:.+]] = tensor.pad %[[SRC]] low[0, 0] high[3, 1]
 // CHECK:           tensor.yield %[[PAD_VAL]]
 // CHECK:         %[[EMPTY:.+]] = tensor.empty() : tensor<8x2xf32>
 // CHECK:         %[[TRANSP:.+]] = linalg.transpose
 // CHECK-SAME:      ins(%[[PAD]] : tensor<8x2xf32>)
 // CHECK-SAME:      outs(%[[EMPTY]] : tensor<8x2xf32>)
 // CHECK-SAME:      permutation = [0, 1]
 // CHECK:         %[[INSERT:.+]] = tensor.insert_slice %[[TRANSP]] into %[[DEST]]
 // CHECK-SAME:      [0, 0, 0, 0] [1, 1, 8, 2] [1, 1, 1, 1]