Generalize cost model interface in stream analysis

PiperOrigin-RevId: 507822487

Author: cky9301

include/tfrt/compiler/stream_analysis.h

@@ -49,6 +49,8 @@
 #ifndef TFRT_COMPILER_STREAM_ANALYSIS_H_
 #define TFRT_COMPILER_STREAM_ANALYSIS_H_
 
+#include <optional>
+
 #include "llvm/ADT/SetVector.h"
 #include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/IR/Block.h"
@@ -119,8 +121,21 @@ class Stream {
 // function.
 class StreamAnalysis {
  public:
-  explicit StreamAnalysis(mlir::func::FuncOp op) : StreamAnalysis(op.front()) {}
-  explicit StreamAnalysis(mlir::Block& block) { AnalyzeBlock(block); }
+  class CostModelInterface {
+   public:
+    virtual ~CostModelInterface();
+
+    // The implementation is expected to return a positive value or
+    // std::nullopt if a cost cannot be computed. If std::nullopt is returned,
+    // stream analysis will use cost threshold as the cost for this op.
+    virtual std::optional<int64_t> GetOperationCost(
+        mlir::Operation* op) const = 0;
+  };
+
+  explicit StreamAnalysis(mlir::func::FuncOp op,
+                          const CostModelInterface* cost_model = nullptr);
+  explicit StreamAnalysis(mlir::Block& block,
+                          const CostModelInterface* cost_model = nullptr);
 
   // Return the stream that contains `op`. An operation can only belong to one
   // stream.
@@ -228,6 +243,8 @@ class StreamAnalysis {
 
   // `stream_map_` contains the finalized op-to-stream mapping.
   llvm::DenseMap<mlir::Operation*, Stream*> stream_map_;
+
+  const CostModelInterface* cost_model_ = nullptr;
 };
 
 }  // namespace compiler

lib/compiler/stream_analysis.cc

@@ -18,6 +18,8 @@
 
 #include "tfrt/compiler/stream_analysis.h"
 
+#include <optional>
+
 #include "tfrt/basic_kernels/opdefs/basic_kernels.h"
 #include "tfrt/basic_kernels/opdefs/types.h"
 #include "tfrt/compiler/opdefs/tfrt_op_interfaces.h"
@@ -67,20 +69,49 @@ bool GetMergeInterDependentStreams(mlir::Block& block) {
   return false;
 }
 
+class DefaultCostModel : public StreamAnalysis::CostModelInterface {
+ public:
+  std::optional<int64_t> GetOperationCost(mlir::Operation* op) const override {
+    // Check if operations defines a cost function.
+    if (auto cost_function = mlir::dyn_cast<CostFunctionInterface>(op)) {
+      int64_t cost = cost_function.cost();
+      assert(cost > 0 && "cost must be a positive value");
+      return cost;
+    }
+
+    return std::nullopt;
+  }
+};
+
+const DefaultCostModel* GetDefaultCostModel() {
+  static const auto* const default_cost_model = new DefaultCostModel();
+  return default_cost_model;
+}
+
 }  // namespace
 
+StreamAnalysis::CostModelInterface::~CostModelInterface() = default;
+
+StreamAnalysis::StreamAnalysis(mlir::func::FuncOp op,
+                               const CostModelInterface* cost_model)
+    : StreamAnalysis(op.front(), cost_model) {}
+StreamAnalysis::StreamAnalysis(mlir::Block& block,
+                               const CostModelInterface* cost_model)
+    : cost_model_(cost_model) {
+  if (cost_model_ == nullptr) cost_model_ = GetDefaultCostModel();
+  AnalyzeBlock(block);
+}
+
 int64_t StreamAnalysis::GetOperationCost(mlir::Operation* op) const {
   // Root has the lowest cost.
   if (op == kRootOperation) return 1;
 
   // A few TFRT kernels are guaranteed to be cheap.
-  if (llvm::isa<ReturnOp, MergeChainsOp>(op)) return 1;
+  if (llvm::isa<mlir::func::ReturnOp, ReturnOp, MergeChainsOp>(op)) return 1;
 
-  // Check if operations defines a cost function.
-  if (auto cost_function = mlir::dyn_cast<CostFunctionInterface>(op)) {
-    int64_t cost = cost_function.cost();
-    assert(cost > 0 && "cost must be a positive value");
-    return cost;
+  assert(cost_model_);
+  if (auto cost = cost_model_->GetOperationCost(op)) {
+    return *cost;
   }
 
   // If there is no cost specified for this operation, We conservatively return