From 47c11d52cd36208c439d0e419f61e28970bc22ea Mon Sep 17 00:00:00 2001 From: pranavm Date: Tue, 3 Dec 2024 13:11:26 -0800 Subject: [PATCH] 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. --- .../include/mlir-tensorrt/Compiler/Client.h | 2 +- .../include/mlir-tensorrt/Compiler/Options.h | 62 -------- .../mlir-tensorrt/Compiler/OptionsProviders.h | 132 ++++++++++++++++++ .../mlir-tensorrt/Compiler/OptionsRegistry.h | 19 +-- .../Compiler/StableHloToExecutable.h | 36 +---- .../compiler/lib/CAPI/Compiler/Compiler.cpp | 26 ++-- .../compiler/lib/Compiler/CMakeLists.txt | 2 + .../lib/Compiler/OptionsProviders.cpp | 38 +++++ .../lib/Compiler/StableHloToExecutable.cpp | 66 ++------- .../TensorRTExtension/TensorRTExtension.cpp | 5 +- .../executor/lib/Support/DeviceInfo.cpp | 18 ++- .../mlir-tensorrt-dialect/Utils/Options.h | 6 + .../Utils/OptionsBundle.h | 14 ++ .../compiler_api/test_options_context.py | 3 +- 14 files changed, 253 insertions(+), 176 deletions(-) delete mode 100644 mlir-tensorrt/compiler/include/mlir-tensorrt/Compiler/Options.h create mode 100644 mlir-tensorrt/compiler/include/mlir-tensorrt/Compiler/OptionsProviders.h create mode 100644 mlir-tensorrt/compiler/lib/Compiler/OptionsProviders.cpp diff --git a/mlir-tensorrt/compiler/include/mlir-tensorrt/Compiler/Client.h b/mlir-tensorrt/compiler/include/mlir-tensorrt/Compiler/Client.h index 253c34759..6824ffad6 100644 --- a/mlir-tensorrt/compiler/include/mlir-tensorrt/Compiler/Client.h +++ b/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" diff --git a/mlir-tensorrt/compiler/include/mlir-tensorrt/Compiler/Options.h b/mlir-tensorrt/compiler/include/mlir-tensorrt/Compiler/Options.h deleted file mode 100644 index cd9c5e393..000000000 --- a/mlir-tensorrt/compiler/include/mlir-tensorrt/Compiler/Options.h +++ /dev/null @@ -1,62 +0,0 @@ -//===- Options.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-tensorrt-dialect/Utils/Options.h" -#include "mlir/Support/LLVM.h" -#include "llvm/Support/CommandLine.h" -#include - -namespace mlirtrt::compiler { - -/// DebugOptions are options that are common to different compiler API -/// interfaces. -struct DebugOptions { - /// 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 llvmDebugTypes = {}; - - void addToOptions(mlir::OptionsContext &context) { - context.addOption("mlir-print-ir-tree-dir", dumpIRPath, llvm::cl::init("")); - context.addOption("debug", enableLLVMDebugFlag); - context.addList("debug-only", llvmDebugTypes, - llvm::cl::ZeroOrMore, - llvm::cl::CommaSeparated); - } -}; - -} // namespace mlirtrt::compiler - -#endif // MLIR_TENSORRT_COMPILER_OPTIONS diff --git a/mlir-tensorrt/compiler/include/mlir-tensorrt/Compiler/OptionsProviders.h b/mlir-tensorrt/compiler/include/mlir-tensorrt/Compiler/OptionsProviders.h new file mode 100644 index 000000000..e3c0d7626 --- /dev/null +++ b/mlir-tensorrt/compiler/include/mlir-tensorrt/Compiler/OptionsProviders.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 + +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 +constexpr bool has_finalize_impl_v = false; + +template +constexpr bool has_finalize_impl_v< + T, std::void_t().finalizeImpl())>> = true; + +// We use CRTP here so we can call `finalizeImpl()` if it's defined or provide +// a default implementation otherwise. +template +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) + return static_cast(this)->finalizeImpl(); + else + return llvm::Error::success(); + } +}; + +/// DebugOptions are options that are common to different compiler API +/// interfaces. +struct DebugOptions : public OptionsProvider { +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 llvmDebugTypes = {}; + +public: + void addToOptions(mlir::OptionsContext &context) { + context.addOption("mlir-print-ir-tree-dir", dumpIRPath, llvm::cl::init("")); + context.addOption("debug", enableLLVMDebugFlag); + context.addList("debug-only", llvmDebugTypes, + llvm::cl::ZeroOrMore, + llvm::cl::CommaSeparated); + } +}; + +struct ExecutorOptions : public OptionsProvider { +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 { +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 diff --git a/mlir-tensorrt/compiler/include/mlir-tensorrt/Compiler/OptionsRegistry.h b/mlir-tensorrt/compiler/include/mlir-tensorrt/Compiler/OptionsRegistry.h index d24a12fd6..8fc387eeb 100644 --- a/mlir-tensorrt/compiler/include/mlir-tensorrt/Compiler/OptionsRegistry.h +++ b/mlir-tensorrt/compiler/include/mlir-tensorrt/Compiler/OptionsRegistry.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 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 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(result.release()); } diff --git a/mlir-tensorrt/compiler/include/mlir-tensorrt/Compiler/StableHloToExecutable.h b/mlir-tensorrt/compiler/include/mlir-tensorrt/Compiler/StableHloToExecutable.h index 39357f11a..946a8d6f2 100644 --- a/mlir-tensorrt/compiler/include/mlir-tensorrt/Compiler/StableHloToExecutable.h +++ b/mlir-tensorrt/compiler/include/mlir-tensorrt/Compiler/StableHloToExecutable.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 { +struct StableHLOToExecutableOptions + : public mlir::OptionsBundle { /// 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 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; diff --git a/mlir-tensorrt/compiler/lib/CAPI/Compiler/Compiler.cpp b/mlir-tensorrt/compiler/lib/CAPI/Compiler/Compiler.cpp index 90920b61a..082e742fe 100644 --- a/mlir-tensorrt/compiler/lib/CAPI/Compiler/Compiler.cpp +++ b/mlir-tensorrt/compiler/lib/CAPI/Compiler/Compiler.cpp @@ -169,10 +169,15 @@ MTRT_Status mtrtStableHloToExecutableOptionsCreate( auto result = std::make_unique(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 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 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(); diff --git a/mlir-tensorrt/compiler/lib/Compiler/CMakeLists.txt b/mlir-tensorrt/compiler/lib/Compiler/CMakeLists.txt index c26cd8ec1..746bd2e81 100644 --- a/mlir-tensorrt/compiler/lib/Compiler/CMakeLists.txt +++ b/mlir-tensorrt/compiler/lib/Compiler/CMakeLists.txt @@ -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 diff --git a/mlir-tensorrt/compiler/lib/Compiler/OptionsProviders.cpp b/mlir-tensorrt/compiler/lib/Compiler/OptionsProviders.cpp new file mode 100644 index 000000000..e86c12088 --- /dev/null +++ b/mlir-tensorrt/compiler/lib/Compiler/OptionsProviders.cpp @@ -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 = getDeviceInformationFromHost(); + + if (!deviceInfo.isOk()) + return llvm::createStringError(deviceInfo.getString()); + + info = *deviceInfo; + } + return llvm::Error::success(); +} diff --git a/mlir-tensorrt/compiler/lib/Compiler/StableHloToExecutable.cpp b/mlir-tensorrt/compiler/lib/Compiler/StableHloToExecutable.cpp index a585be43a..0443208b5 100644 --- a/mlir-tensorrt/compiler/lib/Compiler/StableHloToExecutable.cpp +++ b/mlir-tensorrt/compiler/lib/Compiler/StableHloToExecutable.cpp @@ -30,7 +30,7 @@ #include "mlir-tensorrt-dialect/Target/TranslateToTensorRT.h" #include "mlir-tensorrt-dialect/TensorRT/Transforms/Passes.h" #include "mlir-tensorrt/Compiler/Extension.h" -#include "mlir-tensorrt/Compiler/Options.h" +#include "mlir-tensorrt/Compiler/OptionsProviders.h" #include "mlir-tensorrt/Compiler/OptionsRegistry.h" #include "mlir-tensorrt/Compiler/TensorRTExtension/TensorRTExtension.h" #include "mlir-tensorrt/Conversion/Passes.h" @@ -162,58 +162,11 @@ StableHLOToExecutableOptions::StableHLOToExecutableOptions( disallowHostTensorsInTensorRTClusters, llvm::cl::init(false), llvm::cl::desc("Don't allow TensorRt clusters to contain host tensor " "calculations (but they can still be inputs)")); - addOption("executor-index-bitwidth", executorIndexBitwidth, - llvm::cl::init(64)); - addOption("device-compute-capability", deviceComputeCapability, - llvm::cl::init(64), - llvm::cl::desc("Sets the device compute capbility. Only relevant " - "if '--device-infer-from-host=false'")); - addOption("device-max-shared-memory-per-block-kb", - deviceMaxSharedMemoryPerBlockKb, llvm::cl::init(0)); - addOption("device-max-registers-per-block", deviceMaxRegistersPerBlock, - llvm::cl::init(0)); - addOption("device-infer-from-host", shouldInferDeviceOptionsFromHost, - llvm::cl::init(true), - llvm::cl::desc("Infers device information from host")); + addOption("entrypoint", entrypoint, llvm::cl::init("main"), llvm::cl::desc("entrypoint function name")); } -StableHLOToExecutableOptions &StableHLOToExecutableOptions::setDeviceOptions( - int64_t computeCapability, int64_t maxSharedMemoryPerBlockKb) { - deviceMaxSharedMemoryPerBlockKb = maxSharedMemoryPerBlockKb; - deviceComputeCapability = computeCapability; - return *this; -} - -Status StableHLOToExecutableOptions::inferDeviceOptionsFromHost() { - cudaDeviceProp properties; - cudaError_t err = cudaGetDeviceProperties(&properties, 0); - if (err != cudaSuccess) - return getStatusWithMsg(StatusCode::InternalError, - "failed to get cuda device properties"); - - int ccMajor = 0; - int ccMinor = 0; - err = cudaDeviceGetAttribute( - &ccMajor, cudaDeviceAttr::cudaDevAttrComputeCapabilityMajor, 0); - if (err != cudaSuccess) - return getStatusWithMsg(StatusCode::InternalError, - "failed to get cuda device compute capability"); - err = cudaDeviceGetAttribute( - &ccMinor, cudaDeviceAttr::cudaDevAttrComputeCapabilityMinor, 0); - if (err != cudaSuccess) - return getStatusWithMsg(StatusCode::InternalError, - "failed to get cuda device compute capability"); - - // We want SM version as a single number. - int64_t smVersion = ccMajor * 10 + ccMinor; - this->deviceComputeCapability = smVersion; - this->deviceMaxSharedMemoryPerBlockKb = properties.sharedMemPerBlock / 1024; - this->deviceMaxRegistersPerBlock = properties.regsPerBlock; - return Status::getOk(); -} - std::optional StableHLOToExecutableOptions::getHash() const { // If a callback is provided, we have no way of reliably hashing it. if (layerMetadataCallback) @@ -298,8 +251,9 @@ void StableHloToExecutableTask::buildPostClusteringPipeline( populateExtensionPasses(pm, opts, Phase::ExecutorLowering); ConvertCUDAToExecutorPassOptions cudaToExecutorOpts; - cudaToExecutorOpts.indexBitwidth = opts.executorIndexBitwidth; - cudaToExecutorOpts.usePackedMemRefCConv = opts.executorUsePackedMemRefCConv; + cudaToExecutorOpts.indexBitwidth = opts.get().indexBitwidth; + cudaToExecutorOpts.usePackedMemRefCConv = + opts.get().usePackedMemRefCConv; pm.addPass(createConvertCUDAToExecutorPass(cudaToExecutorOpts)); pm.addPass(createDropNestedModulesPass()); @@ -328,7 +282,7 @@ void StableHloToExecutableTask::populatePassManager( buildPostClusteringPipeline(pm, options); mlir::executor::ConvertStdToExecutorPassOptions stdToExecOpts; - stdToExecOpts.indexBitwidth = options.executorIndexBitwidth; + stdToExecOpts.indexBitwidth = options.get().indexBitwidth; stdToExecOpts.usePackedMemRefCConv = true; mlir::executor::buildExecutorLoweringPipeline(pm, stdToExecOpts); } @@ -513,10 +467,12 @@ static StableHLOToExecutableOptions populateStablehloClusteringPipelineOpts( extensions.getOrCreateExtension(); StableHLOToExecutableOptions opts(std::move(extensions)); - opts.deviceComputeCapability = cliOpts.deviceComputeCapability; - opts.deviceMaxSharedMemoryPerBlockKb = + opts.get().info.computeCapability = + cliOpts.deviceComputeCapability; + opts.get().info.maxSharedMemoryPerBlockKb = cliOpts.deviceMaxSharedMemoryPerBlockKb; - opts.shouldInferDeviceOptionsFromHost = cliOpts.inferDeviceOptionsFromHost; + opts.get().shouldInferFromHost = + cliOpts.inferDeviceOptionsFromHost; opts.entrypoint = cliOpts.entrypoint; return opts; } diff --git a/mlir-tensorrt/compiler/lib/Compiler/TensorRTExtension/TensorRTExtension.cpp b/mlir-tensorrt/compiler/lib/Compiler/TensorRTExtension/TensorRTExtension.cpp index f79d0a0bb..b9e8877b3 100644 --- a/mlir-tensorrt/compiler/lib/Compiler/TensorRTExtension/TensorRTExtension.cpp +++ b/mlir-tensorrt/compiler/lib/Compiler/TensorRTExtension/TensorRTExtension.cpp @@ -75,8 +75,9 @@ void StableHLOToExecutableTensorRTExtension::populatePasses( if (phase == Phase::ExecutorLowering) { ConvertTensorRTRuntimeToExecutorPassOptions toExecutorOpts; - toExecutorOpts.indexBitwidth = options.executorIndexBitwidth; - toExecutorOpts.usePackedMemRefCConv = options.executorUsePackedMemRefCConv; + toExecutorOpts.indexBitwidth = options.get().indexBitwidth; + toExecutorOpts.usePackedMemRefCConv = + options.get().usePackedMemRefCConv; pm.addPass(createConvertTensorRTRuntimeToExecutorPass(toExecutorOpts)); return; } diff --git a/mlir-tensorrt/executor/lib/Support/DeviceInfo.cpp b/mlir-tensorrt/executor/lib/Support/DeviceInfo.cpp index 1cc75d198..9e9d4146f 100644 --- a/mlir-tensorrt/executor/lib/Support/DeviceInfo.cpp +++ b/mlir-tensorrt/executor/lib/Support/DeviceInfo.cpp @@ -23,26 +23,32 @@ using namespace mlirtrt; +static Status makeCudaStringError(cudaError_t errCode, + llvm::StringRef context) { + // Create a detailed error message using llvm::createStringError + return getInternalErrorStatus("{0}: {1}", context, + cudaGetErrorString(errCode)); +} + StatusOr mlirtrt::getDeviceInformationFromHost() { #ifdef MLIR_EXECUTOR_ENABLE_CUDA cudaDeviceProp properties; cudaError_t err = cudaGetDeviceProperties(&properties, 0); if (err != cudaSuccess) - return getStatusWithMsg(StatusCode::InternalError, - "failed to get cuda device properties"); + return makeCudaStringError(err, "failed to get cuda device properties"); int ccMajor = 0; int ccMinor = 0; err = cudaDeviceGetAttribute( &ccMajor, cudaDeviceAttr::cudaDevAttrComputeCapabilityMajor, 0); if (err != cudaSuccess) - return getStatusWithMsg(StatusCode::InternalError, - "failed to get cuda device compute capability"); + return makeCudaStringError(err, + "failed to get cuda device compute capability"); err = cudaDeviceGetAttribute( &ccMinor, cudaDeviceAttr::cudaDevAttrComputeCapabilityMinor, 0); if (err != cudaSuccess) - return getStatusWithMsg(StatusCode::InternalError, - "failed to get cuda device compute capability"); + return makeCudaStringError(err, + "failed to get cuda device compute capability"); // We want SM version as a single number. int64_t smVersion = ccMajor * 10 + ccMinor; diff --git a/mlir-tensorrt/tensorrt/include/mlir-tensorrt-dialect/Utils/Options.h b/mlir-tensorrt/tensorrt/include/mlir-tensorrt-dialect/Utils/Options.h index d6140a858..6c4e4b2a7 100644 --- a/mlir-tensorrt/tensorrt/include/mlir-tensorrt-dialect/Utils/Options.h +++ b/mlir-tensorrt/tensorrt/include/mlir-tensorrt-dialect/Utils/Options.h @@ -155,6 +155,12 @@ class OptionsContext : public llvm::cl::SubCommand { /// as a cache key. virtual std::optional getHash() const; + /// Update options, if needed, after parsing. This method can be used to + /// modify options based on the values of other options or can be used to + /// populate options that were not provided using arbitrarily complex logic + /// (instead of just a default value). + virtual llvm::Error finalize() = 0; + private: struct OptionInfo { std::unique_ptr option; diff --git a/mlir-tensorrt/tensorrt/include/mlir-tensorrt-dialect/Utils/OptionsBundle.h b/mlir-tensorrt/tensorrt/include/mlir-tensorrt-dialect/Utils/OptionsBundle.h index c6597c259..959ad0d50 100644 --- a/mlir-tensorrt/tensorrt/include/mlir-tensorrt-dialect/Utils/OptionsBundle.h +++ b/mlir-tensorrt/tensorrt/include/mlir-tensorrt-dialect/Utils/OptionsBundle.h @@ -21,6 +21,7 @@ #define MLIR_TENSORRT_DIALECT_UTILS_OPTIONS_BUNDLE #include "mlir-tensorrt-dialect/Utils/Options.h" +#include "llvm/Support/Error.h" #include namespace mlir { @@ -48,6 +49,19 @@ class OptionsBundle : public OptionsContext { return std::get(optionProviders); } + llvm::Error finalize() override { + llvm::Error result = llvm::Error::success(); + std::apply( + [&](auto &...optionProvider) { + ((result = std::move(llvm::joinErrors(std::move(result), + optionProvider.finalize()))), + ...); + }, + optionProviders); + + return result; + } + private: std::tuple optionProviders{}; }; diff --git a/mlir-tensorrt/test/python/mlir_tensorrt_compiler/compiler_api/test_options_context.py b/mlir-tensorrt/test/python/mlir_tensorrt_compiler/compiler_api/test_options_context.py index 95b9c591a..fdb93b596 100644 --- a/mlir-tensorrt/test/python/mlir_tensorrt_compiler/compiler_api/test_options_context.py +++ b/mlir-tensorrt/test/python/mlir_tensorrt_compiler/compiler_api/test_options_context.py @@ -16,6 +16,7 @@ opts = api.OptionsContext( client, "stable-hlo-to-executable", + # Set some options explicitly so we can spot check the `print` output. [ "--tensorrt-builder-opt-level=3", "--tensorrt-strongly-typed=false", @@ -27,4 +28,4 @@ # CHECK: InvalidArgument: InvalidArgument: non-existent-options-type is not a valid option type. Valid options were: stable-hlo-to-executable -# CHECK: Options[{{.*}}] +# CHECK: Options[{{.*--tensorrt-workspace-memory-pool-limit=1073741824.*--tensorrt-strongly-typed=false.*--tensorrt-builder-opt-level=3.*}}]