Move some logic from MarkAliasPass to AliasAnalysis.

wujingyue · wujingyue · commit d158799f2dc5 · 2023-12-08T16:23:18.000-08:00
This prepares #1478 for landing. We'll need to use AliasAnalysis before segmentation and during scheduling, so it's good to move common logic into AliasAnalysis.
diff --git a/csrc/optimization/alias_analysis.cpp b/csrc/optimization/alias_analysis.cpp
@@ -352,7 +352,7 @@ void AliasAnalysisResult::add(
 const Val* AliasAnalysisResult::findRoot(const Val* alias) const {
   const TensorView* root = dynamic_cast<const TensorView*>(alias);
   if (root == nullptr) {
-    return nullptr;
+    return alias;
   }
 
   // This can be made faster by path compression at the cost of losing
@@ -363,6 +363,31 @@ const Val* AliasAnalysisResult::findRoot(const Val* alias) const {
   return root;
 }
 
+const TensorView* AliasAnalysisResult::getRoot(
+    const TensorView* fusion_out) const {
+  const auto i = out_to_root_.find(fusion_out);
+  return i == out_to_root_.end() ? nullptr : i->second;
+}
+
+void AliasAnalysisResult::finalize(Fusion* fusion) {
+  for (TensorView* out :
+       ir_utils::filterByType<TensorView>(fusion->outputs())) {
+    // Lazy move: we could check compatibility and only give up when
+    // the allocation domain is incompatible with what we prefer for
+    // aliasing.
+    if (out->hasAllocation()) {
+      continue;
+    }
+
+    const Val* in = findRoot(out);
+    if (!in->isFusionInput()) {
+      continue;
+    }
+
+    out_to_root_[out] = in->as<TensorView>();
+  }
+}
+
 Layout AliasAnalysisResult::preferredLayout(const Val* v) const {
   const TensorView* tv = dynamic_cast<const TensorView*>(v);
   NVF_CHECK(
@@ -378,11 +403,23 @@ Layout AliasAnalysisResult::preferredLayout(const Val* v) const {
 
 std::string AliasAnalysisResult::toString(const int indent_size) const {
   std::stringstream ss;
+  indent(ss, indent_size) << "All aliases:"
+                          << (alias_to_source_.empty() ? " <empty>" : "")
+                          << std::endl;
   for (const auto& [alias, source_and_layout] : alias_to_source_) {
     const auto& [source, layout] = source_and_layout;
-    indent(ss, indent_size) << ir_utils::varName(alias) << " is an alias of "
-                            << ir_utils::varName(source) << " if its layout is "
-                            << layout.toString() << std::endl;
+    indent(ss, indent_size + 1)
+        << ir_utils::varName(alias) << " is an alias of "
+        << ir_utils::varName(source) << " if its layout is "
+        << layout.toString() << std::endl;
+  }
+  indent(ss, indent_size) << "Output aliases only:"
+                          << (out_to_root_.empty() ? " <empty>" : "")
+                          << std::endl;
+  for (const auto& [out, root] : out_to_root_) {
+    indent(ss, indent_size + 1)
+        << ir_utils::varName(out) << " is a transitive alias of "
+        << ir_utils::varName(root) << std::endl;
   }
   return ss.str();
 }
@@ -399,6 +436,7 @@ AliasAnalysisResult findAliases(Fusion* fusion) {
     // results).
     finder.dispatch(expr);
   }
+  analysis.finalize(fusion);
   return analysis;
 }
 
diff --git a/csrc/optimization/alias_analysis.h b/csrc/optimization/alias_analysis.h
@@ -21,6 +21,17 @@ struct Layout {
   std::string toString(int indent_size = 0) const;
 };
 
+// Holds aliases found in a fusion. The expected user flow is
+//
+// ```
+// AliasAnalysisResult analysis;
+// analysis.add(...);
+// ...
+// analysis.add(...);
+// analysis.finalize(fusion);
+//
+// // The user can now call const methods to retrieve information.
+// ```
 class AliasAnalysisResult {
  public:
   AliasAnalysisResult() = default;
@@ -29,27 +40,35 @@ class AliasAnalysisResult {
   AliasAnalysisResult(AliasAnalysisResult&&) = default;
   AliasAnalysisResult& operator=(AliasAnalysisResult&&) = default;
 
-  // Returns itself if `alias` doesn't alias anything.
-  const Val* findRoot(const Val* alias) const;
+  // Marks `source` as the immediate aliasing source of `alias` and sets the
+  // preferred layout.
+  void add(const TensorView* alias, const TensorView* source, Layout&& layout);
+
+  void finalize(Fusion* fusion);
 
   // Returns the preferred layout. If `alias` is not in `preferred_layout_`,
   // returns the `TensorView`'s initial layout.
   Layout preferredLayout(const Val* alias) const;
 
-  // Marks `source` as the immediate aliasing source of `alias` and sets the
-  // preferred layout.
-  void add(const TensorView* alias, const TensorView* source, Layout&& layout);
-
   std::string toString(int indent_size) const;
 
+  const TensorView* getRoot(const TensorView* fusion_out) const;
+
  private:
+  // Walks up `alias_to_source_` to find the root of the chain. Returns itself
+  // if `alias` doesn't alias anything.
+  const Val* findRoot(const Val* alias) const;
+
   // Maps aliases (e.g. the output of a View) to their direct sources (e.g. the
   // input of the same View). Also stores the preferred output layout for the
   // alias. Consider path compression, a common optimization used in
   // disjoint-set data structure, so it's easy to figure out the root of an
   // alias.
   std::unordered_map<const TensorView*, std::pair<const TensorView*, Layout>>
       alias_to_source_;
+
+  // Maps fusion outputs to their aliased fusion inputs.
+  std::unordered_map<const TensorView*, const TensorView*> out_to_root_;
 };
 
 // Finds aliases of the fusion inputs. The analysis should be conservative --
diff --git a/csrc/optimization/mark_alias.cpp b/csrc/optimization/mark_alias.cpp
@@ -14,28 +14,16 @@
 namespace nvfuser::optimization {
 
 void MarkAliasPass::runPass(Fusion* fusion) {
-  const AliasAnalysisResult alias_analysis = findAliases(fusion);
+  const AliasAnalysisResult analysis = findAliases(fusion);
   if (isDebugDumpEnabled(DebugDumpOption::PreSegmenterLogging)) {
     debug() << "Alias analysis result:" << std::endl;
-    debug() << alias_analysis.toString(/*indent_size=*/1) << std::endl;
+    debug() << analysis.toString(/*indent_size=*/1) << std::endl;
   }
 
   for (TensorView* out :
        ir_utils::filterByType<TensorView>(fusion->outputs())) {
-    // Lazy move: we could check compatibility and only give up when
-    // the allocation domain is incompatible with what we prefer for
-    // aliasing.
-    if (out->hasAllocation()) {
-      if (isDebugDumpEnabled(DebugDumpOption::PreSegmenterLogging)) {
-        debug() << "MarkAliasPass skipped " << out->toString()
-                << " because it already has an allocation domain:" << std::endl
-                << out->domain()->toString(1, /*leaf_only=*/false) << std::endl;
-      }
-      continue;
-    }
-
-    const Val* in = alias_analysis.findRoot(out);
-    if (!in->isFusionInput()) {
+    const Val* in = analysis.getRoot(out);
+    if (in == nullptr) {
       continue;
     }
 
@@ -55,7 +43,7 @@ void MarkAliasPass::runPass(Fusion* fusion) {
       continue;
     }
 
-    const Layout out_layout = alias_analysis.preferredLayout(out);
+    const Layout out_layout = analysis.preferredLayout(out);
     if (isDebugDumpEnabled(DebugDumpOption::PreSegmenterLogging)) {
       debug() << "MarkAliasPass changed the layout of " << out->toString()
               << std::endl;
diff --git a/test/test_alias.cpp b/test/test_alias.cpp
@@ -46,7 +46,7 @@ TEST_F(AliasAnalysisTest, View_SymbolicTensor) {
 
   optimization::AliasAnalysisResult alias_analysis =
       optimization::findAliases(&fusion);
-  EXPECT_EQ(alias_analysis.findRoot(out), in);
+  EXPECT_EQ(alias_analysis.getRoot(out), in);
 }
 
 TEST_F(AliasAnalysisTest, ChainOfViews) {
@@ -65,7 +65,7 @@ TEST_F(AliasAnalysisTest, ChainOfViews) {
 
   optimization::AliasAnalysisResult alias_analysis =
       optimization::findAliases(&fusion);
-  EXPECT_EQ(alias_analysis.findRoot(out), in);
+  EXPECT_EQ(alias_analysis.getRoot(out), in);
 }
 
 TEST_F(AliasAnalysisTest, View_Contiguous) {
@@ -82,7 +82,7 @@ TEST_F(AliasAnalysisTest, View_Contiguous) {
 
   optimization::AliasAnalysisResult alias_analysis =
       optimization::findAliases(&fusion);
-  EXPECT_EQ(alias_analysis.findRoot(out), in);
+  EXPECT_EQ(alias_analysis.getRoot(out), in);
   optimization::Layout preferred_layout = alias_analysis.preferredLayout(out);
   EXPECT_THAT(
       preferred_layout.allocation_domain,
@@ -108,7 +108,7 @@ TEST_F(AliasAnalysisTest, View_MergeNonContiguous) {
 
   optimization::AliasAnalysisResult alias_analysis =
       optimization::findAliases(&fusion);
-  EXPECT_EQ(alias_analysis.findRoot(out), out);
+  EXPECT_EQ(alias_analysis.getRoot(out), nullptr);
 }
 
 TEST_F(AliasAnalysisTest, Set) {
@@ -124,7 +124,7 @@ TEST_F(AliasAnalysisTest, Set) {
 
   optimization::AliasAnalysisResult alias_analysis =
       optimization::findAliases(&fusion);
-  EXPECT_EQ(alias_analysis.findRoot(out), in);
+  EXPECT_EQ(alias_analysis.getRoot(out), in);
 
   const std::vector<IterDomain*>& out_rfactor = out->getMaybeRFactorDomain();
   EXPECT_THAT(
@@ -145,7 +145,7 @@ TEST_F(AliasAnalysisTest, Permute) {
 
   optimization::AliasAnalysisResult alias_analysis =
       optimization::findAliases(&fusion);
-  EXPECT_EQ(alias_analysis.findRoot(out), in);
+  EXPECT_EQ(alias_analysis.getRoot(out), in);
 
   const std::vector<IterDomain*>& out_rfactor = out->getMaybeRFactorDomain();
   EXPECT_THAT(
@@ -157,47 +157,49 @@ TEST_F(AliasAnalysisTest, View_SplitExpandedBroadcast) {
   Fusion fusion;
   FusionGuard fg(&fusion);
 
-  TensorView* in = makeContigConcreteTensor({4, 5});
+  TensorView* in = TensorViewBuilder()
+                       .ndims(3)
+                       .dtype(DataType::Float)
+                       .contiguity({true, true, std::nullopt})
+                       .shape({4, 5, 6})
+                       .expanded({false, false, true})
+                       .build();
   fusion.addInput(in);
-  TensorView* broadcast_out = broadcast(in, {false, false, true});
-  TensorView* expand_out = expand(
-      broadcast_out,
-      {IrBuilder::create<Val>(4),
-       IrBuilder::create<Val>(5),
-       IrBuilder::create<Val>(6)});
   // tryStaticReshape used to fail to get the expanded extent, which is 6.
-  TensorView* out = reshape(
-      expand_out, {IrBuilder::create<Val>(40), IrBuilder::create<Val>(3)});
+  // Therefore, we use the `vector<Val*>` version of `reshape` as a regression
+  // test.
+  TensorView* out =
+      reshape(in, {IrBuilder::create<Val>(40), IrBuilder::create<Val>(3)});
   fusion.addOutput(out);
 
   optimization::AliasAnalysisResult alias_analysis =
       optimization::findAliases(&fusion);
-  EXPECT_EQ(alias_analysis.findRoot(out), out);
+  EXPECT_EQ(alias_analysis.getRoot(out), nullptr);
 }
 
 TEST_F(AliasAnalysisTest, View_ForwardExpandedBroadcast) {
   Fusion fusion;
   FusionGuard fg(&fusion);
 
-  TensorView* in = makeContigConcreteTensor({4, 5});
+  TensorView* in = TensorViewBuilder()
+                       .ndims(3)
+                       .dtype(DataType::Float)
+                       .contiguity({true, true, std::nullopt})
+                       .shape({4, 5, 6})
+                       .expanded({false, false, true})
+                       .build();
   fusion.addInput(in);
-  TensorView* broadcast_out = broadcast(in, {false, false, true});
-  TensorView* expand_out = expand(
-      broadcast_out,
-      {IrBuilder::create<Val>(4),
-       IrBuilder::create<Val>(5),
-       IrBuilder::create<Val>(6)});
-  TensorView* out = reshape(expand_out, {4, 5, 6}, {20, -1});
+  TensorView* out = reshape(in, {4, 5, 6}, {20, -1});
   fusion.addOutput(out);
 
   optimization::AliasAnalysisResult alias_analysis =
       optimization::findAliases(&fusion);
-  EXPECT_EQ(alias_analysis.findRoot(out), expand_out);
+  EXPECT_EQ(alias_analysis.getRoot(out), in);
 
   // Verify the last dimension isn't expanded physically.
   FusionExecutor fe;
   at::Tensor in_tensor =
-      at::randn({4, 5}, at::dtype(at::kFloat).device(at::kCUDA, 0));
+      at::randn({4, 5}).cuda().as_strided({4, 5, 6}, {5, 1, 0});
   fe.compileFusion(&fusion, {in_tensor});
   at::Tensor out_tensor = fe.runFusion({in_tensor})[0];
 
@@ -208,20 +210,20 @@ TEST_F(AliasAnalysisTest, View_MergeExpandedBroadcast) {
   Fusion fusion;
   FusionGuard fg(&fusion);
 
-  TensorView* in = makeContigConcreteTensor({4, 5});
+  TensorView* in = TensorViewBuilder()
+                       .ndims(3)
+                       .dtype(DataType::Float)
+                       .contiguity({true, true, std::nullopt})
+                       .shape({4, 5, 6})
+                       .expanded({false, false, true})
+                       .build();
   fusion.addInput(in);
-  TensorView* broadcast_out = broadcast(in, {false, false, true});
-  TensorView* expand_out = expand(
-      broadcast_out,
-      {IrBuilder::create<Val>(4),
-       IrBuilder::create<Val>(5),
-       IrBuilder::create<Val>(6)});
-  TensorView* out = reshape(expand_out, {4, 5, 6}, {4, -1});
+  TensorView* out = reshape(in, {4, 5, 6}, {4, -1});
   fusion.addOutput(out);
 
   optimization::AliasAnalysisResult alias_analysis =
       optimization::findAliases(&fusion);
-  EXPECT_EQ(alias_analysis.findRoot(out), out);
+  EXPECT_EQ(alias_analysis.getRoot(out), nullptr);
 }
 
 TEST_F(AliasAnalysisTest, TrivialSlice) {
@@ -236,7 +238,7 @@ TEST_F(AliasAnalysisTest, TrivialSlice) {
 
   optimization::AliasAnalysisResult alias_analysis =
       optimization::findAliases(&fusion);
-  EXPECT_EQ(alias_analysis.findRoot(out), in);
+  EXPECT_EQ(alias_analysis.getRoot(out), in);
 }
 
 TEST_F(AliasAnalysisTest, MergeTriviallySlicedDimensions) {
@@ -251,7 +253,7 @@ TEST_F(AliasAnalysisTest, MergeTriviallySlicedDimensions) {
 
   optimization::AliasAnalysisResult alias_analysis =
       optimization::findAliases(&fusion);
-  EXPECT_EQ(alias_analysis.findRoot(out), in);
+  EXPECT_EQ(alias_analysis.getRoot(out), in);
 }
 
 TEST_F(AliasAnalysisTest, MergeSlicedDimensions) {
@@ -266,8 +268,7 @@ TEST_F(AliasAnalysisTest, MergeSlicedDimensions) {
 
   optimization::AliasAnalysisResult alias_analysis =
       optimization::findAliases(&fusion);
-  EXPECT_EQ(alias_analysis.findRoot(out), out);
-  EXPECT_EQ(alias_analysis.findRoot(slice_out), in);
+  EXPECT_EQ(alias_analysis.getRoot(out), nullptr);
 }
 
 using AliasTest = NVFuserTest;
