Skip to content

Commit 30f9816

Browse files
christopherbatechristopherbate
authored
[mlir-tensorrt] Integrate internal changes (#651)
-- b7c0ac3c95aac78b563cfcf1926fcfb0de20ef67 by Chris Bate <[email protected]>: [compiler] Fix bufferization for `trtrt.enqueue` This fixes a bug where `trtrt.enqueue` could incorrectly allow the same buffer to be used as both an input and an output. This is technically possible since for many TensorRT programs, the inputs will be read before any outputs are written. However, in general we cannot actually make this assumption, which is codified in the bufferization infrastructure via the `bufferizesToElementwiseAccess` method of the BufferizableOpInterface. Previously, we were returning true for this method for `trtrt.enqueue` despite not analyzing the corresponding function. Since the ability to re-use an input pointer avoids an allocation, it can be critical to good performance in certain cases, especially in loops. Therefore, in the future, we need to add back the ability to specify `true` when some basic analysis reduces the risk of that being incorrect. -- 6fab2aa80a7776fe9afd15a59fcce916b61570cd by Chris Bate <[email protected]>: Refactor registration of dialects and passes - Improve uniformity of registration for dialects/passes/extensions. - Remove some CMake flags which we never set to OFF and probably never should. - Improve format of preprocessor-gaurded code blocks by generating a project-wide header `Features.h` header containing some convenience macros that look a lot nicer than `#ifdef..#endif` blocks. -- 96f45d6c6f3f2f1354984944f194fdc8a857ff34 by Sagar Shelke <[email protected]>: [tensorrt] Fix a bug in `trtSetWeights` function in adaptor. This MR fixes a bug in `trtSetWeights` method in NvInferAdaptor. This method stores weights in weight map in order to keep them alive until engine is built. Previously, a new weight in map was created but original weight was never copied. This created issues like getting zeros in the output to all the way getting NaNs. Change copies original weights into map. -- 08fc170f769d9cedce2541da14b101e0947175e2 by Chris Bate <[email protected]>: [compiler] Add explicit memory space assignment Previously, the 'plan-module-bufferize' pass used the encodings of tensor types to infer the memory space of memref types. However, not every tensor had an encoding. There was a notion of a default memory space (`#plan.memory_space<device>`) which was used to deduce the memory space when no tensor encoding was present. However, this could result in situations where the program could not be bufferized correctly, for example if we had `tensor.cast` operations that added or removed the encoding. In such cases, we are relying on brittle logic of the bufferization infrastructure to somehow deduce the correct memory space, and this may not always work. Therefore, this change adds a new pass `plan-assign-memory-spaces` which explicitly assigns a `#plan.memory_space` encoding to all tensor types. In proceeds in two steps. The first step simply assigns the 'device' space to all tensors. Then the second step performs some minor optimizations to avoid unnecessary host-device transfers after bufferization. An end-to-end test case for the bufferization pipeline is added which was previously failing. Interestingly, the addition of the new pass results in better bufferization for cases where we can't detensorize while loops. -- 1be85552b63d7a06d9e27b7d9bd57e2e62589356 by Christopher Bate <[email protected]>: [compiler] Add additional IR printing flags to the compiler API Exposes `-mlir-elide-elementsattrs-if-larger` and `-mlir-elide-resource-strings-if-larger` to the compiler API via DebugOptions provider. -- cfc7a93555f2778711610468b2f56e9f2fb28148 by Chris Bate <[email protected]>: [tensorrt] Make the TensorRT builder logger global This change makes the TensorRT builder logger a global singleton. This change has been made due to an issue with the TRT API -- the logger does not actually get associated with the TRT object created via `createInferBuilder` or `createInferRuntime` -- instead the TRT API maintains a global logger that is populated with push/op mechanism, except the stack size can only be 1. So if the Builder object lifetime overlaps the Runtime object lifetime, then the Runtime will ignore the logger passed to `createInferRuntime` and instead use the global logger populated by `createInferBuilder`. This is clearly problematic if the API user (e.g. MLIR-TRT compiler) thinks they can destroy the logger used to create the Builder object after the Builder object is destroyed. In such a case, the Runtime will still try to use the builder's logger, which results in use-after-free. To reduce the likelihood of this ocurring, we make the builder logger a global singleton. Note, however, that we can't guarantee that the logger is not destroyed when the compiler library is closed via `dlclose`, and we can't make any guaruntees about global shutdown order. -- 4f215aa9513d7751572cd83d8575f251c5a61edb by Christopher Bate <[email protected]>: NFC: Fix two compiler warnings - 'captured structured bindings are a C++20 extension' - Unused function when MLIR_TRT_ENABLE_NCCL=ON -- 6f4cc4ad0eb6a4c554e5e36fc23a4b3b56ba607e by Christopher Bate <[email protected]>: NFC: update test runner configs to take into account system memory This change updates LIT parallelism setting logic to parallelism based on system memory. GitOrigin-RevId: 6f4cc4ad0eb6a4c554e5e36fc23a4b3b56ba607e
1 parent 3cab972 commit 30f9816

File tree

51 files changed

+1144
-438
lines changed

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

51 files changed

+1144
-438
lines changed

‎mlir-tensorrt/.gitignore‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@
33
*.log
44
**/llvm-project/**
55
**/llvm-project/
6+
CMakeUserPresets.json
67

78
# Docs build artifacts
89
/public/

‎mlir-tensorrt/CMakeLists.txt‎

Lines changed: 29 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -90,6 +90,35 @@ if(PROJECT_IS_TOP_LEVEL)
9090
set(CMAKE_EXPORT_COMPILE_COMMANDS ON)
9191
endif()
9292

93+
# -------------------------------------------------
94+
# Option validation
95+
# -------------------------------------------------
96+
# Write out a header file containing convenience macros for each flag.
97+
function(mtrt_write_feature_flags_header)
98+
set(feature_flags_header
99+
"${CMAKE_CURRENT_BINARY_DIR}/include/mlir-tensorrt/Features.h")
100+
101+
# Generate the header at configure time
102+
file(WRITE "${feature_flags_header}" [[
103+
// Auto-generated feature macros, do not edit.
104+
#ifndef MLIR_TENSORRT_FEATURES_H
105+
#define MLIR_TENSORRT_FEATURES_H
106+
107+
]])
108+
109+
foreach(FEATURE IN LISTS MLIR_TRT_FEATURE_FLAGS)
110+
file(APPEND "${feature_flags_header}" "#ifdef ${FEATURE}\n")
111+
file(APPEND "${feature_flags_header}" "#define IF_${FEATURE}(code) do { code } while (0)\n")
112+
file(APPEND "${feature_flags_header}" "#else\n")
113+
file(APPEND "${feature_flags_header}" "#define IF_${FEATURE}(code) do {} while (0)\n")
114+
file(APPEND "${feature_flags_header}" "#endif // ${FEATURE}\n\n")
115+
endforeach()
116+
file(APPEND "${feature_flags_header}" "#endif // MLIR_TENSORRT_FEATURES_H\n")
117+
endfunction()
118+
119+
mtrt_write_feature_flags_header()
120+
include_directories("${CMAKE_CURRENT_BINARY_DIR}/include")
121+
93122
# -------------------------------------------------
94123
# Setup LLVM/MLIR
95124
# -------------------------------------------------

‎mlir-tensorrt/build_tools/cmake/Targets.cmake‎

Lines changed: 12 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -44,6 +44,16 @@ function(add_mlir_tensorrt_public_c_api_library target)
4444
endif()
4545
endfunction()
4646

47+
# --------------------------------------------------------------
48+
# Adds an upstream MLIR library target to the
49+
# MLIR_TENSORRT_LIBS global property list to capture it as an
50+
# implicit dependency for all final tools and compiler
51+
# end-user products.
52+
# --------------------------------------------------------------
53+
function(add_mlir_tensorrt_compiler_dependency target)
54+
set_property(GLOBAL APPEND PROPERTY MLIR_TENSORRT_LIBS ${target})
55+
endfunction()
56+
4757
# ------------------------------------------------------------------------------
4858
# A wrapper around `add_mlir_dialect_library` that also appends the dialect
4959
# library to the global `MLIR_TENSORRT_DIALECT_LIBS` list property.
@@ -99,11 +109,11 @@ function(add_mlir_tensorrt_backend_library target)
99109
BASE_DIRECTORY ${CMAKE_CURRENT_BINARY_DIR})
100110

