Exposes base OptionsContext class in the API

In order to support more than just the StableHloToExecutable pipeline, we
need to be able to create different option types from the API. This commit
exposes the base `OptionsContext` class in the Python API and includes a mechanism
for child classes to register themselves with the client, allowing them to be
created through a common API.

Author: pranavm-nvidia

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

@@ -51,6 +51,29 @@ static inline bool mtrtCompilerClientIsNull(MTRT_CompilerClient options) {
   return !options.ptr;
 }
 
+//===----------------------------------------------------------------------===//
+// MTRT_OptionsContext
+//===----------------------------------------------------------------------===//
+
+typedef struct MTRT_OptionsContext {
+  void *ptr;
+} MTRT_OptionsContext;
+
+MLIR_CAPI_EXPORTED MTRT_Status mtrtOptionsContextCreateFromArgs(
+    MTRT_CompilerClient client, MTRT_OptionsContext *options,
+    MlirStringRef optionsType, const MlirStringRef *argv, unsigned argc);
+
+MLIR_CAPI_EXPORTED void mtrtOptionsContextPrint(MTRT_OptionsContext options,
+                                                MlirStringCallback append,
+                                                void *userData);
+
+MLIR_CAPI_EXPORTED MTRT_Status
+mtrtOptionsContextDestroy(MTRT_OptionsContext options);
+
+static inline bool mtrtOptionsConextIsNull(MTRT_OptionsContext options) {
+  return !options.ptr;
+}
+
 //===----------------------------------------------------------------------===//
 // MTRT_StableHLOToExecutableOptions
 //===----------------------------------------------------------------------===//

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

@@ -28,13 +28,17 @@
 #define MLIR_TENSORRT_COMPILER_CLIENT
 
 #include "mlir-executor/Support/Status.h"
+#include "mlir-tensorrt-dialect/Utils/Options.h"
 #include "mlir-tensorrt/Compiler/Options.h"
+#include "mlir-tensorrt/Dialect/Plan/IR/Plan.h"
 #include "mlir/IR/MLIRContext.h"
 #include "mlir/Pass/PassManager.h"
 #include "mlir/Support/TypeID.h"
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/Hashing.h"
+#include "llvm/ADT/StringRef.h"
 #include "llvm/Support/ErrorHandling.h"
+#include <functional>
 #include <memory>
 
 namespace mlirtrt::compiler {
@@ -85,6 +89,19 @@ class CompilationTask : public CompilationTaskBase {
 // CompilerClient
 //===----------------------------------------------------------------------===//
 
+class CompilerClient;
+
+using OptionsConstructorFuncT =
+    std::function<StatusOr<std::unique_ptr<mlir::OptionsContext>>(
+        const CompilerClient &client, const llvm::ArrayRef<llvm::StringRef>)>;
+
+void registerOption(const llvm::StringRef optionsType,
+                    OptionsConstructorFuncT func);
+
+StatusOr<std::unique_ptr<mlir::OptionsContext>>
+createOptions(const CompilerClient &client, const llvm::StringRef optionsType,
+              const llvm::ArrayRef<llvm::StringRef> args);
+
 /// C++ users of the MLIR-TensorRT Compiler API should create a CompilerClient
 /// once for each process or thread that will be performing concurrent
 /// compilation work. The CompilerClient holds long-lived resources such as the
@@ -147,6 +164,39 @@ class CompilerClient {
       cachedPassManagers;
 };
 
+/// Helper to register option types with the client
+template <typename OptionsT, typename TaskT>
+StatusOr<std::unique_ptr<mlir::OptionsContext>>
+optionsCreateFromArgs(const CompilerClient &client,
+                      const llvm::ArrayRef<llvm::StringRef> args) {
+  // Load available extensions.
+  mlir::MLIRContext *context = client.getContext();
+  mlir::plan::PlanDialect *planDialect =
+      context->getLoadedDialect<mlir::plan::PlanDialect>();
+  compiler::TaskExtensionRegistry extensions =
+      planDialect->extensionConstructors.getExtensionRegistryForTask<TaskT>();
+
+  auto result = std::make_unique<OptionsT>(std::move(extensions));
+
+  std::string err;
+  if (failed(result->parse(args, err))) {
+    return getInternalErrorStatus(
+        "failed to parse options string \"{0:$[ ]}\" due to error {1}",
+        llvm::iterator_range(args), err);
+  }
+
+  // TODO: Figure out whether to add a method in the base class like
+  // "finalizeOptions" or a callback here, or something else if
+  // `inferDeviceOptionsFromHost` is unique to StableHLO.
+  //
+  // Populate device options from host information.
+  Status inferStatus = result->inferDeviceOptionsFromHost();
+  if (!inferStatus.isOk())
+    return inferStatus;
+
+  return std::unique_ptr<mlir::OptionsContext>(result.release());
+}
+
 } // namespace mlirtrt::compiler
 
 #endif // MLIR_TENSORRT_COMPILER_CLIENT

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

@@ -189,6 +189,9 @@ class StableHloToExecutableTask
                                const StableHLOToExecutableOptions &options);
 };
 
