Preserve allocation domain when inserting segment_set. (#2808)

For #2599.
NVIDIA · Aug 20, 2024 · cf02392 · cf02392
1 parent 46b583a
commit cf02392
Show file tree

Hide file tree

Showing 2 changed files with 35 additions and 0 deletions.
diff --git a/csrc/preseg_passes/mark_aliases_prepare.cpp b/csrc/preseg_passes/mark_aliases_prepare.cpp
@@ -12,6 +12,7 @@
 #include <ops/alias.h>
 #include <options.h>
 #include <preseg_passes/mark_aliases_prepare.h>
+#include <transform_replay.h>
 
 namespace nvfuser::preseg_passes {
 
@@ -99,6 +100,16 @@ void insertSegmentSetAfter(
 
   // The general case.
   TensorView* copy = segment_set(use_of);
+  // Inherit the allocation domain from `use_of`. This is important to pass
+  // AliasTest.Bookend_SegmentSetPreservesAllocation.
+  TensorDomain* replayed_domain =
+      TransformReplay::replayCasP(
+          copy, use_of, -1, TransformReplayOptions().replayAllocation())
+          .first;
+  if (replayed_domain->hasAllocation()) {
+    copy->setAllocationDomain(
+        replayed_domain->allocation(), replayed_domain->contiguity());
+  }
   std::for_each(first_user, last_user, [&](const Use& use) {
     ir_utils::replaceValInExprInputs(use.user, use_of, copy);
   });

diff --git a/tests/cpp/test_alias.cpp b/tests/cpp/test_alias.cpp
@@ -1196,6 +1196,30 @@ TEST_F(AliasTest, InplaceUpdate) {
       UnorderedElementsAre(HeuristicIs(ScheduleHeuristic::PointWise)));
 }
 
+TEST_F(AliasTest, Bookend_SegmentSetPreservesAllocation) {
+  auto fusion = std::make_unique<Fusion>();
+  FusionGuard fg(fusion.get());
+
+  TensorView* in = makeContigConcreteTensor({2, 3});
+  TensorView* permute_out = permute(in, {1, 0});
+  TensorView* compute_out = mul(in, in);
+  fusion->addInput(in);
+  fusion->addOutput(permute_out);
+  fusion->addOutput(compute_out);
+
+  in->setAllocationDomain({in->axis(1), in->axis(0)}, true);
+  permute_out->setAllocationDomain(
+      {permute_out->axis(0), permute_out->axis(1)}, true);
+
+  FusionExecutorCache fec(std::move(fusion));
+  at::Tensor in_tensor = at::randn({3, 2}).cuda().transpose(0, 1);
+  std::vector<at::Tensor> out_tensors = fec.runFusionWithInputs({in_tensor});
+  testValidate(fec.fusion(), out_tensors, {in_tensor}, __LINE__, __FILE__);
+
+  at::Tensor permute_out_tensor = out_tensors[0];
+  EXPECT_TRUE(permute_out_tensor.is_alias_of(in_tensor));
+}
+
 TEST_F(AliasTest, Bookend_InputsAndOutputs) {
   auto fusion = std::make_unique<Fusion>();
   FusionGuard fg(fusion.get());