Fix Attention addLayer, make cmake to work with TRT 10.14

Author: yizhuoz004

mlir-tensorrt/CMakePresets.json

@@ -100,6 +100,20 @@
         "MLIR_TRT_ENABLE_NCCL": "OFF",
         "MLIR_TRT_DOWNLOAD_TENSORRT_VERSION": "$env{DOWNLOAD_TENSORRT_VERSION}"
       }
+    },
+    {
+      "name": "python-wheel-build",
+      "displayName": "Configuration for building the compiler/runtime Python package wheels",
+      "generator": "Ninja",
+      "binaryDir": "build",
+      "inherits": "ninja-llvm",
+      "cacheVariables": {
+        "CMAKE_BUILD_TYPE": "Release",
+        "LLVM_ENABLE_ASSERTIONS": "OFF",
+        "CMAKE_PLATFORM_NO_VERSIONED_SONAME": "ON",
+        "MLIR_TRT_ENABLE_NCCL": "OFF",
+        "MLIR_TRT_DOWNLOAD_TENSORRT_VERSION": "$env{DOWNLOAD_TENSORRT_VERSION}"
+      }
     }
   ]
 }

mlir-tensorrt/build_tools/cmake/MTRTDependencies.cmake

@@ -57,8 +57,8 @@ macro(configure_tensorrt_python_plugin_header)
     find_file(
       trt_python_plugin_header
       NAMES NvInferPythonPlugin.h plugin.h
-      HINTS ${ARG_INSTALL_DIR} ${ARG_INSTALL_DIR}/python/include/impl
-      PATHS ${ARG_INSTALL_DIR} ${ARG_INSTALL_DIR}/python/include/impl
+      HINTS ${ARG_INSTALL_DIR} ${ARG_INSTALL_DIR}/python/include/impl ${ARG_INSTALL_DIR}/include/impl
+      PATHS ${ARG_INSTALL_DIR} ${ARG_INSTALL_DIR}/python/include/impl ${ARG_INSTALL_DIR}/include/impl
       REQUIRED
       NO_CMAKE_PATH NO_DEFAULT_PATH
       NO_CACHE

mlir-tensorrt/build_tools/cmake/TensorRTDownloadURL.cmake

@@ -80,6 +80,10 @@ function(mtrt_get_tensorrt_download_url ARG_VERSION OS_NAME TARGET_ARCH ARG_OUT_
     set(ARG_VERSION "10.12.0.36")
   endif()
 
+  if(ARG_VERSION VERSION_EQUAL "10.14")
+    set(ARG_VERSION "10.14.1.48")
+  endif()
+
   set(downloadable_versions
     "8.6.1.6"
     "9.0.1.4" "9.1.0.4" "9.2.0.5"
@@ -97,6 +101,7 @@ function(mtrt_get_tensorrt_download_url ARG_VERSION OS_NAME TARGET_ARCH ARG_OUT_
     "10.8.0.43"
     "10.9.0.34"
     "10.12.0.36"
+    "10.14.1.48"
   )
 
   if(NOT ARG_VERSION IN_LIST downloadable_versions)
@@ -164,6 +169,8 @@ function(mtrt_get_tensorrt_download_url ARG_VERSION OS_NAME TARGET_ARCH ARG_OUT_
   elseif(ARG_VERSION VERSION_GREATER 10.10
       AND ARG_VERSION VERSION_LESS 10.13)
     set(TRT_CUDA_VERSION 12.9)
+  elseif(ARG_VERSION VERSION_GREATER 10.13)
+    set(TRT_CUDA_VERSION 13.0)
   endif()
 
   # Handle TRT 8 versions.

mlir-tensorrt/compiler/test/python/mlir_tensorrt_compiler/dialects/test_tensorrt.py

@@ -49,6 +49,8 @@ def test_attributes():
             tensorrt.TripLimitAttr.get("kWHILE"),
             tensorrt.FillOperationAttr.get("kRANDOM_UNIFORM"),
             tensorrt.ScatterModeAttr.get("kELEMENT"),
+            tensorrt.AttentionNormalizationOpAttr.get("kSOFTMAX"),
+            tensorrt.DataTypeAttr.get("kFLOAT"),
         ]:
             print(attr)
 
@@ -74,3 +76,5 @@ def test_attributes():
 # CHECK-NEXT: #tensorrt.trip_limit<kWHILE>
 # CHECK-NEXT: #tensorrt.fill_operation<kRANDOM_UNIFORM>
 # CHECK-NEXT: #tensorrt.scatter_mode<kELEMENT>
+# CHECK-NEXT: #tensorrt.attention_normalization_op<kSOFTMAX>
+# CHECK-NEXT: #tensorrt.data_type<kFLOAT>

mlir-tensorrt/compiler/tools/CMakeLists.txt

@@ -21,5 +21,5 @@ set(LLVM_LINK_COMPONENTS
 add_subdirectory(mlir-tensorrt-opt)
 add_subdirectory(mlir-tensorrt-compiler)
 add_subdirectory(mlir-tensorrt-translate)
-add_subdirectory(mlir-tensorrt-lsp-server)
+# add_subdirectory(mlir-tensorrt-lsp-server)
 add_subdirectory(mlir-tensorrt-runner)

mlir-tensorrt/integrations/python/setup_utils.py

@@ -13,7 +13,7 @@
 import subprocess
 import atexit
 
-TENSORRT_VERSION = os.getenv("MLIR_TRT_DOWNLOAD_TENSORRT_VERSION", "10.12")
+TENSORRT_VERSION = os.getenv("MLIR_TRT_DOWNLOAD_TENSORRT_VERSION", "10.14")
 
 
 def log(*args):
@@ -105,8 +105,8 @@ def run_cmake_build(python_package_name: str, python_wheel_staging_dir: Path):
 
     # Environment variable overrides
     cmake_preset = os.environ.get("MLIR_TRT_CMAKE_PRESET", "python-wheel-build")
-    install_prefix = os.environ.get("MLIR_TRT_INSTALL_DIR", None)
-    build_dir = os.environ.get("MLIR_TRT_BUILD_DIR", None)
+    install_prefix = os.environ.get("MLIR_TRT_INSTALL_DIR", "./install")
+    build_dir = os.environ.get("MLIR_TRT_BUILD_DIR", "./build")
     parallel_jobs = os.environ.get("MLIR_TRT_PARALLEL_JOBS", str(os.cpu_count() or 1))
 
     # Additional CMake options from environment

mlir-tensorrt/tensorrt/include/mlir-tensorrt-dialect-c/TensorRTAttributes.h

@@ -188,6 +188,22 @@ DECLARE_ATTR_GETTER_FROM_STRING(ScatterMode)
 DECLARE_IS_ATTR(ScatterMode)
 DECLARE_STRING_GETTER_FROM_ATTR(ScatterMode)
 
+//===----------------------------------------------------------------------===//
+// AttentionNormalizationOp
+//===----------------------------------------------------------------------===//
+
+DECLARE_ATTR_GETTER_FROM_STRING(AttentionNormalizationOp)
+DECLARE_IS_ATTR(AttentionNormalizationOp)
+DECLARE_STRING_GETTER_FROM_ATTR(AttentionNormalizationOp)
+
+//===----------------------------------------------------------------------===//
+// DataType
+//===----------------------------------------------------------------------===//
+
+DECLARE_ATTR_GETTER_FROM_STRING(DataType)
+DECLARE_IS_ATTR(DataType)
+DECLARE_STRING_GETTER_FROM_ATTR(DataType)
+
 #ifdef __cplusplus
 }
 #endif