101111
set(LLVM_TARGET_DEFINITIONS "${SRC_TD}")
102-
112+
103113
string(REPLACE ".td" "Attrs.h.inc" h_inc_file ${BIN_TD})
104114
string(REPLACE ".td" "Attrs.cpp.inc" cpp_inc_file ${BIN_TD})
105115
mlir_tablegen("${h_inc_file}" -gen-attrdef-decls)
106-
mlir_tablegen("${cpp_inc_file}" -gen-attrdef-defs)
116+
mlir_tablegen("${cpp_inc_file}" -gen-attrdef-defs)
107117

108118
add_public_tablegen_target(${target}IncGen)
109119

‎mlir-tensorrt/compiler/CMakeLists.txt‎

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,29 @@ set(MLIR_TENSORRT_COMPILER_INCLUDE_DIRS
99
include_directories("${CMAKE_CURRENT_BINARY_DIR}/include"
1010
"${CMAKE_CURRENT_SOURCE_DIR}/include")
1111

12+
# We use the MLIR_TENSORRT_LIBS global property to aggregate the list of all
13+
# declared compiler libraries. This is helpful for simplifying the link
14+
# dependency declarations for tools that must link "the world" like
15+
# `mlir-tensorrt-opt`.
16+
17+
# Because MLIR has a level of indirection that lets implementation for
18+
# interfaces be provided by separate implementation code
19+
# ("PromisedInterfaces/ExternalModels") which is registered at runtime, it is
20+
# difficult to capture all dependencies for dialects we require in the compiler
21+
# purely through target dependency properties. To see what we require from
22+
# usptream, look at the file `mlir-tensorrt/InitAllDialects.h`. Therefore, we
23+
# manually enumerate some dependencies here, mainly for providing the functions
24+
# registering interface external models.
25+
add_mlir_tensorrt_compiler_dependency(MLIRArithTransforms)
26+
add_mlir_tensorrt_compiler_dependency(MLIRArithValueBoundsOpInterfaceImpl)
27+
add_mlir_tensorrt_compiler_dependency(MLIRAsyncDialect)
28+
add_mlir_tensorrt_compiler_dependency(MLIRBufferizationTransforms)
29+
add_mlir_tensorrt_compiler_dependency(MLIRControlFlowTransforms)
30+
add_mlir_tensorrt_compiler_dependency(MLIRNVVMTarget)
31+
add_mlir_tensorrt_compiler_dependency(MLIRPtrDialect)
32+
add_mlir_tensorrt_compiler_dependency(MLIRTargetLLVM)
33+
add_mlir_tensorrt_compiler_dependency(MLIRTensorTransformOps)
34+
1235
add_subdirectory(include)
1336
add_subdirectory(lib)
1437
add_subdirectory(test)

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

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -113,6 +113,20 @@ struct DebugOptions : public OptionsProvider {
113113
"tree rooted at this directory. Use in conjunction with "
114114
"mlir-print-ir-* flags")};
115115

