Internal testing framework change.

PiperOrigin-RevId: 684605290
openxla · Oct 10, 2024 · 5f76470 · 5f76470
1 parent 704b1d7
commit 5f76470
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 1 deletion.
diff --git a/xla/service/gpu/fusions/triton/BUILD b/xla/service/gpu/fusions/triton/BUILD
@@ -392,6 +392,12 @@ cc_library(
 xla_test(
     name = "triton_fusion_emitter_device_legacy_test",
     srcs = if_gpu_is_configured(["triton_fusion_emitter_device_legacy_test.cc"]),
+    backend_args = if_google(
+        {
+            "gpu_h100": ["--heap_check="],
+        },
+        {},
+    ),
     backends = [
         "gpu_a100",
         "gpu_h100",

diff --git a/xla/service/gpu/model/BUILD b/xla/service/gpu/model/BUILD
@@ -5,7 +5,7 @@ load("//xla:xla.bzl", "xla_cc_test")
 
 # Libraries for performance modeling of HLO.
 load("//xla/tests:build_defs.bzl", "xla_test")
-load("//xla/tsl:tsl.bzl", "internal_visibility")
+load("//xla/tsl:tsl.bzl", "if_google", "internal_visibility")
 load("//xla/tsl:tsl.default.bzl", "get_compatible_with_portable")
 
 package(
@@ -963,6 +963,13 @@ cc_library(
 xla_test(
     name = "hlo_op_profiler_test",
     srcs = ["hlo_op_profiler_test.cc"],
+    # TODO(b/372714955): Fix the memory leak!
+    backend_args = if_google(
+        {
+            "gpu_h100": ["--heap_check="],
+        },
+        {},
+    ),
     backends = ["gpu"],
     local_defines = if_cuda(["GOOGLE_CUDA"]),
     deps = [