Adds ExecutorOptions and DeviceOptions options providers

- Moves `executorIndexBitwidth` and `executorUsePackedMemRefCConv` into its
    own options bundle.

- Adds a `DeviceOptions` provider and updates the OptionsContext and OptionsProviders
    to use `llvm::Error`s instead of `mlirtrt::Status` since the latter is not accessible
    to the OptionsContext.

Author: pranavm-nvidia

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

@@ -28,7 +28,7 @@
 #define MLIR_TENSORRT_COMPILER_CLIENT
 
 #include "mlir-executor/Support/Status.h"
-#include "mlir-tensorrt/Compiler/Options.h"
+#include "mlir-tensorrt/Compiler/OptionsProviders.h"
 #include "mlir/IR/MLIRContext.h"
 #include "mlir/Pass/PassManager.h"
 #include "mlir/Support/TypeID.h"

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

@@ -0,0 +1,132 @@
+//===- OptionsProviders.h ---------------------------------------*- C++ -*-===//
+//
+// SPDX-FileCopyrightText: Copyright 2024 NVIDIA CORPORATION & AFFILIATES.
+// All rights reserved.
+// SPDX-License-Identifier: Apache-2.0
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+//
+//===----------------------------------------------------------------------===//
+///
+/// Data structures and functions for manipulating compiler options.
+///
+//===----------------------------------------------------------------------===//
+#ifndef MLIR_TENSORRT_COMPILER_OPTIONS
+#define MLIR_TENSORRT_COMPILER_OPTIONS
+
+#include "mlir-executor/Support/DeviceInfo.h"
+#include "mlir-tensorrt-dialect/Utils/Options.h"
+#include "mlir/Support/LLVM.h"
+#include "llvm/Support/CommandLine.h"
+#include "llvm/Support/Error.h"
+#include <string>
+
+namespace mlirtrt::compiler {
+
+// Use SFINAE to check whether the `finalizeImpl()` method is defined on a type.
+// If it is, the specialization (where the value is true) will be the better
+// match. Otherwise, we'll get the default value of false.
+template <typename, typename = void>
+constexpr bool has_finalize_impl_v = false;
+
+template <typename T>
+constexpr bool has_finalize_impl_v<
+    T, std::void_t<decltype(std::declval<T>().finalizeImpl())>> = true;
+
+// We use CRTP here so we can call `finalizeImpl()` if it's defined or provide
+// a default implementation otherwise.
+template <typename Derived>
+struct OptionsProvider {
+  /// Modifies options after parsing. This is required since we may need
+  /// to make changes to options based on the values of other options.
+  /// Do *not* override this method; instead, implement `finalizeImpl()`.
+  llvm::Error finalize() {
+    if constexpr (has_finalize_impl_v<Derived>)
+      return static_cast<Derived *>(this)->finalizeImpl();
+    else
+      return llvm::Error::success();
+  }
+};
+
+/// DebugOptions are options that are common to different compiler API
+/// interfaces.
+struct DebugOptions : public OptionsProvider<DebugOptions> {
+public:
+  /// A directory path where the IR will be dumped during compilation
+  /// using the `mlir-print-ir-tree-dir` mechanism.
+  std::string dumpIRPath = "";
+
+  /// Whether the LLVM 'debug' flag that enables execution of code guarded by
+  /// the `LLVM_DEBUG` macro should be set to 'on'. This results in very verbose
+  /// output from the compiler dumped to stderr.
+  bool enableLLVMDebugFlag = false;
+
+  /// A set of names to be given to the LLVM 'debug types' option, akin to
+  /// setting
+  /// `-debug-types=...` from the command line.
+  mlir::SmallVector<std::string> llvmDebugTypes = {};
+
+public:
+  void addToOptions(mlir::OptionsContext &context) {
+    context.addOption("mlir-print-ir-tree-dir", dumpIRPath, llvm::cl::init(""));
+    context.addOption("debug", enableLLVMDebugFlag);
+    context.addList<std::string>("debug-only", llvmDebugTypes,
+                                 llvm::cl::ZeroOrMore,
+                                 llvm::cl::CommaSeparated);
+  }
+};
+
+struct ExecutorOptions : public OptionsProvider<ExecutorOptions> {
+public:
+  /// The host index bit-width.
+  int64_t indexBitwidth{64};
+
+  /// Whether to pass memref's as struct/table in function calls.
+  bool usePackedMemRefCConv{true};
+
+public:
+  void addToOptions(mlir::OptionsContext &context) {
+    context.addOption("executor-index-bitwidth", indexBitwidth,
+                      llvm::cl::init(64));
+  }
+};
+
+struct DeviceOptions : public OptionsProvider<DeviceOptions> {
+public:
+  DeviceInfo info;
+
+  /// Whether to ignore `deviceX` options and instead infer them from the GPUs
+  /// on the host system running the compilation.
+  bool shouldInferFromHost = false;
+
+public:
+  void addToOptions(mlir::OptionsContext &context) {
+    context.addOption(
+        "device-compute-capability", info.computeCapability, llvm::cl::init(60),
+        llvm::cl::desc("Sets the device compute capbility. Only relevant "
+                       "if '--device-infer-from-host=false'"));
+    context.addOption("device-max-shared-memory-per-block-kb",
+                      info.maxSharedMemoryPerBlockKb, llvm::cl::init(48));
+    context.addOption("device-max-registers-per-block",
+                      info.maxRegistersPerBlock, llvm::cl::init(65536));
+    context.addOption("device-infer-from-host", shouldInferFromHost,
+                      llvm::cl::init(true),
+                      llvm::cl::desc("Infers device information from host"));
+  }
+
+  llvm::Error finalizeImpl();
+};
+
+} // namespace mlirtrt::compiler
+
+#endif // MLIR_TENSORRT_COMPILER_OPTIONS

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

