WIP: Adding TRT options/task

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

@@ -1 +1,2 @@
 add_subdirectory(StablehloToExecutable)
+add_subdirectory(TensorRTToExecutable)

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

@@ -78,8 +78,7 @@ optionsCreateFromArgs(mlir::MLIRContext *context,
       [&errMsg](const llvm::StringError &err) { errMsg = err.getMessage(); });
 
   if (errMsg)
-    return getInternalErrorStatus("failed to initialize options: %s",
-                                  errMsg->c_str());
+    return getInternalErrorStatus("failed to initialize options: {0}", *errMsg);
 
   return result;
 }

mlir-tensorrt/compiler/include/mlir-tensorrt/Compiler/TensorRTToExecutable/CMakeLists.txt

@@ -0,0 +1,4 @@
+set(_TABLEGEN_ARGS )
+set(LLVM_TARGET_DEFINITIONS Passes.td)
+mlir_tablegen(Passes.h.inc -gen-pass-decls -name TensorRTToExecutable ${_TABLEGEN_ARGS})
+add_public_tablegen_target(MLIRTensorRTTensorRTToExecutableIncGen)

mlir-tensorrt/compiler/include/mlir-tensorrt/Compiler/TensorRTToExecutable/Passes.h

@@ -0,0 +1,49 @@
+//===- Passes.h ----------------------------------------------===//
+//
+// SPDX-FileCopyrightText: Copyright 2025 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.
+//
+//===----------------------------------------------------------------------===//
+///
+/// Declarations for opt tool pipeline command-line registration for pipelines
+/// related to "tensorrt-to-executable".
+///
+//===----------------------------------------------------------------------===//
+#ifndef MLIR_TENSORRT_COMPILER_TENSORRTTOEXECUTABLE_PASSES
+#define MLIR_TENSORRT_COMPILER_TENSORRTTOEXECUTABLE_PASSES
+
+#include <memory>
+#include <mlir/Pass/Pass.h>
+
+namespace mlirtrt::compiler {
+
+//===----------------------------------------------------------------------===//
+// Add Tablegen'd pass declarations and registration methods.
+//===----------------------------------------------------------------------===//
+#define GEN_PASS_DECL
+#define GEN_PASS_REGISTRATION
+#include "mlir-tensorrt/Compiler/TensorRTToExecutable/Passes.h.inc"
+
+//===----------------------------------------------------------------------===//
+// Pipeline Registrations
+//===----------------------------------------------------------------------===//
+
+/// Register the TensorRT clustering and compilation pipelines.
+void registerTensorRTToExecutablePipelines();
+
+} // namespace mlirtrt::compiler
+
+#endif // MLIR_TENSORRT_COMPILER_TENSORRTTOEXECUTABLE_PASSES

mlir-tensorrt/compiler/include/mlir-tensorrt/Compiler/TensorRTToExecutable/Passes.td

@@ -0,0 +1,38 @@
+//===- Passes.td ----------------------------------------------------------===//
+//
+// SPDX-FileCopyrightText: Copyright 2025 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.
+//
+//===----------------------------------------------------------------------===//
+#ifndef MLIR_TENSORRT_COMPILER_TENSORRTTOEXECUTABLE_PASSES
+#define MLIR_TENSORRT_COMPILER_TENSORRTTOEXECUTABLE_PASSES
+
+include "mlir/Pass/PassBase.td"
+
+//===----------------------------------------------------------------------===//
+// OutlineTensorRTOpPass
+//===----------------------------------------------------------------------===//
+
+def OutlineTensorRTOpPass : Pass<"outline-tensorrt-op",
+      "::mlir::ModuleOp"> {
+  let summary = "Outline all tensorrt ops into a tensorrt module";
+
+  let dependentDialects = [
+    "::mlir::plan::PlanDialect"
+  ];
+}
+
+#endif // MLIR_TENSORRT_COMPILER_TENSORRTTOEXECUTABLE_PASSES

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

