Add information for coordinating segments in python frontend. #3289

rdspring1 · 2024-10-27T17:42:17Z

Overview

This PR adds information necessary for coordinating segments in the python frontend. Changes pulled from #3025.

PR Details

Track the fusion state ids for the inputs, outputs, and extents of a Fusion. Inputs and extents are used to gather tensor arguments and scalars to run a fusion segment, while the outputs are employed to store results between segments.
A map from a CPP value to its corresponding fusion state id, which is needed to map values from original fusion to its segmented fusions.

Implementation Details

FusionState is a lightweight representation of a CPP Fusion.
When calling buildFusionIr, a CPP Fusion is created from the Python FusionDefinition. At this point, the FusionState creates a mapping from CPP Fusion to its State objects.
However, the FusionState is temporary and the CPP Fusion is cached in FusionCache. The information linking the CPP Fusion and Python FusionDefinition is stored in FusionCache.
When we create a new FusionState, we look for a cached CPP Fusion. If it exists, we restore the mapping from the data stored in FusionSchedules.

* Track inputs, outputs, and extents

rdspring1 · 2024-10-30T18:45:46Z

!build

jjsjann123

sorry for the delayed review.

jjsjann123 · 2024-10-30T18:32:05Z

csrc/python_frontend/fusion_state.cpp

+    }
+    TensorView* tv = v->as<TensorView>();
+    std::vector<IterDomain*> logical_dom =
+        TensorDomain::noReductions(tv->getLogicalDomain());


@wujingyue is trying to change how we bind IO buffers to kernels. i.e. we might rethink which domain and how we are going to use here.

Not proposing any change, just trying to raise awareness.

jjsjann123 · 2024-10-30T18:35:27Z

csrc/python_frontend/fusion_state.cpp

