Add implementation and tests related to fusion feasibility analysis.

Author: Alwaysproblem

mlir/optimization/scheduler/CMakeLists.txt

@@ -15,6 +15,8 @@ add_executable(
   lib/BufferAnalysisPass.cpp
   lib/LivenessAnalysisPass.cpp
   lib/MemrefLifetime.cpp
+  lib/FusionFeasibility.cpp
+  lib/LivenessAdapter.cpp
   )
 
 # add_dependencies(lab-scheduler ToyCh6ShapeInferenceInterfaceIncGen

mlir/optimization/scheduler/include/lab/FusionFeasibility.h

@@ -0,0 +1,68 @@
+#pragma once
+
+#include "mlir/Analysis/Liveness.h"
+#include "mlir/IR/Operation.h"
+#include "mlir/IR/Value.h"
+#include "llvm/ADT/SmallVector.h"
+#include <cstdint>
+
+#include "mlir/Analysis/Liveness.h"
+#include "mlir/IR/Operation.h"
+#include "mlir/IR/Value.h"
+#include "llvm/ADT/DenseMap.h"
+#include <cstdint>
+#include <optional>
+
+namespace mlir {
+namespace lab {
+
+enum class FusionReasonKind {
+  // positive
+  DirectUse,
+  SingleUse,
+  SupportedProducer,
+  SupportedConsumer,
+  ElementwiseConsumer,
+  StaticShapeCompatible,
+  NoSideEffect,
+  IntermediateEliminable,
+  PeakMemoryAcceptable,
+  TrafficReductionExpected,
+
+  // negative
+  NullProducer,
+  MultiUseProducer,
+  UnsupportedProducer,
+  UnsupportedConsumer,
+  NonElementwiseConsumer,
+  DynamicShapeUnsupported,
+  ShapeMismatch,
+  SideEffectingProducer,
+  SideEffectingConsumer,
+  IntermediateNotEliminable,
+  PeakMemoryTooHigh
+};
+
+struct FusionFeasibilityResult {
+  Operation *producer = nullptr;
+  Operation *consumer = nullptr;
+  OpOperand *fusedOperand = nullptr;
+
+  bool isFusable = false;
+  bool isProfitable = false;
+
+  int64_t eliminatedIntermediateBytes = 0;
+  int64_t estimatedTrafficSavedBytes = 0;
+  int64_t estimatedPeakBeforeBytes = 0;
+  int64_t estimatedPeakAfterBytes = 0;
+  int64_t extraLivenessGrowthBytes = 0;
+
+  double score = 0.0;
+
+  llvm::SmallVector<FusionReasonKind> reasons;
+};
+
+const char *toString(FusionReasonKind kind);
+
+} // namespace lab
+} // namespace mlir

mlir/optimization/scheduler/include/lab/LabPasses.h

@@ -14,5 +14,6 @@ std::unique_ptr<Pass> createLabMatmulTilePass();
 std::unique_ptr<Pass> createLabPipelinePlanPass();
 std::unique_ptr<Pass> createLabLivenessPass();
 std::unique_ptr<Pass> createLabMemrefLifetimePass();
+std::unique_ptr<Pass> createLabFusionFeasibilityPass();
 
 } // namespace mlir

mlir/optimization/scheduler/include/lab/LivenessAdapter.h

@@ -0,0 +1,49 @@
+#pragma once
+
+#include "mlir/Analysis/Liveness.h"
+#include "mlir/IR/Operation.h"
+#include "mlir/IR/Value.h"
+#include "llvm/ADT/DenseMap.h"
+#include <cstdint>
+#include <optional>
+
+namespace mlir {
+namespace lab {
+
+struct ValueLifetimeInfo {
+  int64_t start = -1;
+  int64_t end = -1;
+  int64_t sizeBytes = 0;
+};
+
+class LivenessAdapter {
+public:
+  virtual ~LivenessAdapter() = default;
+  virtual std::optional<ValueLifetimeInfo> lookup(Value v) const = 0;
+  virtual int64_t getPeakLiveBytes(Operation *scope) const = 0;
+};
+
+class MlirLivenessAdapter final : public LivenessAdapter {
+public:
+  explicit MlirLivenessAdapter(Operation *scope);
+
+  std::optional<ValueLifetimeInfo> lookup(Value v) const override;
+  int64_t getPeakLiveBytes(Operation *scope) const override;
+
+private:
+  Operation *topScope;
+  mlir::Liveness liveness;
+
+  llvm::DenseMap<Operation *, int64_t> opOrder;
+  llvm::DenseMap<Value, ValueLifetimeInfo> valueInfo;
+  int64_t cachedPeakLiveBytes = 0;
+
+  void buildOperationOrder(Operation *scope);
+  void buildValueLifetimeInfo(Operation *scope);
+  void buildPeakLiveBytes(Operation *scope);
+
+  static int64_t getValueSizeBytes(Value v);
+};
+
+} // namespace lab
+} // namespace mlir

mlir/optimization/scheduler/lab-opt.cpp

@@ -19,26 +19,31 @@ int main(int argc, char **argv) {
   registry.insert<mlir::func::FuncDialect, mlir::linalg::LinalgDialect,
                   mlir::arith::ArithDialect, mlir::tensor::TensorDialect,
                   mlir::memref::MemRefDialect, mlir::scf::SCFDialect,
-                  mlir::affine::AffineDialect,
-                  mlir::cf::ControlFlowDialect>();
+                  mlir::affine::AffineDialect, mlir::cf::ControlFlowDialect>();
 
   mlir::registerAllPasses();
   mlir::PassPipelineRegistration<>("lab-op-stats", "Lab Op Stats Pass",
                                    [](mlir::OpPassManager &pm) {
                                      pm.addPass(mlir::createLabOpStatsPass());
                                    });
-  mlir::PassPipelineRegistration<>("lab-buffer-stats", "Lab Buffer Stats Pass",
-                                   [](mlir::OpPassManager &pm) {
-                                     pm.addPass(mlir::createLabBufferStatsPass());
-                                   });
+  mlir::PassPipelineRegistration<>(
+      "lab-buffer-stats", "Lab Buffer Stats Pass", [](mlir::OpPassManager &pm) {
+        pm.addPass(mlir::createLabBufferStatsPass());
+      });
   mlir::PassPipelineRegistration<>("lab-liveness", "Lab Liveness Pass",
                                    [](mlir::OpPassManager &pm) {
                                      pm.addPass(mlir::createLabLivenessPass());
                                    });
-  mlir::PassPipelineRegistration<>("lab-memref-lifetime", "Lab Memref Lifetime Pass",
-                                   [](mlir::OpPassManager &pm) {
-                                     pm.addPass(mlir::createLabMemrefLifetimePass());
-                                   });
+  mlir::PassPipelineRegistration<>(
+      "lab-memref-lifetime", "Lab Memref Lifetime Pass",
+      [](mlir::OpPassManager &pm) {
+        pm.addPass(mlir::createLabMemrefLifetimePass());
+      });
+  mlir::PassPipelineRegistration<>(
+      "lab-fusion-feasibility", "Lab Fusion Feasibility Pass",
+      [](mlir::OpPassManager &pm) {
+        pm.addPass(mlir::createLabFusionFeasibilityPass());
+      });
 
   return mlir::asMainReturnCode(
       mlir::MlirOptMain(argc, argv, "Lab optimizer\n", registry));