116+
//===----------------------------------------------------------------------===//
117+
// Printing Flags
118+
//===----------------------------------------------------------------------===//
119+
120+
Option<unsigned> elideElementsAttrIfLarger{
121+
this->ctx, "mlir-elide-elementsattrs-if-larger",
122+
llvm::cl::desc("Elide ElementsAttrs with \"...\" that have "
123+
"more elements than the given upper limit")};
124+
125+
Option<unsigned> elideResourceStringsIfLarger{
126+
this->ctx, "mlir-elide-resource-strings-if-larger",
127+
llvm::cl::desc(
128+
"Elide printing value of resources if string is too long in chars.")};
129+
116130
//===--------------------------------------------------------------------===//
117131
// Pass Statistics
118132
//===--------------------------------------------------------------------===//

‎mlir-tensorrt/compiler/include/mlir-tensorrt/Conversion/CMakeLists.txt‎

Lines changed: 0 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -6,9 +6,6 @@ set(_TABLEGEN_ARGS )
66
if(MLIR_TRT_ENABLE_HLO)
77
list(APPEND _TABLEGEN_ARGS -DMLIR_TENSORRT_ENABLE_HLO)
88
endif()
9-
if(MLIR_TRT_ENABLE_EXECUTOR)
10-
list(APPEND _TABLEGEN_ARGS -DMLIR_TENSORRT_ENABLE_EXECUTOR)
11-
endif()
129