+  std::vector<Val*> extents = getExtents(fusion_);
+  for (Val* extent : extents) {
+    int64_t num_extents = (int64_t)extents_fid_.size();
+    int64_t extent_fid = -num_extents - 1;


is this a negative index?

All scalars, vectors, and tensors use positive indices. The extents do not exist in the FusionState, so I used the negative numbers exclusively for the extent scalars.

The extents are the size of iterDomain in CPP fusion. We don't track those in FusionDefinition but they can become input arguments to a FusionDefinition after segmentation.

so the negative number here is just an initialization? does the number carry any meaning or does a global -1 would do it just fine?
sorry I might miss the part where extents_fid_ is being used.

There is an ordering component to the extent index.

It is used for the same purpose as collecting the extents in prepareRuntimeOrder.
https://github.com/NVIDIA/Fuser/blob/main/csrc/runtime/fusion_cache_utils.cpp#L199-L208

We're mapping the tensor sizes to the extents like so https://github.com/NVIDIA/Fuser/pull/3025/files#diff-e512bea3b02f75ab1e81b759562879c5867e6e863679d6e7696fa34087dc3dc9R98-R100.

Can we add this in a comment listing the use of negative indices to avoid conflict with other indices.

Added comment.

Got'ya. It's hard to figure out the necessity without looking at the actual use. We can keep it as-is and revisit in follow up PRs.

jjsjann123 · 2024-10-30T18:45:25Z

csrc/python_frontend/fusion_state.cpp

+    // The extent can already exist in the fusion. However, since scalars cannot
+    // be passed between segments, always overwrited existing fids. The original
+    // fusion definition will provide scalar extents.
+    map_value_to_fid_[extent] = extent_fid;


I'm a bit lost here.

iiuc, the map_value_to_fid_ on other values are mapped from the Val* to their index field in FusionState. Here looks like we are trying to create a the same thing for each TensorView's logical domain. Where are we creating the python container for that?

I'm not exposing the TensorView's logical domain to the python frontend, but I am tracking it in the FusionState. We may have to pass the scalar extents of the TensorView's logical domain as an input argument to a fusion segment.

Priya2698 · 2024-10-30T20:51:37Z

Is it possible to add a test demonstrating what new information the FusionState stores and its link to the FusionCache?

csrc/python_frontend/fusion_cache.h

tests/python/test_python_frontend.py

Priya2698 · 2024-10-30T21:37:35Z

LGTM overall.
Do you have a document describing the intended flow of information between FusionCache, FusionState, FusionSchedules, FusionDefinition that will be built for Issue #3025?

jjsjann123

csrc/python_frontend/fusion_state.h

jjsjann123 · 2024-10-30T21:52:16Z

csrc/python_frontend/fusion_state.cpp

+  std::vector<Val*> extents = getExtents(fusion_);
+  for (Val* extent : extents) {
+    int64_t num_extents = (int64_t)extents_fid_.size();
+    int64_t extent_fid = -num_extents - 1;


Got'ya. It's hard to figure out the necessity without looking at the actual use. We can keep it as-is and revisit in follow up PRs.

rdspring1 · 2024-10-31T00:30:18Z

!build

rdspring1 · 2024-10-31T17:17:34Z

Summary

A FusionDefinition holds a series of RecordFunctors. When building the FusionDefinition, we traverse the Trie in the FusionCache. Upon reaching the EndRecord, create the CPP Fusion using the RecordFunctors. Also, generate the mappings from Scalar, Vector, and Tensor states to CPP Val.
Since the FusionDefinition is temporary, store any information in the FusionSchedules associated with this FusionDefinition. The FusionSchedules is associated with the EndRecord leaf in the FusionCache.
Now, say we have a new FusionDefinition python object with the same definition. When building the FusionDefinition, traverse the Trie in the FusionCache again. The CPP Fusion already exists, so we load the CPP Fusion to the FusionDefinition along with the stored information in FusionSchedules.

Information Flow

FusionCache holds FusionSchedules in the leaves of Trie.
FusionSchedules holds Fusion and other information.
FusionDefinition creates CPP Fusion associated with FusionCache.
FusionDefinition contains CPP Fusion and mappings from python frontend and CPP Fusion.
FusionState is the parent class of FusionDefinition.

@Priya2698

Create map value to fusion id

0b7dcdc

* Track inputs, outputs, and extents

rdspring1 added the Python API Issues related to the Python API label Oct 27, 2024

rdspring1 requested review from jacobhinkle, jjsjann123 and Priya2698 October 27, 2024 17:42

jjsjann123 reviewed Oct 30, 2024

View reviewed changes

Priya2698 reviewed Oct 30, 2024

View reviewed changes

csrc/python_frontend/fusion_cache.h Outdated Show resolved Hide resolved

create test_fusion_information

8310ace

rdspring1 force-pushed the user_sched_segmentation_mapping branch from 2cbb9d4 to 8310ace Compare October 30, 2024 21:27

Merge branch 'main' into user_sched_segmentation_mapping

946403a

Priya2698 reviewed Oct 30, 2024

View reviewed changes

tests/python/test_python_frontend.py Outdated Show resolved Hide resolved

Priya2698 approved these changes Oct 30, 2024

View reviewed changes

jjsjann123 approved these changes Oct 30, 2024

View reviewed changes

rdspring1 added 2 commits October 30, 2024 17:29

comments

09af663

Merge branch 'main' into user_sched_segmentation_mapping

f6246fa

rdspring1 merged commit 621e146 into main Oct 31, 2024
35 of 36 checks passed

rdspring1 deleted the user_sched_segmentation_mapping branch October 31, 2024 15:50

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Add information for coordinating segments in python frontend. #3289

Add information for coordinating segments in python frontend. #3289

rdspring1 commented Oct 27, 2024

rdspring1 commented Oct 30, 2024

jjsjann123 left a comment

jjsjann123 Oct 30, 2024

jjsjann123 Oct 30, 2024

rdspring1 Oct 30, 2024

jjsjann123 Oct 30, 2024

rdspring1 Oct 30, 2024

Priya2698 Oct 30, 2024

rdspring1 Oct 30, 2024

jjsjann123 Oct 30, 2024

jjsjann123 Oct 30, 2024

rdspring1 Oct 30, 2024

Priya2698 commented Oct 30, 2024

Priya2698 commented Oct 30, 2024 •

edited

Loading

jjsjann123 left a comment

jjsjann123 Oct 30, 2024

rdspring1 commented Oct 31, 2024

rdspring1 commented Oct 31, 2024

Add information for coordinating segments in python frontend. #3289

Add information for coordinating segments in python frontend. #3289

Conversation

rdspring1 commented Oct 27, 2024

Overview

PR Details

Implementation Details

rdspring1 commented Oct 30, 2024

jjsjann123 left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Priya2698 commented Oct 30, 2024

Priya2698 commented Oct 30, 2024 • edited Loading

jjsjann123 left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

rdspring1 commented Oct 31, 2024

rdspring1 commented Oct 31, 2024

Summary

Information Flow

Priya2698 commented Oct 30, 2024 •

edited

Loading