From 7921e63800159d5bf111721cc983942f067308fa Mon Sep 17 00:00:00 2001
From: Jingyue Wu <wujingyue@gmail.com>
Date: Fri, 23 Aug 2024 23:53:25 +0000
Subject: [PATCH] Fix NVFuserTest.AvoidCachingSliceInput.

---
 tests/cpp/test_gpu3.cpp | 8 ++++++++
 1 file changed, 8 insertions(+)

diff --git a/tests/cpp/test_gpu3.cpp b/tests/cpp/test_gpu3.cpp
index f93501c2215..a1059057275 100644
--- a/tests/cpp/test_gpu3.cpp
+++ b/tests/cpp/test_gpu3.cpp
@@ -33,6 +33,8 @@
 #include <kernel_ir_dispatch.h>
 #include <logical_domain_map.h>
 #include <ops/all_ops.h>
+#include <preseg_passes/mark_aliases_prepare.h>
+#include <preseg_passes/optimization_pass.h>
 #include <scheduler/all_schedulers.h>
 #include <scheduler/reduction_utils.h>
 #include <scheduler/utils.h>
@@ -7805,6 +7807,12 @@ TEST_F(NVFuserTest, AvoidCachingSliceInput) {
   auto fusion = std::make_unique<Fusion>();
   FusionGuard fg(fusion.get());
 
+  // Avoid `slice`s from being bookended. This test is to exercise kernel
+  // caching when a `SliceOp` is applied on a segment input. Bookending
+  // `slice`s would defeat that purpose.
+  preseg_passes::OptimizationPassGuard<preseg_passes::MarkAliasesPreparePass>
+      optimization_guard(false);
+
   // values to trigger the original bug.
   const int64_t eight = 8;
   const int64_t twenty = 20;