[mlir-tensorrt] Integrate internal changes

mlir-tensorrt/.gitignore

@@ -3,6 +3,7 @@
 *.log
 **/llvm-project/**
 **/llvm-project/
+CMakeUserPresets.json
 
 # Docs build artifacts
 /public/

mlir-tensorrt/CMakeLists.txt

@@ -90,6 +90,35 @@ if(PROJECT_IS_TOP_LEVEL)
   set(CMAKE_EXPORT_COMPILE_COMMANDS ON)
 endif()
 
+# -------------------------------------------------
+# Option validation
+# -------------------------------------------------
+# Write out a header file containing convenience macros for each flag.
+function(mtrt_write_feature_flags_header)
+  set(feature_flags_header
+    "${CMAKE_CURRENT_BINARY_DIR}/include/mlir-tensorrt/Features.h")
+
+  # Generate the header at configure time
+  file(WRITE "${feature_flags_header}" [[
+  // Auto-generated feature macros, do not edit.
+  #ifndef MLIR_TENSORRT_FEATURES_H
+  #define MLIR_TENSORRT_FEATURES_H
+
+  ]])
+
+  foreach(FEATURE IN LISTS MLIR_TRT_FEATURE_FLAGS)
+      file(APPEND "${feature_flags_header}" "#ifdef ${FEATURE}\n")
+      file(APPEND "${feature_flags_header}" "#define IF_${FEATURE}(code) do { code } while (0)\n")
+      file(APPEND "${feature_flags_header}" "#else\n")
+      file(APPEND "${feature_flags_header}" "#define IF_${FEATURE}(code) do {} while (0)\n")
+      file(APPEND "${feature_flags_header}" "#endif // ${FEATURE}\n\n")
+  endforeach()
+  file(APPEND "${feature_flags_header}" "#endif // MLIR_TENSORRT_FEATURES_H\n")
+endfunction()
+
+mtrt_write_feature_flags_header()
+include_directories("${CMAKE_CURRENT_BINARY_DIR}/include")
+
 # -------------------------------------------------
 # Setup LLVM/MLIR
 # -------------------------------------------------

mlir-tensorrt/build_tools/cmake/Targets.cmake

@@ -44,6 +44,16 @@ function(add_mlir_tensorrt_public_c_api_library target)
   endif()
 endfunction()
 
+# --------------------------------------------------------------
+# Adds an upstream MLIR library target to the
+# MLIR_TENSORRT_LIBS global property list to capture it as an
+# implicit dependency for all final tools and compiler
+# end-user products.
+# --------------------------------------------------------------
+function(add_mlir_tensorrt_compiler_dependency target)
+  set_property(GLOBAL APPEND PROPERTY MLIR_TENSORRT_LIBS ${target})
+endfunction()
+
 # ------------------------------------------------------------------------------
 # A wrapper around `add_mlir_dialect_library` that also appends the dialect
 # library to the global `MLIR_TENSORRT_DIALECT_LIBS` list property.
@@ -99,11 +109,11 @@ function(add_mlir_tensorrt_backend_library target)
     BASE_DIRECTORY ${CMAKE_CURRENT_BINARY_DIR})
 
   set(LLVM_TARGET_DEFINITIONS "${SRC_TD}")
-  
+
   string(REPLACE ".td" "Attrs.h.inc" h_inc_file ${BIN_TD})
   string(REPLACE ".td" "Attrs.cpp.inc" cpp_inc_file ${BIN_TD})
   mlir_tablegen("${h_inc_file}" -gen-attrdef-decls)
-  mlir_tablegen("${cpp_inc_file}" -gen-attrdef-defs)  
+  mlir_tablegen("${cpp_inc_file}" -gen-attrdef-defs)
 
   add_public_tablegen_target(${target}IncGen)
 

mlir-tensorrt/compiler/CMakeLists.txt

