[mlir][linalg] move isElementwise() to Linalg/Utils (NFC)

[aartbik]

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

152

just looking at the style in this file, it seems to make sense to move this static method up, and keep this region for just the public API, but the end result is the same, so up to you....

[aartbik]

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

22

This change pulls in a dependence for Linalg.

Don't you need a CMake and Bazel file change for this?

[okkwon]

• Expose hasOnlyScalarElementwiseOp as a global function.
• Specify a dependency to FuncDialect in the CMake and Bazel files.

[okkwon]

Thanks for the review. BTW, how can I test the changes on the cmake and bazel files? I usually run "ninja check-mlir" for testing.

[aartbik]

In D128398#3603466, @okkwon wrote:

Thanks for the review. BTW, how can I test the changes on the cmake and bazel files? I usually run "ninja check-mlir" for testing.

I always find building with -DBUILD_SHARED_LIBS=ON the config a very good way of making sure no dependencies are missed.

[aartbik]

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

38

It is a bit strange to refer to r now that you removed the parameter name from the prototype

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

152

Ah, making this part of the public API is another way to "fix" this :-)

[okkwon]

Recover the missing r.

[okkwon]

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

38

It is not intended. Recovered it. Thanks for catching it!

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

 
 namespace linalg {
 class LinalgDependenceGraph;
 
 //===----------------------------------------------------------------------===//
 // General utilities
 //===----------------------------------------------------------------------===//
 
+/// Check if all indexing maps are projected permutations.
+bool allIndexingsAreProjectedPermutation(LinalgOp op);
+
+/// Detect whether `r` has only ConstantOp, ElementwiseMappable and YieldOp.
+bool hasOnlyScalarElementwiseOp(Region &r);
+
+/// Check if a LinalgOp is an element-wise operation.
+bool isElementwise(LinalgOp op);
+
 /// Check if `permutation` is a permutation of the range
 /// `[0, permutation.size())`.
 bool isPermutation(ArrayRef<int64_t> permutation);
 
 /// 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);
 

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

vectorizeOneOp(OpBuilder &b, LinalgOp linalgOp, Operation *op,
   // Build and return the new op.
   return VectorizationResult{
       VectorizationStatus::NewOp,
       b.create(op->getLoc(), op->getName().getIdentifier(),
                llvm::to_vector<4>(vectorizedOperands),
                llvm::to_vector<4>(returnTypes), op->getAttrs())};
 }
 
-/// Detect whether `r` has only ConstantOp, ElementwiseMappable and YieldOp.
-static bool hasOnlyScalarElementwiseOp(Region &r) {
-  if (!llvm::hasSingleElement(r))
-    return false;
-  for (Operation &op : r.front()) {
-    if (!(isa<arith::ConstantOp, func::ConstantOp, linalg::YieldOp,
-              linalg::IndexOp>(op) ||
-          OpTrait::hasElementwiseMappableTraits(&op)) ||
-        llvm::any_of(op.getResultTypes(),
-                     [](Type type) { return !type.isIntOrIndexOrFloat(); }))
-      return false;
-  }
-  return true;
-}
-
-/// Returns `true` if all indexing maps of the linalg op are projected
-/// permutations.
-static bool allIndexingsAreProjectedPermutation(LinalgOp op) {
-  return llvm::all_of(op.getIndexingMaps(), [](AffineMap m) {
-    return m.isProjectedPermutation(/*allowZeroInResults=*/true);
-  });
-}
-
-// Return true if the op is an element-wise linalg op.
-static bool isElementwise(Operation *op) {
-  auto linalgOp = dyn_cast<linalg::LinalgOp>(op);
-  if (!linalgOp)
-    return false;
-  if (linalgOp.getNumLoops() != linalgOp.getNumParallelLoops())
-    return false;
-
-  if (!allIndexingsAreProjectedPermutation(linalgOp))
-    return false;
-
-  // TODO: relax the restrictions on indexing map.
-  for (OpOperand *opOperand : linalgOp.getOutputOperands()) {
-    if (!linalgOp.getTiedIndexingMap(opOperand).isPermutation())
-      return false;
-  }
-  return hasOnlyScalarElementwiseOp(linalgOp->getRegion(0));
-}
-
 /// Generic vectorization function that rewrites the body of a `linalgOp` into
 /// vector form. Generic vectorization proceeds as follows:
 ///   1. Verify the `linalgOp` has one non-empty region.
 ///   2. Values defined above the region are mapped to themselves and will be
 ///   broadcasted on a per-need basis by their consumers.
 ///   3. Each region argument is vectorized into a vector.transfer_read (or 0-d
 ///   load).
 ///   TODO: Reuse opportunities for RAR dependencies.

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

 add_mlir_dialect_library(MLIRLinalgUtils
   Utils.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/Linalg
 
   LINK_LIBS PUBLIC
   MLIRAffineDialect
   MLIRAffineAnalysis
   MLIRAffineUtils
   MLIRArithmeticDialect
+  MLIRFuncDialect
   MLIRIR
   MLIRLinalgDialect
   MLIRSCFDialect
   MLIRPass
   MLIRTensorUtils
   MLIRTransformUtils
   )

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

 
 #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/Arithmetic/IR/Arithmetic.h"
 #include "mlir/Dialect/Arithmetic/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/StaticValueUtils.h"
 #include "mlir/IR/AffineExpr.h"
 #include "mlir/IR/AffineExprVisitor.h"
for (Range range : ranges) {
     ubs.emplace_back(range.size);
     steps.emplace_back(range.stride);
   }
 }
 
 namespace mlir {
 namespace linalg {
 
+bool allIndexingsAreProjectedPermutation(LinalgOp op) {
+  return llvm::all_of(op.getIndexingMaps(), [](AffineMap m) {
+    return m.isProjectedPermutation(/*allowZeroInResults=*/true);
+  });
+}
+
+bool hasOnlyScalarElementwiseOp(Region &r) {
+  if (!llvm::hasSingleElement(r))
+    return false;
+  for (Operation &op : r.front()) {
+    if (!(isa<arith::ConstantOp, func::ConstantOp, linalg::YieldOp,
+              linalg::IndexOp>(op) ||
+          OpTrait::hasElementwiseMappableTraits(&op)) ||
+        llvm::any_of(op.getResultTypes(),
+                     [](Type type) { return !type.isIntOrIndexOrFloat(); }))
+      return false;
+  }
+  return true;
+}
+
+bool isElementwise(LinalgOp op) {
+  if (op.getNumLoops() != op.getNumParallelLoops())
+    return false;
+
+  if (!allIndexingsAreProjectedPermutation(op))
+    return false;
+
+  // TODO: relax the restrictions on indexing map.
+  for (OpOperand *opOperand : op.getOutputOperands()) {
+    if (!op.getTiedIndexingMap(opOperand).isPermutation())
+      return false;
+  }
+  return hasOnlyScalarElementwiseOp(op->getRegion(0));
+}
+
 bool isPermutation(ArrayRef<int64_t> permutation) {
   // Count the number of appearances for all indices.
   SmallVector<int64_t> indexCounts(permutation.size(), 0);
   for (auto index : permutation) {
     // Exit if the index is out-of-range.
     if (index < 0 || index >= static_cast<int64_t>(permutation.size()))
       return false;
     indexCounts[index]++;

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

cc_library(
     deps = [
         ":AffineAnalysis",
         ":AffineDialect",
         ":AffineUtils",
         ":Analysis",
         ":ArithmeticDialect",
         ":ArithmeticUtils",
         ":DialectUtils",
+        ":FuncDialect",
         ":IR",
         ":LinalgAnalysis",
         ":LinalgDialect",
         ":MemRefDialect",
         ":Pass",
         ":SCFDialect",
         ":TensorDialect",
         ":TensorUtils",