[mlir][Linalg] Add more utility functions to LinalgDependenceGraph.

Reviewed By: ThomasRaoux

Differential Revision: https://reviews.llvm.org/D90582

Author: MaheshRavishankar

mlir/include/mlir/Dialect/Linalg/Analysis/DependenceAnalysis.h

@@ -104,12 +104,42 @@ class LinalgDependenceGraph {
                          ArrayRef<DependenceType> depTypes = {
                              DependenceType::RAW, DependenceType::WAW}) const;
 
+  /// Returns true if the `linalgOp` has dependences into it.
+  bool hasDependentOperationsInto(LinalgOp linalgOp,
+                                  ArrayRef<DependenceType> depTypes = {
+                                      DependenceType::RAW,
+                                      DependenceType::WAW}) const;
+
+  /// Returns true if the `linalgOp` has dependences from it.
+  bool hasDependentOperationsFrom(LinalgOp linalgOp,
+                                  ArrayRef<DependenceType> depTypes = {
+                                      DependenceType::RAW,
+                                      DependenceType::WAW}) const;
+
   /// Returns true if the `linalgOp` has dependences into or from it.
   bool hasDependentOperations(LinalgOp linalgOp,
                               ArrayRef<DependenceType> depTypes = {
                                   DependenceType::RAW,
                                   DependenceType::WAW}) const;
 
+  /// Returns all operations that have a dependence into `linalgOp` of types
+  /// listed in `depTypes`.
+  SmallVector<LinalgDependenceGraphElem, 2> getDependentOperationsInto(
+      LinalgOp linalgOp, ArrayRef<DependenceType> depTypes = {
+                             DependenceType::RAW, DependenceType::WAW}) const;
+
+  /// Returns all operations that have a dependence from `linalgOp` of types
+  /// listed in `depTypes`.
+  SmallVector<LinalgDependenceGraphElem, 2> getDependentOperationsFrom(
+      LinalgOp linalgOp, ArrayRef<DependenceType> depTypes = {
+                             DependenceType::RAW, DependenceType::WAW}) const;
+
+  /// Returns all dependent operations (into and from) given `operation`.
+  SmallVector<LinalgDependenceGraphElem, 2>
+  getDependentOperations(LinalgOp linalgOp,
+                         ArrayRef<DependenceType> depTypes = {
+                             DependenceType::RAW, DependenceType::WAW}) const;
+
 private:
   // Keep dependences in both directions, this is not just a performance gain
   // but it also reduces usage errors.

mlir/lib/Dialect/Linalg/Analysis/DependenceAnalysis.cpp

@@ -251,13 +251,64 @@ bool LinalgDependenceGraph::hasDependenceFrom(
   return false;
 }
 
-bool LinalgDependenceGraph::hasDependentOperations(
+bool LinalgDependenceGraph::hasDependentOperationsFrom(
+    LinalgOp linalgOp,
+    ArrayRef<LinalgDependenceGraph::DependenceType> depTypes) const {
+  for (auto dep : depTypes) {
+    if (!getDependencesFrom(linalgOp, dep).empty())
+      return true;
+  }
+  return false;
+}
+
+bool LinalgDependenceGraph::hasDependentOperationsInto(
     LinalgOp linalgOp,
     ArrayRef<LinalgDependenceGraph::DependenceType> depTypes) const {
   for (auto dep : depTypes) {
-    if (!getDependencesFrom(linalgOp, dep).empty() ||
-        !getDependencesInto(linalgOp, dep).empty())
+    if (!getDependencesInto(linalgOp, dep).empty())
       return true;
   }
   return false;
 }
+
+bool LinalgDependenceGraph::hasDependentOperations(
+    LinalgOp linalgOp, ArrayRef<DependenceType> depTypes) const {
+  return hasDependentOperationsInto(linalgOp, depTypes) ||
+         hasDependentOperationsFrom(linalgOp, depTypes);
+}
+
+SmallVector<LinalgDependenceGraph::LinalgDependenceGraphElem, 2>
+LinalgDependenceGraph::getDependentOperationsInto(
+    LinalgOp linalgOp, ArrayRef<DependenceType> depTypes) const {
+  SmallVector<LinalgDependenceGraph::LinalgDependenceGraphElem, 2>
+      dependentOperations;
+  for (auto dependenceType : depTypes) {
+    auto dependencies = getDependencesInto(linalgOp, dependenceType);
+    dependentOperations.append(dependencies.begin(), dependencies.end());
+  }
+  return dependentOperations;
+}
+
+SmallVector<LinalgDependenceGraph::LinalgDependenceGraphElem, 2>
+LinalgDependenceGraph::getDependentOperationsFrom(
+    LinalgOp linalgOp, ArrayRef<DependenceType> depTypes) const {
+  SmallVector<LinalgDependenceGraph::LinalgDependenceGraphElem, 2>
+      dependentOperations;
+  for (auto dependenceType : depTypes) {
+    auto dependencies = getDependencesFrom(linalgOp, dependenceType);
+    dependentOperations.append(dependencies.begin(), dependencies.end());
+  }
+  return dependentOperations;
+}
+
+/// Returns all dependent operations (into and from) given `operation`.
+SmallVector<LinalgDependenceGraph::LinalgDependenceGraphElem, 2>
+LinalgDependenceGraph::getDependentOperations(
+    LinalgOp linalgOp, ArrayRef<DependenceType> depTypes) const {
+  SmallVector<LinalgDependenceGraphElem, 2> dependentOperations =
+      getDependentOperationsInto(linalgOp, depTypes);
+  SmallVector<LinalgDependenceGraphElem, 2> t =
+      getDependentOperationsFrom(linalgOp, depTypes);
+  dependentOperations.append(t.begin(), t.end());
+  return dependentOperations;
+}