feat: implement Hyperblock Fusion optimization

include/TaskflowDialect/Analysis/HyperblockDependencyAnalysis.h

@@ -0,0 +1,99 @@
+// HyperblockDependencyAnalysis.h - Analyzes dependencies between hyperblocks.
+//
+// This file provides utilities for analyzing data dependencies between
+// hyperblocks within a Taskflow task.
+
+#ifndef TASKFLOW_ANALYSIS_HYPERBLOCK_DEPENDENCY_ANALYSIS_H
+#define TASKFLOW_ANALYSIS_HYPERBLOCK_DEPENDENCY_ANALYSIS_H
+
+#include "TaskflowDialect/TaskflowOps.h"
+#include "llvm/ADT/DenseMap.h"
+#include "llvm/ADT/DenseSet.h"
+#include "llvm/ADT/SmallVector.h"
+
+namespace mlir {
+namespace taskflow {
+
+/// Represents the type of data dependency between hyperblocks.
+enum class DependencyType {
+  None,
+  RAW,  // Read-After-Write.
+  WAR,  // Write-After-Read.
+  WAW   // Write-After-Write.
+};
+
+/// Represents a dependency edge between two hyperblocks.
+struct HyperblockDependencyEdge {
+  TaskflowHyperblockOp source;
+  TaskflowHyperblockOp target;
+  DependencyType type;
+  Value memref;  // The memory location causing the dependency.
+};
+
+/// Analyzes dependencies between hyperblocks within a task.
+class HyperblockDependencyGraph {
+public:
+  /// Builds the dependency graph from a task operation.
+  void buildFromTask(TaskflowTaskOp taskOp);
+
+  /// Clears all stored dependency information.
+  void clear();
+
+  /// Returns true if there is any dependency from source to target.
+  bool hasDependency(TaskflowHyperblockOp source,
+                     TaskflowHyperblockOp target) const;
+
+  /// Returns all dependencies from source to target.
+  llvm::SmallVector<HyperblockDependencyEdge>
+  getDependencies(TaskflowHyperblockOp source,
+                  TaskflowHyperblockOp target) const;
+
+  /// Returns all predecessors of a hyperblock (hyperblocks it depends on).
+  llvm::SmallVector<TaskflowHyperblockOp>
+  getPredecessors(TaskflowHyperblockOp op) const;
+
+  /// Returns all successors of a hyperblock (hyperblocks that depend on it).
+  llvm::SmallVector<TaskflowHyperblockOp>
+  getSuccessors(TaskflowHyperblockOp op) const;
+
+  /// Checks if two hyperblocks can be fused without creating circular deps.
+  bool canFuse(TaskflowHyperblockOp a, TaskflowHyperblockOp b) const;
+
+  /// Checks if two hyperblocks have compatible counter structures.
+  bool areCountersCompatible(TaskflowHyperblockOp a, TaskflowHyperblockOp b,
+                             int maxBoundDiff) const;
+
+  /// Returns all hyperblocks in the analyzed task.
+  const llvm::SmallVector<TaskflowHyperblockOp> &getHyperblocks() const {
+    return hyperblocks_;
+  }
+
+private:
+  /// Collects memory reads from a hyperblock.
+  llvm::DenseSet<Value> collectReads(TaskflowHyperblockOp op) const;
+
+  /// Collects memory writes from a hyperblock.
+  llvm::DenseSet<Value> collectWrites(TaskflowHyperblockOp op) const;
+
+  /// Adds a dependency edge to the graph.
+  void addEdge(TaskflowHyperblockOp source, TaskflowHyperblockOp target,
+               DependencyType type, Value memref);
+
+  /// All hyperblocks in program order.
+  llvm::SmallVector<TaskflowHyperblockOp> hyperblocks_;
+
+  /// Maps each hyperblock to its predecessor edges.
+  llvm::DenseMap<TaskflowHyperblockOp,
+                 llvm::SmallVector<HyperblockDependencyEdge>>
+      predecessorEdges_;
+
+  /// Maps each hyperblock to its successor edges.
+  llvm::DenseMap<TaskflowHyperblockOp,
+                 llvm::SmallVector<HyperblockDependencyEdge>>
+      successorEdges_;
+};
+
+} // namespace taskflow
+} // namespace mlir
+
+#endif // TASKFLOW_ANALYSIS_HYPERBLOCK_DEPENDENCY_ANALYSIS_H