@@ -0,0 +1,97 @@
+//===- TensorRTToExecutable.h -----------------------------------*- C++ -*-===//
+//
+// SPDX-FileCopyrightText: Copyright 2025 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.
+//
+//===----------------------------------------------------------------------===//
+#ifndef MLIR_TENSORRT_COMPILER_TENSORRTTOEXECUTABLE
+#define MLIR_TENSORRT_COMPILER_TENSORRTTOEXECUTABLE
+
+// TODO (pranavm): MLIR_TRT_TARGET_TENSORRT is only needed because we pull in
+// the TranslateToTensorRT.h header. If we move the translation options, we
+// won't need it.
+#ifdef MLIR_TRT_TARGET_TENSORRT
+#include "mlir-tensorrt-dialect/Target/TranslateToTensorRT.h"
+
+#include "mlir-tensorrt-dialect/Utils/OptionsBundle.h"
+#include "mlir-tensorrt/Compiler/Client.h"
+
+namespace mlirtrt::compiler {
+
+//===----------------------------------------------------------------------===//
+// TensorRTToExecutableOptions
+//===----------------------------------------------------------------------===//
+
+class TensorRTToExecutableTask;
+
+// TODO (pranavm): Figure out a better way to reuse TRT translation options -
+// maybe move to options providers?
+struct TensorRTOptions : public OptionsProvider<TensorRTOptions> {
+public:
+  using OptionsProvider::OptionsProvider;
+  mlir::tensorrt::TensorRTTranslationOptions options;
+
+  TensorRTOptions(mlir::OptionsContext &ctx) : OptionsProvider(ctx) {
+    options.addToOptions(ctx);
+  }
+};
+
+struct TensorRTToExecutableOptions
+    : public mlir::OptionsBundle<DeviceOptions, DebugOptions, ExecutorOptions,
+                                 TensorRTOptions> {
+  // Default initialization does not require any extensions.
+  TensorRTToExecutableOptions() = default;
+
+  TensorRTToExecutableOptions(TaskExtensionRegistry extensions);
+
+  Option<std::string> entrypoint{this, "entrypoint", llvm::cl::init("main"),
+                                 llvm::cl::desc("entrypoint function name")};
+};
+
+//===----------------------------------------------------------------------===//
+// TensorRTToExecutableTask
+//===----------------------------------------------------------------------===//
+
+class TensorRTToExecutableTask
+    : public CompilationTask<TensorRTToExecutableTask,
+                             TensorRTToExecutableOptions> {
+public:
+  TensorRTToExecutableTask(mlir::MLIRContext *ctx,
+                            const TensorRTToExecutableOptions &options);
+
+  /// Build the clustering pipeline that occurs on TensorRT Ops.
+  static void
+  buildTensorRTClusteringPipeline(mlir::OpPassManager &pm,
+                                   const TensorRTToExecutableOptions &options);
+
+  /// Build the compilation pipeline that runs after clustering.
+  static void
+  buildPostClusteringPipeline(mlir::OpPassManager &pm,
+                              const TensorRTToExecutableOptions &options);
+
+  static void populatePassManager(mlir::PassManager &pm,
+                                  const TensorRTToExecutableOptions &options);
+};
+
+/// Register the task/options with the client's registry.
+void registerTensorRTToExecutableTask();
+
+} // namespace mlirtrt::compiler
+
+MLIR_DECLARE_EXPLICIT_TYPE_ID(mlirtrt::compiler::TensorRTToExecutableTask)
+
+#endif
+#endif // MLIR_TENSORRT_COMPILER_TENSORRTTOEXECUTABLE

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

@@ -31,6 +31,7 @@
 #ifdef MLIR_TRT_ENABLE_HLO
 #include "mlir-tensorrt/Compiler/StablehloToExecutable/Passes.h"
 #include "mlir-tensorrt/Compiler/StablehloToExecutable/StablehloToExecutable.h"
+#include "mlir-tensorrt/Compiler/TensorRTToExecutable/Passes.h"
 #include "mlir-tensorrt/Dialect/Plan/Transforms/Passes.h"
 #include "mlir-tensorrt/Dialect/StableHloExt/Transforms/Passes.h"
 #include "mlir-tensorrt/Pipelines/StableHloInputPipelines.h"
@@ -63,6 +64,10 @@ inline void registerAllMlirTensorRtPasses() {
   plan::registerPlanDialectPipelines();
 #endif // MLIR_TRT_ENABLE_HLO
 
+#ifdef MLIR_TRT_TARGET_TENSORRT
+  mlirtrt::compiler::registerTensorRTToExecutablePipelines();
+#endif // MLIR_TRT_TARGET_TENSORRT
+
 #ifdef MLIR_TRT_ENABLE_EXECUTOR
   registerConvertCUDAToExecutorPass();
   bufferization::registerBufferizationPasses();

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

@@ -6,4 +6,5 @@ add_mlir_tensorrt_public_c_api_library(MLIRTensorRTCAPIRegisterAllDialects
     MLIRTensorRTRegistration
     MLIRCAPIIR
     MLIRTensorRTCompilerStableHloToExecutable
+    MLIRTensorRTCompilerTensorRTToExecutable
 )

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

@@ -24,6 +24,7 @@
 
 #include "mlir-tensorrt-c/Compiler/Registration/RegisterAllDialects.h"
 #include "mlir-tensorrt/Compiler/StablehloToExecutable/StablehloToExecutable.h"
+#include "mlir-tensorrt/Compiler/TensorRTToExecutable/TensorRTToExecutable.h"
 #include "mlir-tensorrt/Registration/RegisterMlirTensorRtDialects.h"
 #include "mlir-tensorrt/Registration/RegisterMlirTensorRtPasses.h"
 #include "mlir/CAPI/IR.h"
@@ -38,4 +39,5 @@ void mtrtCompilerRegisterPasses() {
 
 void mtrtCompilerRegisterTasks() {
   mlirtrt::compiler::registerStableHloToExecutableTask();
+  mlirtrt::compiler::registerTensorRTToExecutableTask();
 }

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

@@ -15,4 +15,5 @@ add_mlir_tensorrt_library(MLIRTensorRTCompilerClient
     MLIRTensorRTSupportDeviceInfo
 )
 
-add_subdirectory(StablehloToExecutable)
+add_subdirectory(StablehloToExecutable)
+add_subdirectory(TensorRTToExecutable)

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

@@ -0,0 +1,17 @@
+add_mlir_tensorrt_library(MLIRTensorRTCompilerTensorRTToExecutable
+    TensorRTToExecutable.cpp
+    Passes.cpp
+
+    PARTIAL_SOURCES_INTENDED
+
+    DEPENDS
+    MLIRTensorRTTensorRTToExecutableIncGen
+
+    LINK_LIBS PUBLIC
+    MLIRIR
+    MLIRTensorRTRegistration
+    MLIRTensorRTTargetLua
+    MLIRTensorRTOptionUtils
+    MLIRTensorRTTargetTensorRT
+    MLIRTensorRTCompilerClient
+    )