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/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 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.
+//
+//===----------------------------------------------------------------------===//
+///
+/// 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,35 @@
+//===- Passes.td ----------------------------------------------------------===//
+//
+// 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_PASSES
+#define MLIR_TENSORRT_COMPILER_TENSORRTTOEXECUTABLE_PASSES
+
+include "mlir/Pass/PassBase.td"
+
+//===----------------------------------------------------------------------===//
+// OutlineTensorRTOpPass
+//===----------------------------------------------------------------------===//
+// TODO: what are the dependent dialects? what are the options?
+
+def OutlineTensorRTOpPass : Pass<"outline-tensorrt-op",
+      "::mlir::ModuleOp"> {
+  let summary = "Outline all tensorrt ops into a tensorrt module";
+}
+
+#endif // MLIR_TENSORRT_COMPILER_TENSORRTTOEXECUTABLE_PASSES

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

@@ -0,0 +1,107 @@
+//===- 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-executor/Runtime/API/API.h"
+#include "mlir-executor/Support/Status.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/IR/BuiltinOps.h"
+#include "mlir/Pass/PassManager.h"
+#include "mlir/Support/TypeID.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) {}
+
+  void addToOptions(mlir::OptionsContext &context) {
+    options.addToOptions(context);
+  }
+};
+
+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

@@ -23,6 +23,7 @@
 #define REGISTRATION_REGISTERMLIRTENSORRTPASSES_H
 
 #include "mlir-tensorrt-dialect/TensorRT/Transforms/Passes.h"
+#include "mlir-tensorrt/Compiler/TensorRTToExecutable/TensorRTToExecutable.h"
 #include "mlir-tensorrt/Conversion/Passes.h"
 #include "mlir-tensorrt/Transforms/Passes.h"
 #include "mlir/Conversion/Passes.h"
@@ -53,6 +54,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::registerStablehloToExecutablePasses();
   mlirtrt::compiler::registerStablehloToExecutablePipelines();

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/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
+    )