@@ -9,6 +9,29 @@ set(MLIR_TENSORRT_COMPILER_INCLUDE_DIRS
 include_directories("${CMAKE_CURRENT_BINARY_DIR}/include"
                     "${CMAKE_CURRENT_SOURCE_DIR}/include")
 
+# We use the MLIR_TENSORRT_LIBS global property to aggregate the list of all
+# declared compiler libraries. This is helpful for simplifying the link
+# dependency declarations for tools that must link "the world" like
+# `mlir-tensorrt-opt`.
+
+# Because MLIR has a level of indirection that lets implementation for
+# interfaces be provided by separate implementation code
+# ("PromisedInterfaces/ExternalModels") which is registered at runtime, it is
+# difficult to capture all dependencies for dialects we require in the compiler
+# purely through target dependency properties. To see what we require from
+# usptream, look at the file `mlir-tensorrt/InitAllDialects.h`. Therefore, we
+# manually enumerate some dependencies here, mainly for providing the functions
+# registering interface external models.
+add_mlir_tensorrt_compiler_dependency(MLIRArithTransforms)
+add_mlir_tensorrt_compiler_dependency(MLIRArithValueBoundsOpInterfaceImpl)
+add_mlir_tensorrt_compiler_dependency(MLIRAsyncDialect)
+add_mlir_tensorrt_compiler_dependency(MLIRBufferizationTransforms)
+add_mlir_tensorrt_compiler_dependency(MLIRControlFlowTransforms)
+add_mlir_tensorrt_compiler_dependency(MLIRNVVMTarget)
+add_mlir_tensorrt_compiler_dependency(MLIRPtrDialect)
+add_mlir_tensorrt_compiler_dependency(MLIRTargetLLVM)
+add_mlir_tensorrt_compiler_dependency(MLIRTensorTransformOps)
+
 add_subdirectory(include)
 add_subdirectory(lib)
 add_subdirectory(test)

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

@@ -113,6 +113,20 @@ struct DebugOptions : public OptionsProvider {
                      "tree rooted at this directory. Use in conjunction with "
                      "mlir-print-ir-* flags")};
 
+  //===----------------------------------------------------------------------===//
+  // Printing Flags
+  //===----------------------------------------------------------------------===//
+
+  Option<unsigned> elideElementsAttrIfLarger{
+      this->ctx, "mlir-elide-elementsattrs-if-larger",
+      llvm::cl::desc("Elide ElementsAttrs with \"...\" that have "
+                     "more elements than the given upper limit")};
+
+  Option<unsigned> elideResourceStringsIfLarger{
+      this->ctx, "mlir-elide-resource-strings-if-larger",
+      llvm::cl::desc(
+          "Elide printing value of resources if string is too long in chars.")};
+
   //===--------------------------------------------------------------------===//
   // Pass Statistics
   //===--------------------------------------------------------------------===//

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

@@ -6,9 +6,6 @@ set(_TABLEGEN_ARGS )
 if(MLIR_TRT_ENABLE_HLO)
   list(APPEND _TABLEGEN_ARGS -DMLIR_TENSORRT_ENABLE_HLO)
 endif()
-if(MLIR_TRT_ENABLE_EXECUTOR)
-  list(APPEND _TABLEGEN_ARGS -DMLIR_TENSORRT_ENABLE_EXECUTOR)
-endif()
 
 mlir_tablegen(Passes.h.inc -gen-pass-decls -name MLIRTensorRTConversion ${_TABLEGEN_ARGS})
 add_public_tablegen_target(MLIRTensorRTConversionPassIncGen)

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

@@ -125,8 +125,6 @@ def ConvertTensorRTToEmitCPass : Pass<"convert-tensorrt-to-emitc",
   let dependentDialects = ["::mlir::emitc::EmitCDialect"];
 }
 
-#ifdef MLIR_TENSORRT_ENABLE_EXECUTOR
-
 //===----------------------------------------------------------------------===//
 // ConvertMemRefToCUDAPass
 //===----------------------------------------------------------------------===//
@@ -312,9 +310,6 @@ def ConvertTensorRTRuntimeToExecutorPass : Pass<"convert-tensorrt-runtime-to-exe
   let options = ConvertToExecutorOptions;
 }
 
-
-#endif // MLIR_TENSORRT_ENABLE_EXECUTOR
-
 #ifdef MLIR_TENSORRT_ENABLE_HLO
 //===----------------------------------------------------------------------===//
 // ConvertStablehloToScfPass

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

@@ -1,10 +1,7 @@
 if(MLIR_TRT_ENABLE_HLO)
   add_subdirectory(StablehloExt)
-  add_subdirectory(Plan)
 endif()
 
+add_subdirectory(CUDA)
+add_subdirectory(Plan)
 add_subdirectory(TensorRTRuntime)
-
-if(MLIR_TRT_ENABLE_EXECUTOR)
-  add_subdirectory(CUDA)
-endif()

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

@@ -428,6 +428,27 @@ def PostClusteringValidationPass : Pass<"post-clustering-validation", "func::Fun
   }];
 }
 