+/// Register the task/options with the client's registry.
+void registerStableHloToExecutableTask();
+
 //===----------------------------------------------------------------------===//
 // Pipeline Registrations
 //===----------------------------------------------------------------------===//

mlir-tensorrt/compiler/include/mlir-tensorrt/Registration/RegisterMlirTensorRtPasses.h

@@ -53,6 +53,7 @@ inline void registerAllMlirTensorRtPasses() {
   mlir::registerConvertPDLToPDLInterp();
 
 #ifdef MLIR_TRT_ENABLE_HLO
+  mlirtrt::compiler::registerStableHloToExecutableTask();
   mlirtrt::compiler::registerStablehloClusteringPipelines();
   registerStableHloInputPipelines();
   stablehlo_ext::registerStableHloExtPasses();

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

@@ -26,11 +26,13 @@
 #include "mlir-c/Support.h"
 #include "mlir-executor-c/Support/Status.h"
 #include "mlir-tensorrt-dialect/Target/TranslateToTensorRT.h"
+#include "mlir-tensorrt-dialect/Utils/Options.h"
 #include "mlir-tensorrt/Compiler/Extension.h"
 #include "mlir-tensorrt/Compiler/StableHloToExecutable.h"
 #include "mlir-tensorrt/Compiler/TensorRTExtension/TensorRTExtension.h"
 #include "mlir-tensorrt/Dialect/Plan/IR/Plan.h"
 #include "mlir/CAPI/IR.h"
+#include "mlir/CAPI/Utils.h"
 #include "llvm/ADT/StringExtras.h"
 
 using namespace mlirtrt;
@@ -44,6 +46,7 @@ using namespace mlir;
 DEFINE_C_API_PTR_METHODS(MTRT_CompilerClient, CompilerClient)
 DEFINE_C_API_PTR_METHODS(MTRT_StableHLOToExecutableOptions,
                          StableHLOToExecutableOptions)
+DEFINE_C_API_PTR_METHODS(MTRT_OptionsContext, OptionsContext)
 #if defined(__GNUC__) || defined(__clang__)
 #pragma GCC diagnostic pop
 #endif
@@ -99,6 +102,40 @@ MTRT_Status mtrtCompilerClientDestroy(MTRT_CompilerClient client) {
   return mtrtStatusGetOk();
 }
 
+//===----------------------------------------------------------------------===//
+// MTRT_OptionsContext
+//===----------------------------------------------------------------------===//
+
+MLIR_CAPI_EXPORTED MTRT_Status mtrtOptionsContextCreateFromArgs(
+    MTRT_CompilerClient client, MTRT_OptionsContext *options,
+    MlirStringRef optionsType, const MlirStringRef *argv, unsigned argc) {
+  std::vector<llvm::StringRef> argvStrRef(argc);
+  for (unsigned i = 0; i < argc; i++)
+    argvStrRef[i] = llvm::StringRef(argv[i].data, argv[i].length);
+
+  auto result = createOptions(
+      *unwrap(client), llvm::StringRef(optionsType.data, optionsType.length),
+      argvStrRef);
+  if (!result.isOk())
+    return wrap(result.getStatus());
+
+  *options = wrap(result->release());
+  return mtrtStatusGetOk();
+}
+
+MLIR_CAPI_EXPORTED void mtrtOptionsContextPrint(MTRT_OptionsContext options,
+                                                MlirStringCallback append,
+                                                void *userData) {
+  mlir::detail::CallbackOstream stream(append, userData);
+  unwrap(options)->print(stream);
+}
+
+MLIR_CAPI_EXPORTED MTRT_Status
+mtrtOptionsContextDestroy(MTRT_OptionsContext options) {
+  delete unwrap(options);
+  return mtrtStatusGetOk();
+}
+
 //===----------------------------------------------------------------------===//
 // MTRT_StableHLOToExecutableOptions
 //===----------------------------------------------------------------------===//

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

@@ -24,6 +24,9 @@
 #include "mlir-tensorrt/Compiler/Client.h"
 #include "mlir/IR/BuiltinOps.h"
 #include "mlir/Support/FileUtilities.h"
+#include "llvm/ADT/ArrayRef.h"
+#include "llvm/ADT/StringRef.h"
+#include "llvm/Support/ManagedStatic.h"
 
 using namespace mlirtrt;
 using namespace mlirtrt::compiler;
@@ -46,6 +49,25 @@ CompilationTaskBase::~CompilationTaskBase() {}
 // CompilerClient
 //===----------------------------------------------------------------------===//
 
+namespace mlirtrt::compiler {
+static llvm::ManagedStatic<llvm::StringMap<OptionsConstructorFuncT>> registry{};
+
+void registerOption(const llvm::StringRef optionsType,
+                    OptionsConstructorFuncT func) {
+  (*registry)[optionsType] = std::move(func);
+}
+
+StatusOr<std::unique_ptr<mlir::OptionsContext>>
+createOptions(const CompilerClient &client, const llvm::StringRef optionsType,
+              const llvm::ArrayRef<llvm::StringRef> args) {
+  if (!registry->contains(optionsType))
+    return getInvalidArgStatus(
+        "{0} is not a valid option type. Valid options were: {1:$[ ]}",
+        optionsType, llvm::iterator_range(registry->keys()));
+  return (*registry)[optionsType](client, args);
+}
+} // namespace mlirtrt::compiler
+
 StatusOr<std::unique_ptr<CompilerClient>>
 CompilerClient::create(MLIRContext *context) {
   context->disableMultithreading();

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

@@ -501,6 +501,12 @@ static StableHLOToExecutableOptions populateStablehloClusteringPipelineOpts(
   return opts;
 }
 
+void mlirtrt::compiler::registerStableHloToExecutableTask() {
+  registerOption("stable-hlo-to-executable",
+                 optionsCreateFromArgs<StableHLOToExecutableOptions,
+                                       StableHloToExecutableTask>);
+}
+
 void mlirtrt::compiler::registerStablehloClusteringPipelines() {
   PassRegistration<HloToStdPass>();
   PassRegistration<HloToArithDynamicPipelinePass>();

mlir-tensorrt/python/bindings/Compiler/CompilerPyBind.cpp

@@ -17,6 +17,7 @@
 #include "mlir-tensorrt-c/Compiler/Compiler.h"
 #include "mlir/Bindings/Python/PybindAdaptors.h"
 #include "pybind11/pybind11.h"
+#include "llvm/ADT/StringRef.h"
 #include "llvm/Support/DynamicLibrary.h"
 #include "llvm/Support/raw_ostream.h"
 #include <pybind11/attr.h>
@@ -52,6 +53,17 @@ class PyCompilerClient
       mtrtCompilerClientIsNull, mtrtCompilerClientDestroy};
 };
 
+/// Python object type wrapper for `MTRT_OptionsContext`.
+class PyOptionsContext
+    : public PyMTRTWrapper<PyOptionsContext, MTRT_OptionsContext> {
+public:
+  using PyMTRTWrapper::PyMTRTWrapper;
+  DECLARE_WRAPPER_CONSTRUCTORS(PyOptionsContext);
+
+  static constexpr auto kMethodTable = CAPITable<MTRT_OptionsContext>{
+      mtrtOptionsConextIsNull, mtrtOptionsContextDestroy};
+};
+
 /// Python object type wrapper for `MTRT_StableHLOToExecutableOptions`.
 class PyStableHLOToExecutableOptions
     : public PyMTRTWrapper<PyStableHLOToExecutableOptions,
@@ -240,6 +252,36 @@ PYBIND11_MODULE(_api, m) {
         return new PyCompilerClient(client);
       }));
 