include/TaskflowDialect/TaskflowPasses.h

@@ -1,4 +1,4 @@
-// TaskflowPasses.h - Header file for Taskflow passes
+// TaskflowPasses.h - Header file for Taskflow passes.
 
 #ifndef TASKFLOW_PASSES_H
 #define TASKFLOW_PASSES_H
@@ -10,15 +10,23 @@
 #include "mlir/Pass/PassRegistry.h"
 
 #include <memory>
+
 namespace mlir {
 namespace taskflow {
-// Passes defined in TaskflowPasses.td
+
+// Passes defined in TaskflowPasses.td.
 #define GEN_PASS_DECL
 #include "TaskflowDialect/TaskflowPasses.h.inc"
+
+/// Creates a pass that constructs hyperblocks and counter chains from tasks.
 std::unique_ptr<mlir::Pass> createConstructHyperblockFromTaskPass();
 
+/// Creates a pass that optimizes the task graph by fusing hyperblocks and tasks.
+std::unique_ptr<mlir::Pass> createOptimizeTaskGraphPass();
+
 #define GEN_PASS_REGISTRATION
 #include "TaskflowDialect/TaskflowPasses.h.inc"
+
 } // namespace taskflow
 } // namespace mlir
 

include/TaskflowDialect/TaskflowPasses.td

@@ -15,4 +15,26 @@ def ConstructHyperblockFromTask : Pass<"construct-hyperblock-from-task", "func::
   }];
   let constructor = "taskflow::createConstructHyperblockFromTaskPass()";
 }
+
+def OptimizeTaskGraph : Pass<"optimize-task-graph", "func::FuncOp"> {
+  let summary = "Optimizes Taskflow task graph by fusing hyperblocks and tasks.";
+  let description = [{
+    Performs the following optimizations on the Taskflow task graph:
+    1. Hyperblock Fusion: Merges hyperblocks with compatible counter structures.
+       Supports loop peeling when counter bound differences are small.
+    2. Task Fusion: Merges producer-consumer tasks to reduce data transfer
+       overhead between tasks.
+    3. Dead Hyperblock Elimination: Removes unused hyperblocks.
+  }];
+  let constructor = "taskflow::createOptimizeTaskGraphPass()";
+  let options = [
+    Option<"enableHyperblockFusion", "enable-hyperblock-fusion", "bool",
+           /*default=*/"true", "Enables hyperblock fusion optimization.">,
+    Option<"enableTaskFusion", "enable-task-fusion", "bool",
+           /*default=*/"false", "Enables task fusion optimization (not yet implemented).">,
+    Option<"maxBoundDiffForPeeling", "max-bound-diff", "int",
+           /*default=*/"2", "Specifies max loop bound difference for peeling.">
+  ];
+}
+
 #endif // TASKFLOW_PASSES_TD

lib/TaskflowDialect/Analysis/CMakeLists.txt

@@ -0,0 +1,10 @@
+add_mlir_library(MLIRTaskflowAnalysis
+    HyperblockDependencyAnalysis.cpp
+    # TaskDependencyAnalysis.cpp
+
+    LINK_LIBS PUBLIC
+    MLIRIR
+    MLIRSupport
+    MLIRMemRefDialect
+    MLIRTaskflow
+)

lib/TaskflowDialect/Analysis/HyperblockDependencyAnalysis.cpp