@@ -32,6 +32,7 @@
 #include "mlir-tensorrt/Dialect/Plan/IR/Plan.h"
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/StringRef.h"
+#include "llvm/Support/Error.h"
 #include <functional>
 
 namespace mlirtrt::compiler {
@@ -71,14 +72,16 @@ optionsCreateFromArgs(const CompilerClient &client,
         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;
+  llvm::Error finalizeStatus = result->finalize();
+
+  std::optional<std::string> errMsg{};
+  llvm::handleAllErrors(
+      std::move(finalizeStatus),
+      [&errMsg](const llvm::StringError &err) { errMsg = err.getMessage(); });
+
+  if (errMsg)
+    return getInternalErrorStatus("failed to initialize options: %s",
+                                  errMsg->c_str());
 
   return std::unique_ptr<mlir::OptionsContext>(result.release());
 }

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

@@ -36,7 +36,7 @@
 #include "mlir-tensorrt-dialect/Utils/OptionsBundle.h"
 #include "mlir-tensorrt/Compiler/Client.h"
 #include "mlir-tensorrt/Compiler/Extension.h"
-#include "mlir-tensorrt/Compiler/Options.h"
+#include "mlir-tensorrt/Compiler/OptionsProviders.h"
 #include "mlir/IR/BuiltinOps.h"
 #include "mlir/Pass/PassManager.h"
 #include "mlir/Support/TypeID.h"
@@ -51,46 +51,16 @@ namespace mlirtrt::compiler {
 
 class StableHloToExecutableTask;
 
-struct StableHLOToExecutableOptions : public mlir::OptionsBundle<DebugOptions> {
+struct StableHLOToExecutableOptions
+    : public mlir::OptionsBundle<DebugOptions, ExecutorOptions, DeviceOptions> {
   /// Initializes the options. The extensions in the provided registry
   /// must be extensions for the StableHloToExecutable task.
   StableHLOToExecutableOptions(TaskExtensionRegistry extensions);
 
-  /// Set the target device compute capability (SM version) and max shared
-  /// memory per block (in kilobytes). The `maxSharedMemoryPerBlockKb` is the
-  /// maximum shared memory per block allowed for kernels and is passed to the
-  /// TensorRT builder.
-  StableHLOToExecutableOptions &
-  setDeviceOptions(int64_t computeCapability,
-                   int64_t maxSharedMemoryPerBlockKb);
-
-  /// Infer target device information from the first visible CUDA device on the
-  /// host executing this code.
-  Status inferDeviceOptionsFromHost();
-
   /// Return the hash of the options. Returns `nullopt` when the TensorRT
   /// layer metadata callback is set since that can't be reliably hashed.
   std::optional<llvm::hash_code> getHash() const override;
 
-  /// The host index bit-width.
-  int64_t executorIndexBitwidth{64};
-
-  /// Whether to pass memref's as struct/table in function calls.
-  bool executorUsePackedMemRefCConv{true};
-
-  /// Target device compute capability (SM version)
-  int64_t deviceComputeCapability;
-
-  /// Target device max shared memory per block (kilobytes)
-  int64_t deviceMaxSharedMemoryPerBlockKb;
-
-  /// Target device maximum 4-byte register sper block.
-  uint64_t deviceMaxRegistersPerBlock;
-
-  /// Whether to ignore `deviceX` options and instead infer them from the GPUs
-  /// on the host system running the compilation.
-  bool shouldInferDeviceOptionsFromHost = false;
-
   /// Whether to disallow host tensors in TensorRT clusters.
   bool disallowHostTensorsInTensorRTClusters = false;
 

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

@@ -169,10 +169,15 @@ MTRT_Status mtrtStableHloToExecutableOptionsCreate(
   auto result =
       std::make_unique<StableHLOToExecutableOptions>(std::move(extensions));
 
-  /// Populate device options from host information.
-  Status inferStatus = result->inferDeviceOptionsFromHost();
-  if (!inferStatus.isOk())
-    return wrap(inferStatus);
+  llvm::Error finalizeStatus = result->finalize();
+
+  std::optional<std::string> errMsg{};
+  llvm::handleAllErrors(
+      std::move(finalizeStatus),
+      [&errMsg](const llvm::StringError &err) { errMsg = err.getMessage(); });
+
+  if (errMsg)
+    return wrap(getInternalErrorStatus(errMsg->c_str()));
 
   *options = wrap(result.release());
   return mtrtStatusGetOk();
@@ -209,10 +214,15 @@ MTRT_Status mtrtStableHloToExecutableOptionsCreateFromArgs(
         "failed to parse options string {0} due to error: {1}", line, err));
   }
 
-  /// Populate device options from host information.
-  Status inferStatus = result->inferDeviceOptionsFromHost();
-  if (!inferStatus.isOk())
-    return wrap(inferStatus);
+  llvm::Error finalizeStatus = result->finalize();
+
+  std::optional<std::string> errMsg{};
+  llvm::handleAllErrors(
+      std::move(finalizeStatus),
+      [&errMsg](const llvm::StringError &err) { errMsg = err.getMessage(); });
+
+  if (errMsg)
+    return wrap(getInternalErrorStatus(errMsg->c_str()));
 
   *options = wrap(result.release());
   return mtrtStatusGetOk();

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

@@ -2,6 +2,7 @@ add_mlir_tensorrt_library(MLIRTensorRTCompilerClient
     Client.cpp
     Extension.cpp
     OptionsRegistry.cpp
+    OptionsProviders.cpp
     PARTIAL_SOURCES_INTENDED
 
     LINK_LIBS PUBLIC
@@ -11,6 +12,7 @@ add_mlir_tensorrt_library(MLIRTensorRTCompilerClient
     MLIRTensorRTOptionUtils
     MLIRTensorRTTargetTensorRT
     StablehloLinalgTransforms
+    MLIRTensorRTSupportDeviceInfo
 )
 
 add_mlir_tensorrt_library(MLIRTensorRTCompilerStableHloToExecutable

mlir-tensorrt/compiler/lib/Compiler/CMakeLists.txt

@@ -0,0 +1,38 @@
+//===- OptionsProviders.cpp -------------------------------------*- C++ -*-===//
+//
+// SPDX-FileCopyrightText: Copyright 2024 NVIDIA CORPORATION & AFFILIATES.
+// All rights reserved.
+// SPDX-License-Identifier: Apache-2.0
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+//
+//===----------------------------------------------------------------------===//
+///
+/// Data structures and functions for manipulating compiler options.
+///
+//===----------------------------------------------------------------------===//
+#include "mlir-tensorrt/Compiler/OptionsProviders.h"
+#include "mlir-executor/Support/DeviceInfo.h"
+#include "llvm/Support/Error.h"
+
+llvm::Error mlirtrt::compiler::DeviceOptions::finalizeImpl() {
+  if (shouldInferFromHost) {
+    StatusOr<DeviceInfo> deviceInfo = getDeviceInformationFromHost();
+
+    if (!deviceInfo.isOk())
+      return llvm::createStringError(deviceInfo.getString());
+
+    info = *deviceInfo;
+  }
+  return llvm::Error::success();
+}