NVIDIA / Fuser

A Fusion Code Generator for NVIDIA GPUs (commonly known as "nvFuser")
Other
271 stars 53 forks source link

Preserve additional IDs when mutating TensorDomains in OptOutMutator #3410

Closed jacobhinkle closed 6 days ago

jacobhinkle commented 6 days ago

Fixes #3409

This PR does the following:

  1. Creates tests/cpp/test_mutator.cpp and moves the OptOutMutator tests from test_dynamic_transform.cpp into that file.
  2. Adds a new test that repros #3409
  3. Adds an additional argument const std::vector<IterDomain*>& additional_ids to one of the TensorDomain constructors.
  4. Uses that constructor in OptOutMutator::mutate(TensorDomain*) to pass a mutated additional_ids.

A bit off-topic maybe: implementing this got me thinking more about traversals and whether or not StmtSort::getStmts(fusion, /*traverse_members=*/true) will still guarantee full traversals now that the loop domain is not necessarily the set of leaf IterDomains in the TensorDomain: https://github.com/NVIDIA/Fuser/blob/f1a7e23604dc8639bec488e472909eca8bdbbb31/csrc/iter_visitor.cpp#L71-L74 I think in another PR we might want to modify this code to ensure that we traverse all of domain->allIDs() in topological order.

jacobhinkle commented 6 days ago

!test

jacobhinkle commented 6 days ago

!build