+  py::class_<PyOptionsContext>(m, "OptionsContext", py::module_local())
+      .def(py::init<>([](PyCompilerClient &client,
+                         const std::string &optionsType,
+                         const std::vector<std::string> &args) {
+             std::vector<MlirStringRef> refs(args.size());
+             for (unsigned i = 0; i < args.size(); i++)
+               refs[i] = mlirStringRefCreate(args[i].data(), args[i].size());
+
+             MTRT_OptionsContext options;
+             MTRT_Status s = mtrtOptionsContextCreateFromArgs(
+                 client, &options,
+                 mlirStringRefCreate(optionsType.data(), optionsType.size()),
+                 refs.data(), refs.size());
+             THROW_IF_MTRT_ERROR(s);
+             return new PyOptionsContext(options);
+           }),
+           py::arg("client"), py::arg("options_type"), py::arg("args"))
+
+      .def("__repr__", [](PyOptionsContext &self) {
+        auto callback = [](MlirStringRef data, void *initialString) {
+          *reinterpret_cast<std::string *>(initialString) +=
+              llvm::StringRef(data.data, data.length);
+        };
+
+        std::string result("Options[");
+        mtrtOptionsContextPrint(self, callback, &result);
+        result += "]";
+        return result;
+      });
+
   py::class_<PyStableHLOToExecutableOptions>(m, "StableHLOToExecutableOptions",
                                              py::module_local())
       .def(py::init<>([](PyCompilerClient &client,

mlir-tensorrt/python/bindings/Utils.h

@@ -163,7 +163,7 @@ class PyMTRTWrapper {
 
   static py::object createFromCapsule(py::object capsule) {
     if constexpr (cFuncTable.capsuleToCApi == nullptr) {
-      throw py::value_error("boject cannot be converted from opaque capsule");
+      throw py::value_error("object cannot be converted from opaque capsule");
     } else {
       MTRT_StableHLOToExecutableOptions cObj =
           cFuncTable.capsuleToCApi(capsule.ptr());

mlir-tensorrt/test/python/mlir_tensorrt_compiler/compiler_api/test_options_context.py

@@ -0,0 +1,30 @@
+# REQUIRES: host-has-at-least-1-gpus
+# RUN: %PYTHON %s 2>&1 | FileCheck %s
+
+import mlir_tensorrt.compiler.api as api
+from mlir_tensorrt.compiler.ir import *
+
+
+with Context() as context:
+    client = api.CompilerClient(context)
+    # Try to create a non-existent option type
+    try:
+        opts = api.OptionsContext(client, "non-existent-options-type", [])
+    except Exception as err:
+        print(err)
+
+    opts = api.OptionsContext(
+        client,
+        "stable-hlo-to-executable",
+        [
+            "--tensorrt-builder-opt-level=3",
+            "--tensorrt-strongly-typed=false",
+            "--tensorrt-workspace-memory-pool-limit=1gb",
+        ],
+    )
+
+    print(opts)
+
+
+# CHECK: InvalidArgument: InvalidArgument: non-existent-options-type is not a valid option type. Valid options were: stable-hlo-to-executable
+# CHECK: --tensorrt-timing-cache-path= --device-infer-from-host=true --debug-only= --executor-index-bitwidth=64 --entrypoint=main --plan-clustering-disallow-host-tensors-in-tensorrt-clusters=false --tensorrt-workspace-memory-pool-limit=1073741824 --device-max-registers-per-block=65536 --tensorrt-strongly-typed=false --tensorrt-layer-info-dir= --device-compute-capability=86 --debug=false --mlir-print-ir-tree-dir= --disable-tensorrt-extension=false --tensorrt-builder-opt-level=3 --tensorrt-engines-dir=