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,8 +33,32 @@
 #include "mlir/Pass/PassManager.h"
 #include "mlir/Support/TypeID.h"
 #include "llvm/ADT/DenseMap.h"
+#include "llvm/ADT/DenseMapInfo.h"
+#include "llvm/ADT/Hashing.h"
 #include <memory>
 
+namespace llvm {
+
+template <>
+struct DenseMapInfo<std::optional<llvm::hash_code>> {
+  static inline std::optional<llvm::hash_code> getEmptyKey() { return {}; }
+  static inline std::optional<llvm::hash_code> getTombstoneKey() {
+    // TODO: how to provide a tombstone key?
+    return {};
+  }
+  static unsigned getHashValue(const std::optional<llvm::hash_code> &val) {
+    if (!val)
+      return 0U;
+    return hash_value(*val);
+  }
+  static bool isEqual(const std::optional<llvm::hash_code> &LHS,
+                      const std::optional<llvm::hash_code> &RHS) {
+    return LHS == RHS;
+  }
+};
+
+} // namespace llvm
+
 namespace mlirtrt::compiler {
 
 //===----------------------------------------------------------------------===//
@@ -103,10 +127,11 @@ class CompilerClient {
   /// string representation of the options.
   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();
+    auto key = std::make_pair(mlir::TypeID::get<CompilationTaskType>(), hash);
     auto it = cachedPassManagers.find(key);
-    if (it == cachedPassManagers.end() || options.shouldInvalidateCache()) {
+
+    if (!hash || it == cachedPassManagers.end()) {
       auto pm = std::make_unique<CompilationTaskType>(context, options);
       setupPassManagerLogging(*pm, options.debugOptions);
       auto *ptr = pm.get();
@@ -131,7 +156,8 @@ class CompilerClient {
   mlir::MLIRContext *context;
 
   /// Key pair of [PassManager Kind, Options Hash].
-  using PassManagerKey = std::pair<mlir::TypeID, llvm::hash_code>;
+  using PassManagerKey =
+      std::pair<mlir::TypeID, std::optional<llvm::hash_code>>;
 
   /// A registry of pass managers for specific kinds of tasks. The map is
   /// indexed by the TypeID of the PassManager kind and the hash of the 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

@@ -271,11 +271,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();
 }
 
 //===----------------------------------------------------------------------===//

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;
   {