@@ -0,0 +1,223 @@
+// HyperblockDependencyAnalysis.cpp - Implements hyperblock dependency analysis.
+
+#include "TaskflowDialect/Analysis/HyperblockDependencyAnalysis.h"
+#include "mlir/Dialect/MemRef/IR/MemRef.h"
+
+using namespace mlir;
+using namespace mlir::taskflow;
+
+void HyperblockDependencyGraph::buildFromTask(TaskflowTaskOp taskOp) {
+  clear();
+
+  // Collects all hyperblocks in program order.
+  taskOp.getBody().walk([&](TaskflowHyperblockOp op) {
+    hyperblocks_.push_back(op);
+  });
+
+  // Builds dependency edges between all pairs of hyperblocks.
+  for (size_t i = 0; i < hyperblocks_.size(); ++i) {
+    auto hbI = hyperblocks_[i];
+    auto writesI = collectWrites(hbI);
+    auto readsI = collectReads(hbI);
+
+    for (size_t j = i + 1; j < hyperblocks_.size(); ++j) {
+      auto hbJ = hyperblocks_[j];
+      auto writesJ = collectWrites(hbJ);
+      auto readsJ = collectReads(hbJ);
+
+      // Checks RAW: I writes, J reads.
+      for (Value memref : writesI) {
+        if (readsJ.contains(memref)) {
+          addEdge(hbI, hbJ, DependencyType::RAW, memref);
+        }
+      }
+
+      // Checks WAR: I reads, J writes.
+      for (Value memref : readsI) {
+        if (writesJ.contains(memref)) {
+          addEdge(hbI, hbJ, DependencyType::WAR, memref);
+        }
+      }
+
+      // Checks WAW: I writes, J writes.
+      for (Value memref : writesI) {
+        if (writesJ.contains(memref)) {
+          addEdge(hbI, hbJ, DependencyType::WAW, memref);
+        }
+      }
+    }
+  }
+}
+
+void HyperblockDependencyGraph::clear() {
+  hyperblocks_.clear();
+  predecessorEdges_.clear();
+  successorEdges_.clear();
+}
+
+bool HyperblockDependencyGraph::hasDependency(
+    TaskflowHyperblockOp source, TaskflowHyperblockOp target) const {
+  auto it = successorEdges_.find(source);
+  if (it == successorEdges_.end()) {
+    return false;
+  }
+  for (const auto &edge : it->second) {
+    if (edge.target == target) {
+      return true;
+    }
+  }
+  return false;
+}
+
+llvm::SmallVector<HyperblockDependencyEdge>
+HyperblockDependencyGraph::getDependencies(TaskflowHyperblockOp source,
+                                            TaskflowHyperblockOp target) const {
+  llvm::SmallVector<HyperblockDependencyEdge> result;
+  auto it = successorEdges_.find(source);
+  if (it != successorEdges_.end()) {
+    for (const auto &edge : it->second) {
+      if (edge.target == target) {
+        result.push_back(edge);
+      }
+    }
+  }
+  return result;
+}
+
+llvm::SmallVector<TaskflowHyperblockOp>
+HyperblockDependencyGraph::getPredecessors(TaskflowHyperblockOp op) const {
+  llvm::SmallVector<TaskflowHyperblockOp> result;
+  llvm::DenseSet<TaskflowHyperblockOp> seen;
+
+  auto it = predecessorEdges_.find(op);
+  if (it != predecessorEdges_.end()) {
+    for (const auto &edge : it->second) {
+      if (!seen.contains(edge.source)) {
+        seen.insert(edge.source);
+        result.push_back(edge.source);
+      }
+    }
+  }
+  return result;
+}
+
+llvm::SmallVector<TaskflowHyperblockOp>
+HyperblockDependencyGraph::getSuccessors(TaskflowHyperblockOp op) const {
+  llvm::SmallVector<TaskflowHyperblockOp> result;
+  llvm::DenseSet<TaskflowHyperblockOp> seen;
+
+  auto it = successorEdges_.find(op);
+  if (it != successorEdges_.end()) {
+    for (const auto &edge : it->second) {
+      if (!seen.contains(edge.target)) {
+        seen.insert(edge.target);
+        result.push_back(edge.target);
+      }
+    }
+  }
+  return result;
+}
+
+bool HyperblockDependencyGraph::canFuse(TaskflowHyperblockOp a,
+                                         TaskflowHyperblockOp b) const {
+  // Fusing two hyperblocks (A and B) is safe only if it does not violate 
+  // intermediate dependencies. Specifically, if there is a block C between 
+  // A and B in program order, we cannot fuse A and B if A -> C and C -> B.
+  // Fusing A and B would effectively move B before C, breaking C -> B.
+
+  // Finds positions in program order.
+  int posA = -1, posB = -1;
+  for (size_t i = 0; i < hyperblocks_.size(); ++i) {
+    if (hyperblocks_[i] == a) posA = i;
+    if (hyperblocks_[i] == b) posB = i;
+  }
+
+  if (posA < 0 || posB < 0) {
+    return false;
+  }
+
+  // Ensures a comes before b for fusion (or they are adjacent).
+  if (posA > posB) {
+    std::swap(a, b);
+    std::swap(posA, posB);
+  }
+
+  // Checks if there are any hyperblocks between a and b that depend on a
+  // and b depends on them (would create cycle after fusion).
+  for (size_t i = posA + 1; i < static_cast<size_t>(posB); ++i) {
+    auto middle = hyperblocks_[i];
+    if (hasDependency(a, middle) && hasDependency(middle, b)) {
+      return false;  // Fusion would break dependency chain.
+    }
+  }
+
+  return true;
+}
+
+bool HyperblockDependencyGraph::areCountersCompatible(
+    TaskflowHyperblockOp a, TaskflowHyperblockOp b, int maxBoundDiff) const {
+  auto indicesA = a.getIndices();
+  auto indicesB = b.getIndices();
+
+  // Requires same number of indices.
+  if (indicesA.size() != indicesB.size()) {
+    return false;
+  }
+
+  // Checks each counter pair.
+  for (size_t i = 0; i < indicesA.size(); ++i) {
+    auto counterA = indicesA[i].getDefiningOp<TaskflowCounterOp>();
+    auto counterB = indicesB[i].getDefiningOp<TaskflowCounterOp>();
+
+    if (!counterA || !counterB) {
+      return false;
+    }
+
+    int64_t lowerA = counterA.getLowerBound().getSExtValue();
+    int64_t upperA = counterA.getUpperBound().getSExtValue();
+    int64_t stepA = counterA.getStep().getSExtValue();
+
+    int64_t lowerB = counterB.getLowerBound().getSExtValue();
+    int64_t upperB = counterB.getUpperBound().getSExtValue();
+    int64_t stepB = counterB.getStep().getSExtValue();
+
+    // Requires same lower bound and step.
+    if (lowerA != lowerB || stepA != stepB) {
+      return false;
+    }
+
+    // Checks upper bound difference.
+    int diff = std::abs(static_cast<int>(upperA - upperB));
+    if (diff > maxBoundDiff) {
+      return false;
+    }
+  }
+
+  return true;
+}
+
+llvm::DenseSet<Value>
+HyperblockDependencyGraph::collectReads(TaskflowHyperblockOp op) const {
+  llvm::DenseSet<Value> reads;
+  op.getBody().walk([&](memref::LoadOp loadOp) {
+    reads.insert(loadOp.getMemRef());
+  });
+  return reads;
+}
+
+llvm::DenseSet<Value>
+HyperblockDependencyGraph::collectWrites(TaskflowHyperblockOp op) const {
+  llvm::DenseSet<Value> writes;
+  op.getBody().walk([&](memref::StoreOp storeOp) {
+    writes.insert(storeOp.getMemRef());
+  });
+  return writes;
+}
+
+void HyperblockDependencyGraph::addEdge(TaskflowHyperblockOp source,
+                                         TaskflowHyperblockOp target,
+                                         DependencyType type, Value memref) {
+  HyperblockDependencyEdge edge{source, target, type, memref};
+  successorEdges_[source].push_back(edge);
+  predecessorEdges_[target].push_back(edge);
+}