+//===----------------------------------------------------------------------===//
+// PlanAssignMemorySpacesPass
+//===----------------------------------------------------------------------===//
+
+def PlanAssignMemorySpacesPass : Pass<"plan-assign-memory-spaces",
+                                    "::mlir::ModuleOp"> {
+  let summary = "assigns memory spaces encodings to tensor types";
+
+  let description = [{
+    This pass applies a type conversion that adds a '#plan.memory_space'
+    attribute to all tensor types in the top-level module that do not already
+    have an encoding.
+  }];
+
+  let dependentDialects = [
+    "::mlir::plan::PlanDialect",
+    "::mlir::bufferization::BufferizationDialect",
+    "::mlir::tensor::TensorDialect"
+  ];
+}
+
 //===----------------------------------------------------------------------===//
 // PlanAllocTensorsPass
 //===----------------------------------------------------------------------===//

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

@@ -0,0 +1,199 @@
+//===- InitAllDialects.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.
+//
+//===----------------------------------------------------------------------===//
+///
+/// Registration methods for MLIR dialects.
+///
+//===----------------------------------------------------------------------===//
+#ifndef MLIR_TENSORRT_INIT_ALL_DIALECTS
+#define MLIR_TENSORRT_INIT_ALL_DIALECTS
+
+#include "mlir-executor/Executor/IR/Executor.h"
+#include "mlir-tensorrt-dialect/TensorRT/IR/TensorRTDialect.h"
+#include "mlir-tensorrt-dialect/TensorRT/Target/TensorRTEncodingImpl.h"
+#include "mlir-tensorrt/Backends/Host/HostBackend.h"
+#include "mlir-tensorrt/Backends/TensorRT/TensorRTBackend.h"
+#include "mlir-tensorrt/Dialect/CUDA/IR/CUDADialect.h"
+#include "mlir-tensorrt/Dialect/CUDA/Transforms/BufferizableOpInterfaceImpl.h"
+#include "mlir-tensorrt/Dialect/Plan/IR/Plan.h"
+#include "mlir-tensorrt/Dialect/StablehloExt/IR/StableHloExt.h"
+#include "mlir-tensorrt/Dialect/TensorRTRuntime/IR/TensorRTRuntime.h"
+#include "mlir-tensorrt/Dialect/TensorRTRuntime/Transforms/BufferizableOpInterfaceImpl.h"
+#include "mlir-tensorrt/Features.h"
+#include "mlir/Dialect/Affine/IR/AffineOps.h"
+#include "mlir/Dialect/Affine/IR/ValueBoundsOpInterfaceImpl.h"
+#include "mlir/Dialect/Arith/IR/ValueBoundsOpInterfaceImpl.h"
+#include "mlir/Dialect/Arith/Transforms/BufferDeallocationOpInterfaceImpl.h"
+#include "mlir/Dialect/Arith/Transforms/BufferViewFlowOpInterfaceImpl.h"
+#include "mlir/Dialect/Arith/Transforms/BufferizableOpInterfaceImpl.h"
+#include "mlir/Dialect/Async/IR/Async.h"
+#include "mlir/Dialect/Bufferization/IR/Bufferization.h"
+#include "mlir/Dialect/Bufferization/Transforms/FuncBufferizableOpInterfaceImpl.h"
+#include "mlir/Dialect/Complex/IR/Complex.h"
+#include "mlir/Dialect/ControlFlow/IR/ControlFlowOps.h"
+#include "mlir/Dialect/ControlFlow/Transforms/BufferDeallocationOpInterfaceImpl.h"
+#include "mlir/Dialect/ControlFlow/Transforms/BufferizableOpInterfaceImpl.h"
+#include "mlir/Dialect/DLTI/DLTI.h"
+#include "mlir/Dialect/EmitC/IR/EmitC.h"
+#include "mlir/Dialect/GPU/IR/GPUDialect.h"
+#include "mlir/Dialect/Index/IR/IndexDialect.h"
+#include "mlir/Dialect/LLVMIR/LLVMDialect.h"
+#include "mlir/Dialect/LLVMIR/NVVMDialect.h"
+#include "mlir/Dialect/LLVMIR/Transforms/InlinerInterfaceImpl.h"
+#include "mlir/Dialect/Linalg/IR/Linalg.h"
+#include "mlir/Dialect/Linalg/IR/ValueBoundsOpInterfaceImpl.h"
+#include "mlir/Dialect/Linalg/Transforms/AllInterfaces.h"
+#include "mlir/Dialect/Linalg/Transforms/BufferizableOpInterfaceImpl.h"
+#include "mlir/Dialect/Linalg/Transforms/RuntimeOpVerification.h"
+#include "mlir/Dialect/Linalg/Transforms/SubsetInsertionOpInterfaceImpl.h"
+#include "mlir/Dialect/Linalg/Transforms/TilingInterfaceImpl.h"
+#include "mlir/Dialect/Math/IR/Math.h"
+#include "mlir/Dialect/MemRef/IR/MemRef.h"
+#include "mlir/Dialect/MemRef/IR/MemRefMemorySlot.h"
+#include "mlir/Dialect/MemRef/IR/ValueBoundsOpInterfaceImpl.h"
+#include "mlir/Dialect/MemRef/Transforms/AllocationOpInterfaceImpl.h"
+#include "mlir/Dialect/MemRef/Transforms/BufferViewFlowOpInterfaceImpl.h"
+#include "mlir/Dialect/MemRef/Transforms/RuntimeOpVerification.h"
+#include "mlir/Dialect/PDL/IR/PDL.h"
+#include "mlir/Dialect/PDLInterp/IR/PDLInterp.h"
+#include "mlir/Dialect/Ptr/IR/PtrDialect.h"
+#include "mlir/Dialect/Quant/IR/Quant.h"
+#include "mlir/Dialect/SCF/IR/SCF.h"
+#include "mlir/Dialect/SCF/IR/ValueBoundsOpInterfaceImpl.h"
+#include "mlir/Dialect/SCF/Transforms/BufferDeallocationOpInterfaceImpl.h"
+#include "mlir/Dialect/SCF/Transforms/BufferizableOpInterfaceImpl.h"
+#include "mlir/Dialect/Shape/IR/Shape.h"
+#include "mlir/Dialect/Tensor/IR/Tensor.h"
+#include "mlir/Dialect/Tensor/IR/TensorInferTypeOpInterfaceImpl.h"
+#include "mlir/Dialect/Tensor/IR/TensorTilingInterfaceImpl.h"
+#include "mlir/Dialect/Tensor/IR/ValueBoundsOpInterfaceImpl.h"
+#include "mlir/Dialect/Tensor/TransformOps/TensorTransformOps.h"
+#include "mlir/Dialect/Tensor/Transforms/BufferizableOpInterfaceImpl.h"
+#include "mlir/Dialect/Tensor/Transforms/SubsetInsertionOpInterfaceImpl.h"
+#include "mlir/Dialect/Transform/IR/TransformDialect.h"
+#include "mlir/Dialect/UB/IR/UBOps.h"
+#include "mlir/Dialect/Vector/IR/ValueBoundsOpInterfaceImpl.h"
+#include "mlir/Dialect/Vector/IR/VectorOps.h"
+#include "mlir/Dialect/Vector/Transforms/BufferizableOpInterfaceImpl.h"
+#include "mlir/Dialect/Vector/Transforms/SubsetOpInterfaceImpl.h"
+#include "mlir/Interfaces/CastInterfaces.h"
+#include "mlir/Target/LLVM/NVVM/Target.h"
+
+#ifdef MLIR_TRT_ENABLE_HLO
+#include "stablehlo/dialect/ChloOps.h"
+#include "stablehlo/dialect/StablehloOps.h"
+#include "stablehlo/dialect/VhloOps.h"
+#endif
+
+namespace mlirtrt::compiler {
+
+inline void registerAllDialects(mlir::DialectRegistry &registry) {
+  // clang-format off
+  registry.insert<
+      mlir::affine::AffineDialect,
+      mlir::arith::ArithDialect,
+      mlir::async::AsyncDialect,
+      mlir::bufferization::BufferizationDialect,
+      mlir::cf::ControlFlowDialect,
+      mlir::complex::ComplexDialect,
+      mlir::cuda::CUDADialect,
+      mlir::DLTIDialect,
+      mlir::emitc::EmitCDialect,
+      mlir::executor::ExecutorDialect,
+      mlir::func::FuncDialect,
+      mlir::gpu::GPUDialect,
+      mlir::index::IndexDialect,
+      mlir::linalg::LinalgDialect,
+      mlir::LLVM::LLVMDialect,
+      mlir::math::MathDialect,
+      mlir::memref::MemRefDialect,
+      mlir::NVVM::NVVMDialect,
+      mlir::pdl_interp::PDLInterpDialect,
+      mlir::pdl::PDLDialect,
+      mlir::plan::PlanDialect,
+      mlir::ptr::PtrDialect,
+      mlir::quant::QuantDialect,
+      mlir::scf::SCFDialect,
+      mlir::shape::ShapeDialect,
+      mlir::tensor::TensorDialect,
+      mlir::tensorrt::TensorRTDialect,
+      mlir::transform::TransformDialect,
+      mlir::trtrt::TensorRTRuntimeDialect,
+      mlir::ub::UBDialect,
+      mlir::vector::VectorDialect
+    >();
+  // clang-format on
+
+  IF_MLIR_TRT_ENABLE_HLO({
+    registry.insert<mlir::stablehlo::StablehloDialect>();
+    registry.insert<mlir::chlo::ChloDialect>();
+    registry.insert<mlir::vhlo::VhloDialect>();
+  });
+
+  // Register all external models.
+  mlir::affine::registerValueBoundsOpInterfaceExternalModels(registry);
+  mlir::arith::registerBufferDeallocationOpInterfaceExternalModels(registry);
+  mlir::arith::registerBufferizableOpInterfaceExternalModels(registry);
+  mlir::arith::registerBufferViewFlowOpInterfaceExternalModels(registry);
+  mlir::arith::registerValueBoundsOpInterfaceExternalModels(registry);
+  mlir::bufferization::func_ext::registerBufferizableOpInterfaceExternalModels(
+      registry);
+  mlir::builtin::registerCastOpInterfaceExternalModels(registry);
+  mlir::cf::registerBufferDeallocationOpInterfaceExternalModels(registry);
+  mlir::cf::registerBufferizableOpInterfaceExternalModels(registry);
+  mlir::cuda::registerBufferizableOpInterfaceExternalModels(registry);
+  mlir::linalg::registerBufferizableOpInterfaceExternalModels(registry);
+  mlir::linalg::registerRuntimeVerifiableOpInterfaceExternalModels(registry);
+  mlir::linalg::registerSubsetOpInterfaceExternalModels(registry);
+  mlir::linalg::registerTilingInterfaceExternalModels(registry);
+  mlir::linalg::registerValueBoundsOpInterfaceExternalModels(registry);
+  mlir::LLVM::registerInlinerInterface(registry);
+  mlir::memref::registerAllocationOpInterfaceExternalModels(registry);
+  mlir::memref::registerBufferViewFlowOpInterfaceExternalModels(registry);
+  mlir::memref::registerMemorySlotExternalModels(registry);
+  mlir::memref::registerRuntimeVerifiableOpInterfaceExternalModels(registry);
+  mlir::memref::registerValueBoundsOpInterfaceExternalModels(registry);
+  mlir::NVVM::registerInlinerInterface(registry);
+  mlir::NVVM::registerNVVMTargetInterfaceExternalModels(registry);
+  mlir::scf::registerBufferDeallocationOpInterfaceExternalModels(registry);
+  mlir::scf::registerBufferizableOpInterfaceExternalModels(registry);
+  mlir::scf::registerValueBoundsOpInterfaceExternalModels(registry);
+  mlir::tensor::registerBufferizableOpInterfaceExternalModels(registry);
+  mlir::tensor::registerFindPayloadReplacementOpInterfaceExternalModels(
+      registry);
+  mlir::tensor::registerInferTypeOpInterfaceExternalModels(registry);
+  mlir::tensor::registerSubsetOpInterfaceExternalModels(registry);
+  mlir::tensor::registerTilingInterfaceExternalModels(registry);
+  mlir::tensor::registerValueBoundsOpInterfaceExternalModels(registry);
+  mlir::tensorrt::registerTensorKindOpInterfaceExternalModels(registry);
+  mlir::tensorrt::registerTensorRTEncodingOpInterfaceExternalModels(registry);
+  mlir::trtrt::registerBufferizableOpInterfaceExternalModels(registry);
+  mlir::vector::registerBufferizableOpInterfaceExternalModels(registry);
+  mlir::vector::registerSubsetOpInterfaceExternalModels(registry);
+  mlir::vector::registerValueBoundsOpInterfaceExternalModels(registry);
+
+  IF_MLIR_TRT_ENABLE_HLO({
+    mlir::stablehlo::registerTensorKindOpInterfaceExternalModels(registry);
+    mlir::stablehlo::registerTypeInferenceExternalModels(registry);
+  });
+}
+
+} // namespace mlirtrt::compiler
+
+#endif // MLIR_TENSORRT_INIT_ALL_DIALECTS