Addresses review comments

Author: pranavm-nvidia

mlir-tensorrt/compiler/include/mlir-tensorrt-c/Compiler/Compiler.h

@@ -60,6 +60,9 @@ typedef struct MTRT_StableHLOToExecutableOptions {
   void *ptr;
 } MTRT_StableHLOToExecutableOptions;
 
+/// A callback that allows the user to customize the metadata set for layers
+/// corresponding to each MLIR operation. The callback should invoke the
+/// provided append function in order to manipulate the result string.
 typedef void (*MTRT_MetadataCallback)(MlirOperation op,
                                       MlirStringCallback append,
                                       void *appendCtx, void *userData);
@@ -81,6 +84,8 @@ MLIR_CAPI_EXPORTED MTRT_Status mtrtStableHloToExecutableOptionsSetDebugOptions(
     const char **debugTypes, size_t debugTypeSizes,
     const char *dumpIrTreeDir = nullptr, const char *dumpTensorRTDir = nullptr);
 
+/// Sets the layer metadata callback. The `userData` argument is passed along
+/// to the callback when it is invoked.
 MLIR_CAPI_EXPORTED MTRT_Status
 mtrtStableHloToExecutableOptionsSetTensorRTTranslationMetadataCallback(
     MTRT_StableHLOToExecutableOptions options, MTRT_MetadataCallback callback,

mlir-tensorrt/compiler/include/mlir-tensorrt/Compiler/Client.h

@@ -33,6 +33,7 @@
 #include "mlir/Pass/PassManager.h"
 #include "mlir/Support/TypeID.h"
 #include "llvm/ADT/DenseMap.h"
+#include "llvm/ADT/Hashing.h"
 #include <memory>
 
 namespace mlirtrt::compiler {
@@ -101,12 +102,18 @@ class CompilerClient {
   /// Create or retrieve a cached PassManager of the given derived type using
   /// the provided options. PassManagers are cached by type and a hash of the
   /// string representation of the options.
+  /// This function should only be called if the options have a valid hash.
   template <typename CompilationTaskType, typename OptionsType>
   mlir::PassManager &getOrCreatePassManager(const OptionsType &options) {
-    auto key = std::make_pair(mlir::TypeID::get<CompilationTaskType>(),
-                              options.getHash());
+    auto hash = options.getHash();
+
+    assert(hash);
+
+    auto key =
+        std::make_pair(mlir::TypeID::get<CompilationTaskType>(), hash.value());
     auto it = cachedPassManagers.find(key);
-    if (it == cachedPassManagers.end() || options.shouldInvalidateCache()) {
+
+    if (it == cachedPassManagers.end()) {
       auto pm = std::make_unique<CompilationTaskType>(context, options);
       setupPassManagerLogging(*pm, options.debugOptions);
       auto *ptr = pm.get();
@@ -119,7 +126,6 @@ class CompilerClient {
   /// Return the MLIRContext associated with the client.
   mlir::MLIRContext *getContext() const { return context; }
 
-private:
   /// Helper for setting the correct logging options on cached PassManagers.
   static void setupPassManagerLogging(mlir::PassManager &pm,
                                       const DebugOptions &options);

mlir-tensorrt/compiler/include/mlir-tensorrt/Compiler/StableHloToExecutable.h

@@ -105,14 +105,7 @@ struct StableHLOToExecutableOptions : public mlir::OptionsContext {
   /// Get the mutable DebugOptions.
   DebugOptions &getDebugOptions() { return debugOptions; }
 
-  llvm::hash_code getHash() const override;
-
-  bool shouldInvalidateCache() const override {
-    // If a callback is provided, we have no way of verifying whether it is
-    // equivalent to a callback from another set of options. Therefore, we are
-    // forced to invalidate the cache entry if it is present at all.
-    return static_cast<bool>(layerMetadataCallback);
-  }
+  std::optional<llvm::hash_code> getHash() const override;
 
   /// The host index bit-width.
   int64_t executorIndexBitwidth{64};

mlir-tensorrt/compiler/lib/Compiler/StableHloToExecutable.cpp

@@ -40,13 +40,15 @@
 #include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Dialect/Bufferization/Transforms/Passes.h"
 #include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/Pass/Pass.h"
 #include "mlir/Pass/PassManager.h"
 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
 #include "mlir/Transforms/Passes.h"
 #include "stablehlo/dialect/StablehloOps.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/raw_ostream.h"
+#include <functional>
 #include <memory>
 
 #define DEBUG_TYPE "compiler-api"
@@ -271,11 +273,12 @@ Status StableHLOToExecutableOptions::inferDeviceOptionsFromHost() {
   return Status::getOk();
 }
 
-llvm::hash_code StableHLOToExecutableOptions::getHash() const {
-  llvm::hash_code hash = OptionsContext::getHash();
+std::optional<llvm::hash_code> StableHLOToExecutableOptions::getHash() const {
+  // If a callback is provided, we have no way of reliably hashing it.
   if (layerMetadataCallback)
-    return llvm::hash_combine(hash, &layerMetadataCallback);
-  return hash;
+    return std::nullopt;
+
+  return OptionsContext::getHash();
 }
 
 //===----------------------------------------------------------------------===//
@@ -473,11 +476,19 @@ StableHloToExecutableTask::compileStableHLOToExecutable(
   }
 #endif
 
-  mlir::PassManager &runner =
-      client.getOrCreatePassManager<StableHloToExecutableTask>(options);
+  mlir::PassManager *runner;
+
+  if (options.getHash())
+    runner = &client.getOrCreatePassManager<StableHloToExecutableTask>(options);
+  else {
+    auto pm = std::make_unique<StableHloToExecutableTask>(client.getContext(),
+                                                          options);
+    CompilerClient::setupPassManagerLogging(*pm, options.debugOptions);
+    runner = pm.get();
+  }
 
   // Setup pass manager
-  if (failed(runner.run(module)))
+  if (failed(runner->run(module)))
     return getInternalErrorStatus(
         "failed to run compilation on module with symbol name: {0}",
         module.getName() ? *module.getName() : "no-symbol-name");

mlir-tensorrt/tensorrt/include/mlir-tensorrt-dialect/Utils/Options.h

@@ -149,7 +149,7 @@ class OptionsContext : public llvm::cl::SubCommand {
   void print(llvm::raw_ostream &os) const;
 
   /// Get a hash derived from the string representation of the options.
-  virtual llvm::hash_code getHash() const;
+  virtual std::optional<llvm::hash_code> getHash() const;
 
   virtual bool shouldInvalidateCache() const { return false; }
 

mlir-tensorrt/tensorrt/lib/Utils/Options.cpp

@@ -63,7 +63,7 @@ void OptionsContext::print(llvm::raw_ostream &os) const {
       " ");
 }
 
-llvm::hash_code OptionsContext::getHash() const {
+std::optional<llvm::hash_code> OptionsContext::getHash() const {
   // We hash by just hashing the string representation.
   llvm::SmallString<128> str;
   {