bazel test //tensorrt/test:tensorrt_mlir_dialect_tests

NVIDIA · Oct 21, 2024 · 2bf0785 · 2bf0785
1 parent 9f12bef
commit 2bf0785
Show file tree

Hide file tree

Showing 8 changed files with 82 additions and 8 deletions.
diff --git a/mlir-tensorrt/BUILD b/mlir-tensorrt/BUILD
@@ -691,6 +691,19 @@ cc_binary(
     ],
 )
 
+cc_binary(
+    name = "tensorrt-opt",
+    srcs = ["tensorrt/tensorrt-opt/tensorrt-opt.cpp"],
+    copts = [
+        "-DMLIR_TRT_TARGET_TENSORRT",
+    ],
+    linkopts = ["-Wl,-rpath,/opt/src/mlir-tensorrt/bazel-mlir-tensorrt/external/tensorrt10_x86/targets/x86_64-linux-gnu/lib"],
+    deps = [
+        ":TensorRTRegistration",
+        "@llvm-project//mlir:MlirOptLib",
+    ],
+)
+
 cc_binary(
     name = "mlir-tensorrt-opt",
     srcs = ["tools/MlirTensorRtOpt.cpp"],

diff --git a/mlir-tensorrt/WORKSPACE b/mlir-tensorrt/WORKSPACE
@@ -30,7 +30,6 @@ rules_cuda_dependencies()
 
 register_detected_cuda_toolchains()
 
-
 # -------------------------- #
 #    Hermetic Python Setup   #
 # -------------------------- #

diff --git a/mlir-tensorrt/tensorrt/tensorrt-opt/tensorrt-opt.cpp b/mlir-tensorrt/tensorrt/tensorrt-opt/tensorrt-opt.cpp
@@ -33,17 +33,21 @@
 #include "mlir/Tools/mlir-opt/MlirOptMain.h"
 #include "mlir/Transforms/Passes.h"
 
+#ifdef MLIR_TRT_ENABLE_TESTING
 namespace mlir {
 void registerTestTensorKindAnalysisPass();
 }
+#endif // MLIR_TRT_ENABLE_TESTING
 
 int main(int argc, char **argv) {
   mlir::DialectRegistry registry;
   registry.insert<mlir::tensorrt::TensorRTDialect, mlir::func::FuncDialect,
                   mlir::tensor::TensorDialect, mlir::arith::ArithDialect,
                   mlir::affine::AffineDialect, mlir::quant::QuantizationDialect,
                   mlir::scf::SCFDialect>();
+#ifdef MLIR_TRT_ENABLE_TESTING
   mlir::registerTestTensorKindAnalysisPass();
+#endif // MLIR_TRT_ENABLE_TESTING
   mlir::func::registerInlinerExtension(registry);
   mlir::tensorrt::registerTensorRTTranslationCLOpts();
   mlir::tensorrt::registerTensorRTPasses();

diff --git a/mlir-tensorrt/tensorrt/test/BUILD b/mlir-tensorrt/tensorrt/test/BUILD
@@ -0,0 +1,61 @@
+# Licensed under the Apache License v2.0 with LLVM Exceptions.
+# See https://llvm.org/LICENSE.txt for license information.
+# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+# Also available under a BSD-style license. See LICENSE.
+
+load("@bazel_skylib//rules:expand_template.bzl", "expand_template")
+load("@llvm-project//llvm:lit_test.bzl", "lit_test")
+
+expand_template(
+    name = "lit_site_cfg_py",
+    testonly = True,
+    out = "lit.site.cfg.py",
+    substitutions = {
+        "@LIT_SITE_CFG_IN_HEADER@": "# Autogenerated, do not edit.",
+        "\"@LLVM_TOOLS_DIR@\"": "os.path.join(os.environ['TEST_SRCDIR'], 'llvm-project', 'llvm')",
+        "\"@MLIR_TENSORRT_DIALECT_BINARY_DIR@\"": "os.path.join(os.environ['TEST_SRCDIR'], 'tensorrt-mlir')",
+        "\"@MLIR_TENSORRT_DIALECT_SOURCE_DIR@\"": "os.path.join(os.environ['TEST_SRCDIR'], 'tensorrt-mlir', 'tensorrt')",
+        "@MLIR_ENABLE_BINDINGS_PYTHON@": "False",
+        "@ENABLE_ASAN@": "False",
+    },
+    template = "lit.site.cfg.py.in",
+)
+
+# Common data used by most lit tests.
+filegroup(
+    name = "lit_data",
+    testonly = True,
+    data = [
+        "lit.cfg.py",
+        "lit.site.cfg.py",
+        "@llvm-project//llvm:FileCheck",
+        "@llvm-project//llvm:count",
+        "@llvm-project//llvm:not",
+    ],
+)
+
+[
+    lit_test(
+        name = "%s.test" % src,
+        srcs = [src],
+        data = [
+            "//:tensorrt-opt",
+            "//tensorrt/test:lit_data",
+        ],
+        tags = ["tensorrt_mlir_dialect_tests"],
+    )
+    for src in glob(
+        [
+            "TensorRT/*.mlir",
+        ],
+        exclude = [
+            # exclude MLIR_TRT_ENABLE_TESTING related tests
+            "TensorRT/tensor-kind-analysis.mlir",
+        ],
+    )
+]
+
+test_suite(
+    name = "tensorrt_mlir_dialect_tests",
+    tags = ["tensorrt_mlir_dialect_tests"],
+)
diff --git a/mlir-tensorrt/tensorrt/test/lit.cfg.py b/mlir-tensorrt/tensorrt/test/lit.cfg.py
@@ -40,9 +40,7 @@
 # subdirectories contain auxiliary inputs for various tests in their parent
 # directories.
 config.excludes = ["Inputs", "Examples", "CMakeLists.txt", "README.txt", "LICENSE.txt"]
-config.tensorrt_dialect_tools_dir = os.path.join(
-    config.tensorrt_dialect_obj_root, "bin"
-)
+config.tensorrt_dialect_tools_dir = os.path.join(config.tensorrt_dialect_obj_root)
 config.tensorrt_dialect_libs_dir = os.path.join(config.tensorrt_dialect_obj_root, "lib")
 config.substitutions.append(
     ("%tensorrt_dialect_libs", config.tensorrt_dialect_libs_dir)

diff --git a/mlir-tensorrt/tensorrt/test/lit.site.cfg.py.in b/mlir-tensorrt/tensorrt/test/lit.site.cfg.py.in
@@ -12,4 +12,4 @@ import lit.llvm
 lit.llvm.initialize(lit_config, config)
 
 # Let the main config do the real work.
-lit_config.load_config(config, "@MLIR_TENSORRT_DIALECT_SOURCE_DIR@/test/lit.cfg.py")
+lit_config.load_config(config, "@MLIR_TENSORRT_DIALECT_SOURCE_DIR@" + "/test/lit.cfg.py")
diff --git a/mlir-tensorrt/test/BUILD b/mlir-tensorrt/test/BUILD
@@ -39,9 +39,9 @@ filegroup(
     name = "lit_data",
     testonly = True,
     data = [
+        "gpu_tools.py",
         "lit.cfg.py",
         "lit.site.cfg.py",
-        "gpu_tools.py",
         "@llvm-project//llvm:FileCheck",
         "@llvm-project//llvm:count",
         "@llvm-project//llvm:not",
@@ -61,7 +61,6 @@ filegroup(
             requirement("nvidia-ml-py"),
             requirement("click"),
             requirement("numpy"),
-
         ],
     )
     for src in glob([

diff --git a/mlir-tensorrt/version.bzl b/mlir-tensorrt/version.bzl
@@ -3,4 +3,4 @@
 MLIR_TENSORRT_VERSION_MAJOR = "0"
 MLIR_TENSORRT_VERSION_MINOR = "1"
 MLIR_TENSORRT_VERSION_PATCH = "34"
-MLIR_TENSORRT_VERSION = "{}.{}.{}".format(MLIR_TENSORRT_VERSION_MAJOR, MLIR_TENSORRT_VERSION_MINOR, MLIR_TENSORRT_VERSION_PATCH)
+MLIR_TENSORRT_VERSION = "{}.{}.{}".format(MLIR_TENSORRT_VERSION_MAJOR, MLIR_TENSORRT_VERSION_MINOR, MLIR_TENSORRT_VERSION_PATCH)