[mlir][Linalg] NFC - Extract an IndexingUtils from Linalg/Utils

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 AffineForOp;

 bool isElementwise(LinalgOp op);
 
 /// Check if iterator type has "parallel" semantics.
 bool isParallelIterator(utils::IteratorType iteratorType);
 
 /// Check if iterator type  has "reduction" semantics.
 bool isReductionIterator(utils::IteratorType iteratorType);
 
-/// Helper function that creates a memref::DimOp or tensor::DimOp depending on
-/// the type of `source`.
-Value createOrFoldDimOp(OpBuilder &b, Location loc, Value source, int64_t dim);
-OpFoldResult createFoldedDimOp(OpBuilder &b, Location loc, Value source,
-                               int64_t dim);
-
-/// Given an operation, retrieves the value of each dynamic dimension through
-/// constructing the necessary DimOp operators.
-SmallVector<Value, 4> getDynOperands(Location loc, Value val, OpBuilder &b);
-
 /// Create a tensor::PadOp that pads `source` to the size of the statically
 /// sized `type` whose static sizes are assumed to be greater than the dynamic
 /// `source` size. The padding introduces trailing `pad` values until the
 /// target size is met. If `source` is defined by one or more LinalgOps that
 /// have been padded with the same value and sizes, return their padded result
 /// instead of creating a tensor::PadOp.
 ///
 /// Example:

 
 /// Tensor counterpart of `fuseProducerOfBuffer`.
 /// This implements the fusion part of the "tileAndFuse on tensors"
 /// transformation and thus requires the `consumerOpOperand` to be a
 /// `extract_slice` op (generally obtained by applying the tiling
 /// transformation).
 FailureOr<FusionInfo> fuseProducerOfTensor(OpBuilder &b,
                                            OpOperand &consumerOpOperand);
+
 /// Tensor counterpart of `fuseProducerOfBuffer`.
 /// This implements the fusion part of the "tileAndFuse on tensors"
 /// transformation and thus requires the `consumerOpOperand` to be a
 /// `extract_slice` op (generally obtained by applying the tiling
 /// transformation). Assumes `producerOfTensor` is a Linalg op that produces
 /// `consumerOpOperand`.
 FailureOr<FusionInfo> fuseProducerOfTensor(OpBuilder &b,
                                            OpResult producerOpResult,

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

for (Type t : rankedTensorTypes) {
     }
     if (found)
       continue;
 
     // Extract static / dynamic shape mix from the first operand.
     Value firstOperand = operands.front();
     auto rankedTensorType = t.cast<RankedTensorType>();
     auto staticShape = llvm::to_vector<4>(rankedTensorType.getShape());
-    auto dynamicShape = linalg::getDynOperands(loc, firstOperand, b);
+    auto dynamicShape = linalg::createDynamicDimensions(b, loc, firstOperand);
 
     res.push_back(b.create<tensor::EmptyOp>(
         loc, staticShape, rankedTensorType.getElementType(), dynamicShape));
   }
   return res;
 }
 
 namespace {

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
   MLIRArithDialect
   MLIRFuncDialect
   MLIRIR
   MLIRLinalgDialect
   MLIRSCFDialect
   MLIRPass
   MLIRTensorUtils
   MLIRTransformUtils
-  )

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 (val.getType().isa<UnrankedMemRefType, MemRefType>())
+    return b.createOrFold<memref::DimOp>(loc, val, dim);
+  if (val.getType().isa<UnrankedTensorType, RankedTensorType>())
+    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 = val.getType().cast<ShapedType>();
+  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 = val.getType().cast<ShapedType>();
+  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 = val.getType().cast<ShapedType>();
+  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/lib/Dialect/Linalg/Utils/Utils.cpp

 bool isParallelIterator(utils::IteratorType iteratorType) {
   return iteratorType == utils::IteratorType::parallel;
 }
 
 bool isReductionIterator(utils::IteratorType iteratorType) {
   return iteratorType == utils::IteratorType::reduction;
 }
 
-/// Helper function that creates a memref::DimOp or tensor::DimOp depending on
-/// the type of `source`.
-Value createOrFoldDimOp(OpBuilder &b, Location loc, Value source, int64_t dim) {
-  if (source.getType().isa<UnrankedMemRefType, MemRefType>())
-    return b.createOrFold<memref::DimOp>(loc, source, dim);
-  if (source.getType().isa<UnrankedTensorType, RankedTensorType>())
-    return b.createOrFold<tensor::DimOp>(loc, source, dim);
-  llvm_unreachable("Expected MemRefType or TensorType");
-}
-
-OpFoldResult createFoldedDimOp(OpBuilder &b, Location loc, Value source,
-                               int64_t dim) {
-  auto shapedType = source.getType().cast<ShapedType>();
-  if (!shapedType.hasRank() || shapedType.isDynamicDim(dim))
-    return createOrFoldDimOp(b, loc, source, dim);
-  return b.getIndexAttr(shapedType.getDimSize(dim));
-}
-
-/// Given an operation, retrieves the value of each dynamic dimension through
-/// constructing the necessary DimOp operators.
-SmallVector<Value, 4> getDynOperands(Location loc, Value val, OpBuilder &b) {
-  SmallVector<Value, 4> dynOperands;
-  auto shapedType = val.getType().cast<ShapedType>();
-  for (const auto &dim : llvm::enumerate(shapedType.getShape())) {
-    if (dim.value() == ShapedType::kDynamic)
-      dynOperands.push_back(createOrFoldDimOp(b, loc, val, dim.index()));
-  }
-  return dynOperands;
-}
-
 Value makeComposedPadHighOp(OpBuilder &b, Location loc, RankedTensorType type,
                             Value source, Value pad, bool nofold) {
   // Exit if `source` is not defined by an ExtractSliceOp.
   auto sliceOp = source.getDefiningOp<tensor::ExtractSliceOp>();
   if (!sliceOp)
     return tensor::createPadHighOp(type, source, pad, nofold, loc, b);
 
   // Search the `source` use-def chain for padded LinalgOps.