1310
mlir_tablegen(Passes.h.inc -gen-pass-decls -name MLIRTensorRTConversion ${_TABLEGEN_ARGS})
1411
add_public_tablegen_target(MLIRTensorRTConversionPassIncGen)

‎mlir-tensorrt/compiler/include/mlir-tensorrt/Conversion/Passes.td‎

Lines changed: 0 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -125,8 +125,6 @@ def ConvertTensorRTToEmitCPass : Pass<"convert-tensorrt-to-emitc",
125125
let dependentDialects = ["::mlir::emitc::EmitCDialect"];
126126
}
127127

128-
#ifdef MLIR_TENSORRT_ENABLE_EXECUTOR
129-
130128
//===----------------------------------------------------------------------===//
131129
// ConvertMemRefToCUDAPass
132130
//===----------------------------------------------------------------------===//
@@ -312,9 +310,6 @@ def ConvertTensorRTRuntimeToExecutorPass : Pass<"convert-tensorrt-runtime-to-exe
312310
let options = ConvertToExecutorOptions;
313311
}
314312

315-
316-
#endif // MLIR_TENSORRT_ENABLE_EXECUTOR
317-
318313
#ifdef MLIR_TENSORRT_ENABLE_HLO
319314
//===----------------------------------------------------------------------===//
320315
// ConvertStablehloToScfPass

‎mlir-tensorrt/compiler/include/mlir-tensorrt/Dialect/CMakeLists.txt‎

Lines changed: 2 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,7 @@
11
if(MLIR_TRT_ENABLE_HLO)
22
add_subdirectory(StablehloExt)
3-
add_subdirectory(Plan)
43
endif()
54

5+
add_subdirectory(CUDA)
6+
add_subdirectory(Plan)
67
add_subdirectory(TensorRTRuntime)
7-
8-
if(MLIR_TRT_ENABLE_EXECUTOR)
9-
add_subdirectory(CUDA)
10-
endif()

‎mlir-tensorrt/compiler/include/mlir-tensorrt/Dialect/Plan/Transforms/Passes.td‎

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -428,6 +428,27 @@ def PostClusteringValidationPass : Pass<"post-clustering-validation", "func::Fun
428428
}];
429429
}
430430

431+
//===----------------------------------------------------------------------===//
432+
// PlanAssignMemorySpacesPass
433+
//===----------------------------------------------------------------------===//
434+
435+
def PlanAssignMemorySpacesPass : Pass<"plan-assign-memory-spaces",
436+
"::mlir::ModuleOp"> {
437+
let summary = "assigns memory spaces encodings to tensor types";
438+
439+
let description = [{
440+
This pass applies a type conversion that adds a '#plan.memory_space'
441+
attribute to all tensor types in the top-level module that do not already
442+
have an encoding.
443+
}];
444+
445+
let dependentDialects = [
446+
"::mlir::plan::PlanDialect",
447+
"::mlir::bufferization::BufferizationDialect",
448+
"::mlir::tensor::TensorDialect"
449+
];
450+
}
451+
431452
//===----------------------------------------------------------------------===//
432453
// PlanAllocTensorsPass
433454
//===----------------------------------------------------------------------===//

‎mlir-tensorrt/compiler/include/mlir-tensorrt/InitAllDialects.h‎

