WIP: Adding TRT options/task

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

@@ -128,6 +128,19 @@ struct DeviceOptions : public OptionsProvider<DeviceOptions> {
   llvm::Error finalizeImpl();
 };
 
+struct CommonCompilationOptions
+    : public OptionsProvider<CommonCompilationOptions> {
+public:
+  /// Entrypoint function name.
+  std::string entrypoint = "main";
+
+public:
+  void addToOptions(mlir::OptionsContext &context) {
+    context.addOption("entrypoint", entrypoint, llvm::cl::init("main"),
+                      llvm::cl::desc("entrypoint function name"));
+  }
+};
+
 } // namespace mlirtrt::compiler
 
 #endif // MLIR_TENSORRT_COMPILER_OPTIONS

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

@@ -0,0 +1,30 @@
+//===- PassManagerUtils.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.
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir-tensorrt/Compiler/OptionsProviders.h"
+#include "mlir/Pass/PassManager.h"
+
+//===----------------------------------------------------------------------===//
+// Common helpers
+//===----------------------------------------------------------------------===//
+
+mlir::LogicalResult
+setupPassManager(mlir::PassManager &pm,
+                 const mlirtrt::compiler::DebugOptions &options);

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

@@ -52,7 +52,8 @@ namespace mlirtrt::compiler {
 class StablehloToExecutableTask;
 
 struct StablehloToExecutableOptions
-    : public mlir::OptionsBundle<DebugOptions, ExecutorOptions, DeviceOptions> {
+    : public mlir::OptionsBundle<DebugOptions, ExecutorOptions, DeviceOptions,
+                                 CommonCompilationOptions> {
   /// Initializes the options. The extensions in the provided registry
   /// must be extensions for the StableHloToExecutable task.
   StablehloToExecutableOptions(TaskExtensionRegistry extensions);
@@ -64,9 +65,6 @@ struct StablehloToExecutableOptions
   /// and backend types that support allocating results.
   bool enableNonDPSReturns = false;
 
-  /// Entrypoint function name.
-  std::string entrypoint = "main";
-
   /// Base class for extensions associated with StableHloToExecutableTask.
   class ExtensionBase : public TaskExtensionBase {
   public:

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

@@ -0,0 +1,88 @@
+//===- TensorRTToExecutable.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.
+//
+//===----------------------------------------------------------------------===//
+#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/Options.h"
+#include "mlir-tensorrt-dialect/Utils/OptionsBundle.h"
+#include "mlir-tensorrt/Compiler/Client.h"
+#include "mlir-tensorrt/Compiler/Extension.h"
+#include "mlir-tensorrt/Compiler/OptionsProviders.h"
+#include "mlir/Support/TypeID.h"
+
+namespace mlirtrt::compiler {
+
+//===----------------------------------------------------------------------===//
+// TensorRTToExecutableOptions
+//===----------------------------------------------------------------------===//
+
+// TODO (pranavm): Figure out a better way to reuse TRT translation options -
+// maybe move to options providers?
+struct TensorRTOptions
+    : public mlirtrt::compiler::OptionsProvider<TensorRTOptions> {
+  mlir::tensorrt::TensorRTTranslationOptions options;
+
+  void addToOptions(mlir::OptionsContext &context) {
+    options.addToOptions(context);
+  }
+};
+
+struct TensorRTToExecutableOptions
+    : public mlir::OptionsBundle<DeviceOptions, DebugOptions, ExecutorOptions,
+                                 CommonCompilationOptions, TensorRTOptions> {
+
+  TensorRTToExecutableOptions(TaskExtensionRegistry extensions);
+};
+
+//===----------------------------------------------------------------------===//
+// TensorRTToExecutableTask
+//===----------------------------------------------------------------------===//
+
+class TensorRTToExecutableTask
+    : public CompilationTask<TensorRTToExecutableTask,
+                             TensorRTToExecutableOptions> {
+public:
+  using Base::Base;
+
+  static void populatePassManager(mlir::PassManager &pm,
+                                  const TensorRTToExecutableOptions &options);
+};
+
+/// Register the task/options with the client's registry.
+void registerTensorRTToExecutableTask();
+
+//===----------------------------------------------------------------------===//
+// Pipeline Registrations
+//===----------------------------------------------------------------------===//
+
+// TODO (pranavm): How to do pipeline registration?
+// void registerTensorRTPipelines();
+
+} // 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

@@ -23,6 +23,7 @@
 #define REGISTRATION_REGISTERMLIRTENSORRTPASSES_H
 
 #include "mlir-tensorrt-dialect/TensorRT/Transforms/Passes.h"
+#include "mlir-tensorrt/Compiler/TensorRTToExecutable.h"
 #include "mlir-tensorrt/Conversion/Passes.h"
 #include "mlir-tensorrt/Transforms/Passes.h"
 #include "mlir/Conversion/Passes.h"
@@ -52,6 +53,12 @@ inline void registerAllMlirTensorRtPasses() {
   mlir::registerTransformsPasses();
   mlir::registerConvertPDLToPDLInterp();
 
+  // TODO (pranavm): Check if this needs to be conditional - the TRT passes
+  // above are not.
+#ifdef MLIR_TRT_TARGET_TENSORRT
+  mlirtrt::compiler::registerTensorRTToExecutableTask();
+#endif
+
 #ifdef MLIR_TRT_ENABLE_HLO
   mlirtrt::compiler::registerStablehloClusteringPipelines();
   registerStableHloInputPipelines();

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

@@ -3,6 +3,7 @@ add_mlir_tensorrt_library(MLIRTensorRTCompilerClient
     Extension.cpp
     OptionsRegistry.cpp
     OptionsProviders.cpp
+    PassManagerUtils.cpp
     PARTIAL_SOURCES_INTENDED
 
     LINK_LIBS PUBLIC
@@ -19,6 +20,8 @@ add_mlir_tensorrt_library(MLIRTensorRTCompilerStableHloToExecutable
     StableHloToExecutable.cpp
     # TODO: TensorRTExtension should be an independent library.
     TensorRTExtension/TensorRTExtension.cpp
+    # TODO (pranavm): TensorRTToExecutable should probably be a separate library
+    TensorRTToExecutable.cpp
 
     PARTIAL_SOURCES_INTENDED
 

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

@@ -0,0 +1,43 @@
+//===- PassManagerUtils.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.
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir-tensorrt/Compiler/PassManagerUtils.h"
+
+using namespace mlirtrt::compiler;
+using namespace mlir;
+
+//===----------------------------------------------------------------------===//
+// Common helpers
+//===----------------------------------------------------------------------===//
+
+mlir::LogicalResult setupPassManager(mlir::PassManager &pm,
+                                     const DebugOptions &options) {
+  pm.enableVerifier(true);
+  mlir::applyDefaultTimingPassManagerCLOptions(pm);
+  if (failed(mlir::applyPassManagerCLOptions(pm)))
+    return mlir::failure();
+  if (!options.dumpIRPath.empty()) {
+    pm.enableIRPrintingToFileTree(
+        [](Pass *, Operation *) { return false; },
+        [](Pass *, Operation *) { return true; }, true, false, false,
+        options.dumpIRPath, OpPrintingFlags().elideLargeElementsAttrs(32));
+  }
+  return mlir::success();
+}

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

@@ -28,10 +28,10 @@
 #include "mlir-executor/Support/Status.h"
 #include "mlir-executor/Target/Lua/TranslateToRuntimeExecutable.h"
 #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/OptionsProviders.h"
 #include "mlir-tensorrt/Compiler/OptionsRegistry.h"
+#include "mlir-tensorrt/Compiler/PassManagerUtils.h"
 #include "mlir-tensorrt/Compiler/TensorRTExtension/TensorRTExtension.h"
 #include "mlir-tensorrt/Conversion/Passes.h"
 #include "mlir-tensorrt/Dialect/Plan/Transforms/Passes.h"
@@ -60,25 +60,6 @@ using namespace mlir;
 
 #ifdef MLIR_TRT_ENABLE_HLO
 
-//===----------------------------------------------------------------------===//
-// Common helpers
-//===----------------------------------------------------------------------===//
-
-static mlir::LogicalResult setupPassManager(mlir::PassManager &pm,
-                                            const DebugOptions &options) {
-  pm.enableVerifier(true);
-  mlir::applyDefaultTimingPassManagerCLOptions(pm);
-  if (failed(mlir::applyPassManagerCLOptions(pm)))
-    return mlir::failure();
-  if (!options.dumpIRPath.empty()) {
-    pm.enableIRPrintingToFileTree(
-        [](Pass *, Operation *) { return false; },
-        [](Pass *, Operation *) { return true; }, true, false, false,
-        options.dumpIRPath, OpPrintingFlags().elideLargeElementsAttrs(32));
-  }
-  return mlir::success();
-}
-
 //===----------------------------------------------------------------------===//
 // Adhoc Passes
 //===----------------------------------------------------------------------===//
@@ -162,9 +143,6 @@ 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("entrypoint", entrypoint, llvm::cl::init("main"),
-            llvm::cl::desc("entrypoint function name"));
 }
 
 //===----------------------------------------------------------------------===//
@@ -189,7 +167,7 @@ void StablehloToExecutableTask::buildStablehloClusteringPipeline(
   populateExtensionPasses(pm, opts, Phase::PreClustering);
 
   plan::StablehloClusteringPassOptions clusteringOpts{};
-  clusteringOpts.entrypoint = opts.entrypoint;
+  clusteringOpts.entrypoint = opts.get<CommonCompilationOptions>().entrypoint;
   plan::buildPlanSegmentationPipeline(pm, clusteringOpts);
 
   // Compile outlined funcs marked with `cluster.host`. The HLO in these
@@ -465,7 +443,7 @@ static StablehloToExecutableOptions populateStablehloClusteringPipelineOpts(
       cliOpts.deviceMaxSharedMemoryPerBlockKb;
   opts.get<DeviceOptions>().shouldInferFromHost =
       cliOpts.inferDeviceOptionsFromHost;
-  opts.entrypoint = cliOpts.entrypoint;
+  opts.get<CommonCompilationOptions>().entrypoint = cliOpts.entrypoint;
   return opts;
 }