Lines changed: 199 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,199 @@
1+
//===- InitAllDialects.h ----------------------------------------*- C++ -*-===//
2+
//
3+
// SPDX-FileCopyrightText: Copyright 2025 NVIDIA CORPORATION & AFFILIATES.
4+
// All rights reserved.
5+
// SPDX-License-Identifier: Apache-2.0
6+
//
7+
// Licensed under the Apache License, Version 2.0 (the "License");
8+
// you may not use this file except in compliance with the License.
9+
// You may obtain a copy of the License at
10+
//
11+
// http://www.apache.org/licenses/LICENSE-2.0
12+
//
13+
// Unless required by applicable law or agreed to in writing, software
14+
// distributed under the License is distributed on an "AS IS" BASIS,
15+
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
16+
// See the License for the specific language governing permissions and
17+
// limitations under the License.
18+
//
19+
//===----------------------------------------------------------------------===//
20+
///
21+
/// Registration methods for MLIR dialects.
22+
///
23+
//===----------------------------------------------------------------------===//
24+
#ifndef MLIR_TENSORRT_INIT_ALL_DIALECTS
25+
#define MLIR_TENSORRT_INIT_ALL_DIALECTS
26+
27+
#include "mlir-executor/Executor/IR/Executor.h"
28+
#include "mlir-tensorrt-dialect/TensorRT/IR/TensorRTDialect.h"
29+
#include "mlir-tensorrt-dialect/TensorRT/Target/TensorRTEncodingImpl.h"
30+
#include "mlir-tensorrt/Backends/Host/HostBackend.h"
31+
#include "mlir-tensorrt/Backends/TensorRT/TensorRTBackend.h"
32+
#include "mlir-tensorrt/Dialect/CUDA/IR/CUDADialect.h"
33+
#include "mlir-tensorrt/Dialect/CUDA/Transforms/BufferizableOpInterfaceImpl.h"
34+
#include "mlir-tensorrt/Dialect/Plan/IR/Plan.h"
35+
#include "mlir-tensorrt/Dialect/StablehloExt/IR/StableHloExt.h"
36+
#include "mlir-tensorrt/Dialect/TensorRTRuntime/IR/TensorRTRuntime.h"
37+
#include "mlir-tensorrt/Dialect/TensorRTRuntime/Transforms/BufferizableOpInterfaceImpl.h"
38+
#include "mlir-tensorrt/Features.h"
39+
#include "mlir/Dialect/Affine/IR/AffineOps.h"
40+
#include "mlir/Dialect/Affine/IR/ValueBoundsOpInterfaceImpl.h"
41+
#include "mlir/Dialect/Arith/IR/ValueBoundsOpInterfaceImpl.h"
42+
#include "mlir/Dialect/Arith/Transforms/BufferDeallocationOpInterfaceImpl.h"
43+
#include "mlir/Dialect/Arith/Transforms/BufferViewFlowOpInterfaceImpl.h"
44+
#include "mlir/Dialect/Arith/Transforms/BufferizableOpInterfaceImpl.h"
45+
#include "mlir/Dialect/Async/IR/Async.h"
46+
#include "mlir/Dialect/Bufferization/IR/Bufferization.h"
47+
#include "mlir/Dialect/Bufferization/Transforms/FuncBufferizableOpInterfaceImpl.h"
48+
#include "mlir/Dialect/Complex/IR/Complex.h"
49+
#include "mlir/Dialect/ControlFlow/IR/ControlFlowOps.h"
50+
#include "mlir/Dialect/ControlFlow/Transforms/BufferDeallocationOpInterfaceImpl.h"
51+
#include "mlir/Dialect/ControlFlow/Transforms/BufferizableOpInterfaceImpl.h"
52+
#include "mlir/Dialect/DLTI/DLTI.h"
53+
#include "mlir/Dialect/EmitC/IR/EmitC.h"
54+
#include "mlir/Dialect/GPU/IR/GPUDialect.h"
55+
#include "mlir/Dialect/Index/IR/IndexDialect.h"
56+
#include "mlir/Dialect/LLVMIR/LLVMDialect.h"
57+
#include "mlir/Dialect/LLVMIR/NVVMDialect.h"
58+
#include "mlir/Dialect/LLVMIR/Transforms/InlinerInterfaceImpl.h"
59+
#include "mlir/Dialect/Linalg/IR/Linalg.h"
60+
#include "mlir/Dialect/Linalg/IR/ValueBoundsOpInterfaceImpl.h"
61+
#include "mlir/Dialect/Linalg/Transforms/AllInterfaces.h"
62+
#include "mlir/Dialect/Linalg/Transforms/BufferizableOpInterfaceImpl.h"
63+
#include "mlir/Dialect/Linalg/Transforms/RuntimeOpVerification.h"
64+
#include "mlir/Dialect/Linalg/Transforms/SubsetInsertionOpInterfaceImpl.h"
65+
#include "mlir/Dialect/Linalg/Transforms/TilingInterfaceImpl.h"
66+
#include "mlir/Dialect/Math/IR/Math.h"
67+
#include "mlir/Dialect/MemRef/IR/MemRef.h"
68+
#include "mlir/Dialect/MemRef/IR/MemRefMemorySlot.h"
69+
#include "mlir/Dialect/MemRef/IR/ValueBoundsOpInterfaceImpl.h"
70+
#include "mlir/Dialect/MemRef/Transforms/AllocationOpInterfaceImpl.h"
71+
#include "mlir/Dialect/MemRef/Transforms/BufferViewFlowOpInterfaceImpl.h"
72+
#include "mlir/Dialect/MemRef/Transforms/RuntimeOpVerification.h"
73+
#include "mlir/Dialect/PDL/IR/PDL.h"
74+
#include "mlir/Dialect/PDLInterp/IR/PDLInterp.h"
75+
#include "mlir/Dialect/Ptr/IR/PtrDialect.h"
76+
#include "mlir/Dialect/Quant/IR/Quant.h"
77+
#include "mlir/Dialect/SCF/IR/SCF.h"
78+
#include "mlir/Dialect/SCF/IR/ValueBoundsOpInterfaceImpl.h"
79+
#include "mlir/Dialect/SCF/Transforms/BufferDeallocationOpInterfaceImpl.h"
80+
#include "mlir/Dialect/SCF/Transforms/BufferizableOpInterfaceImpl.h"
81+
#include "mlir/Dialect/Shape/IR/Shape.h"
82+
#include "mlir/Dialect/Tensor/IR/Tensor.h"
83+
#include "mlir/Dialect/Tensor/IR/TensorInferTypeOpInterfaceImpl.h"
84+
#include "mlir/Dialect/Tensor/IR/TensorTilingInterfaceImpl.h"
85+
#include "mlir/Dialect/Tensor/IR/ValueBoundsOpInterfaceImpl.h"
86+
#include "mlir/Dialect/Tensor/TransformOps/TensorTransformOps.h"
87+
#include "mlir/Dialect/Tensor/Transforms/BufferizableOpInterfaceImpl.h"
88+
#include "mlir/Dialect/Tensor/Transforms/SubsetInsertionOpInterfaceImpl.h"
89+
#include "mlir/Dialect/Transform/IR/TransformDialect.h"
90+
#include "mlir/Dialect/UB/IR/UBOps.h"
91+
#include "mlir/Dialect/Vector/IR/ValueBoundsOpInterfaceImpl.h"
92+
#include "mlir/Dialect/Vector/IR/VectorOps.h"
93+
#include "mlir/Dialect/Vector/Transforms/BufferizableOpInterfaceImpl.h"
94+
#include "mlir/Dialect/Vector/Transforms/SubsetOpInterfaceImpl.h"
95+
#include "mlir/Interfaces/CastInterfaces.h"
96+
#include "mlir/Target/LLVM/NVVM/Target.h"
97+
98+
#ifdef MLIR_TRT_ENABLE_HLO
99+
#include "stablehlo/dialect/ChloOps.h"
100+
#include "stablehlo/dialect/StablehloOps.h"
101+
#include "stablehlo/dialect/VhloOps.h"
102+
#endif
103+
104+
namespace mlirtrt::compiler {
105+
106+
inline void registerAllDialects(mlir::DialectRegistry &registry) {
107+
// clang-format off
108+
registry.insert<
109+
mlir::affine::AffineDialect,
110+
mlir::arith::ArithDialect,
111+
mlir::async::AsyncDialect,
112+
mlir::bufferization::BufferizationDialect,
113+
mlir::cf::ControlFlowDialect,
114+
mlir::complex::ComplexDialect,
115+
mlir::cuda::CUDADialect,
116+
mlir::DLTIDialect,
117+
mlir::emitc::EmitCDialect,
118+
mlir::executor::ExecutorDialect,
119+
mlir::func::FuncDialect,
120+
mlir::gpu::GPUDialect,
121+
mlir::index::IndexDialect,
122+
mlir::linalg::LinalgDialect,
123+
mlir::LLVM::LLVMDialect,
124+
mlir::math::MathDialect,
125+
mlir::memref::MemRefDialect,
126+
mlir::NVVM::NVVMDialect,
127+
mlir::pdl_interp::PDLInterpDialect,
128+
mlir::pdl::PDLDialect,
129+
mlir::plan::PlanDialect,
130+
mlir::ptr::PtrDialect,
131+
mlir::quant::QuantDialect,
132+
mlir::scf::SCFDialect,
133+
mlir::shape::ShapeDialect,
134+
mlir::tensor::TensorDialect,
135+
mlir::tensorrt::TensorRTDialect,
136+
mlir::transform::TransformDialect,
137+
mlir::trtrt::TensorRTRuntimeDialect,
138+
mlir::ub::UBDialect,
139+
mlir::vector::VectorDialect
140+
>();
141+
// clang-format on
142+
143+
IF_MLIR_TRT_ENABLE_HLO({
144+
registry.insert<mlir::stablehlo::StablehloDialect>();
145+
registry.insert<mlir::chlo::ChloDialect>();
146+
registry.insert<mlir::vhlo::VhloDialect>();
147+
});
148+
149+
// Register all external models.
150+
mlir::affine::registerValueBoundsOpInterfaceExternalModels(registry);
151+
mlir::arith::registerBufferDeallocationOpInterfaceExternalModels(registry);
152+
mlir::arith::registerBufferizableOpInterfaceExternalModels(registry);
153+
mlir::arith::registerBufferViewFlowOpInterfaceExternalModels(registry);
154+
mlir::arith::registerValueBoundsOpInterfaceExternalModels(registry);
155+
mlir::bufferization::func_ext::registerBufferizableOpInterfaceExternalModels(
156+
registry);
157+
mlir::builtin::registerCastOpInterfaceExternalModels(registry);
158+
mlir::cf::registerBufferDeallocationOpInterfaceExternalModels(registry);
159+
mlir::cf::registerBufferizableOpInterfaceExternalModels(registry);
160+
mlir::cuda::registerBufferizableOpInterfaceExternalModels(registry);
161+
mlir::linalg::registerBufferizableOpInterfaceExternalModels(registry);
162+
mlir::linalg::registerRuntimeVerifiableOpInterfaceExternalModels(registry);
163+
mlir::linalg::registerSubsetOpInterfaceExternalModels(registry);
164+
mlir::linalg::registerTilingInterfaceExternalModels(registry);
165+
mlir::linalg::registerValueBoundsOpInterfaceExternalModels(registry);
166+
mlir::LLVM::registerInlinerInterface(registry);
167+
mlir::memref::registerAllocationOpInterfaceExternalModels(registry);
168+
mlir::memref::registerBufferViewFlowOpInterfaceExternalModels(registry);
169+
mlir::memref::registerMemorySlotExternalModels(registry);
170+
mlir::memref::registerRuntimeVerifiableOpInterfaceExternalModels(registry);
171+
mlir::memref::registerValueBoundsOpInterfaceExternalModels(registry);
172+
mlir::NVVM::registerInlinerInterface(registry);
173+
mlir::NVVM::registerNVVMTargetInterfaceExternalModels(registry);
174+
mlir::scf::registerBufferDeallocationOpInterfaceExternalModels(registry);
175+
mlir::scf::registerBufferizableOpInterfaceExternalModels(registry);
176+
mlir::scf::registerValueBoundsOpInterfaceExternalModels(registry);
177+
mlir::tensor::registerBufferizableOpInterfaceExternalModels(registry);
178+
mlir::tensor::registerFindPayloadReplacementOpInterfaceExternalModels(
179+
registry);
180+
mlir::tensor::registerInferTypeOpInterfaceExternalModels(registry);
181+
mlir::tensor::registerSubsetOpInterfaceExternalModels(registry);
182+
mlir::tensor::registerTilingInterfaceExternalModels(registry);
183+
mlir::tensor::registerValueBoundsOpInterfaceExternalModels(registry);
184+
mlir::tensorrt::registerTensorKindOpInterfaceExternalModels(registry);
185+
mlir::tensorrt::registerTensorRTEncodingOpInterfaceExternalModels(registry);
186+
mlir::trtrt::registerBufferizableOpInterfaceExternalModels(registry);
187+
mlir::vector::registerBufferizableOpInterfaceExternalModels(registry);
188+
mlir::vector::registerSubsetOpInterfaceExternalModels(registry);
189+
mlir::vector::registerValueBoundsOpInterfaceExternalModels(registry);
190+
191+
IF_MLIR_TRT_ENABLE_HLO({
192+
mlir::stablehlo::registerTensorKindOpInterfaceExternalModels(registry);
193+
mlir::stablehlo::registerTypeInferenceExternalModels(registry);
194+
});
195+
}
196+
197+
} // namespace mlirtrt::compiler
198+
199+
#endif // MLIR_TENSORRT_INIT_ALL_DIALECTS

0 commit comments

Comments
 (0)