TransformReplay::selfReplay replays contiguity #5316

wujingyue · 2025-10-03T02:39:55Z

Fixes bugs like #5356

We should probably add some knobs so callers decide what to replay, in a separate PR. So far, this function tries to replay everything (namely, loop, allocation and contiguity), which seems to work fine.

wujingyue · 2025-10-03T02:40:12Z

!test

github-actions · 2025-10-03T02:40:39Z

Review updated until commit 23ee554

Description

Fix contiguity replay in TransformReplay::selfReplay
Refactor selfReplay to handle loop and allocation consistently
Improve error handling and validation in replay logic
Add test for contiguity replay with empty allocation

Changes walkthrough 📝

Relevant files

Bug fix

nodes.cpp `Improve contiguity validation` csrc/ir/nodes.cpp Replace NVF_CHECK with NVF_CHECK_EQ for clearer error message Validate contiguity vector size against allocation domain	+3/-6
transform_replay.cpp `Refactor and fix selfReplay logic` csrc/transform_replay.cpp Refactor selfReplay into a lambda to build IterDomainMap Remove ignore_reductions flag and simplify reduction handling Fix contiguity replay by properly mapping and setting contiguity Add proper handling of symbolic and broadcast domains	+130/-85

Tests

test_alias.cpp `Add test for accumulate slices aliasing` tests/cpp/test_alias.cpp Add AccumulateSlices test for aliasing with symbolic domains Test replay of contiguity from broadcast to symbolic ID	+55/-25
test_replay.cpp `Add contiguity replay test` tests/cpp/test_replay.cpp Add ContiguityWithEmptyAllocation test Verify contiguity replay when allocation domain is empty	+16/-0

PR Reviewer Guide 🔍

Here are some key observations to aid the review process:

🧪 PR contains tests

⚡ Recommended focus areas for review

Contiguity Replay Logic

The new code introduces a complex replay mechanism for contiguity, particularly when handling symbolic and broadcast dimensions. The logic for setting contiguity on symbolic IDs may need validation, especially the decision to default to true when contiguity is null. This could impact correctness in cases involving symbolic dimensions.

} else {
  const std::vector<IterDomain*>& new_logical = new_self->logical();
  const auto new_rank = std::ssize(new_logical);
  std::vector<std::optional<bool>> new_contiguities(new_rank, std::nullopt);

  int new_pos = 0;
  for (auto [id, contiguity] : zip(self->logical(), self->contiguity())) {
    IterDomain* new_id = getOrDefault(replay, id);
    if (new_id == nullptr) {
      continue;
    }

    // Find the corresponding contiguity in new_logical. Mapped IterDomains
    // in self->logical() and new_logical follow the same order. So it's safe
    // to only increment `new_pos`.
    while (new_pos < new_rank && new_logical.at(new_pos) != new_id) {
      new_pos++;
    }
    NVF_ERROR_LT(
        new_pos,
        new_rank,
        "Failed to find ",
        new_id->toString(),
        " in ",
        new_logical);
    std::optional<bool>& new_contiguity = new_contiguities.at(new_pos);

    new_contiguity = contiguity;
    // When used during or before concretization, TransformReplay::selfReplay
    // can be applied to replay transformations from symbolic dimensions to
    // concrete dimensions, or in the reverse direction. Therefore,
    // `new_contiguity` is not always identical to `contiguity`.
    if (new_id->isBroadcast()) {
      new_contiguity = std::nullopt;
    } else if (new_id->isSymbolic()) {
      // See AliasTest.AccumulateSlices for an example. aliasOutputToInput is
      // called before concretization and tries to replay contiguity from a
      // broadcast IterDomain to a symbolic IterDomain. However, a symbolic
      // IterDomain can't have contiguity null.
      if (!new_contiguity.has_value()) {
        new_contiguity = true;
      }
    }
  }

  new_self->setContiguity(new_contiguities);

Reduction Axis Handling

The updated code removes the ignore_reductions flag and instead uses mapped_new_ids to filter reduction axes. This change alters the control flow significantly and may affect how reduction dimensions are replayed. The interaction between mapped_new_ids and reduction axes should be carefully reviewed for correctness.

if (self->loop() != self->logical()) {
  std::vector<IterDomain*> new_loop;
  for (auto* new_id : new_self->logical()) {
    if (mapped_new_ids.count(new_id) == 0) {
      NVF_ERROR(
          new_id->isReduction(),
          new_id->toString(),
          " should be a reduction.");
      new_loop.push_back(new_id);
    }

NVF_CHECK_EQ Usage

The use of NVF_CHECK_EQ in validateContiguity has been introduced to replace the previous NVF_CHECK with manual size comparison. While this simplifies the code, it may obscure debugging information if the check fails. The trade-off between conciseness and debuggability should be considered.

NVF_CHECK_EQ(
    contiguity.size(),
    allocation_domain.size(),
    "Invalid contiguity information provided, incorrect size.");

This reverts commit a41c05e.

wujingyue · 2025-10-06T03:49:37Z

!test

wujingyue · 2025-10-06T06:28:06Z

!test

wujingyue · 2025-10-06T17:42:42Z

!test

wujingyue · 2025-10-06T20:01:43Z

!test

wujingyue · 2025-10-06T21:16:45Z

!test

wujingyue · 2025-10-06T23:23:17Z

!test

wujingyue · 2025-10-07T04:19:39Z

!test

wujingyue · 2025-10-07T04:21:11Z

tests/cpp/test_alias.cpp

+  fusion->addInput(in);
+  fusion->addInput(i);
+  fusion->addOutput(acc_out);
+  fusion->aliasOutputToInput(acc_out, acc_in, AllocationType::ReuseBuffer);


This calls TransformReplay::selfReplay to replay transformations from concrete dimensions to symbolic dimensions.

wujingyue · 2025-10-07T04:21:47Z

tests/cpp/test_alias.cpp

  auto out_tensors = executor_cache.runFusionWithInputs({in_tensor});
  ASSERT_EQ(out_tensors.size(), 1);
-  at::Tensor out_tensor = out_tensors[0].as<at::Tensor>();
+  auto out_tensor = out_tensors[0].as<at::Tensor>();


Use auto because as<at::Tensor> already specifies the type.

csrc/transform_replay.cpp

wujingyue · 2025-10-07T22:35:44Z

!test

wujingyue · 2025-10-08T01:10:03Z

!test

naoyam · 2025-10-08T19:05:56Z

It would be really helpful to have some quick PR introduction to remind the context. My mental capacity is not big enough to remember everything currently going on.

Priya2698 · 2025-10-08T23:42:17Z

Thanks @wujingyue for the clarifications.
Changes LGTM, I'll defer approval to @naoyam since he had concerns about changes in this function.

wujingyue · 2025-10-09T01:53:25Z

It would be really helpful to have some quick PR introduction to remind the context. My mental capacity is not big enough to remember everything currently going on.

Done

wujingyue · 2025-10-09T05:35:11Z

!test

naoyam · 2025-10-17T02:25:05Z

csrc/transform_replay.cpp

    }

    new_self->setAllocationDomain(new_allocation, new_contiguities);
+  } else {


This feels a little unexpected to me because, even though there's nothing to replay for the allocation, the contiguity could be modified.

What would happen if self doesn't have an allocation domain but new_self does? Would it work?

naoyam · 2025-10-17T02:25:20Z

!test --diff

naoyam · 2025-10-17T02:28:09Z

Fixes bugs like #5356

We should probably add some knobs so callers decide what to replay, in a separate PR. So far, this function tries to replay everything (namely, loop, allocation and contiguity), which seems to work fine.

Changes of contiguity may not cause test failures but could result in, e.g., different vectorizations. Started the diff check just in case.

Replay loop unconditionally

a41c05e

wujingyue added 2 commits October 2, 2025 20:43

WIP

3c09559

Revert "Replay loop unconditionally"

ddbcbbd

This reverts commit a41c05e.

wujingyue changed the title ~~Replay loop unconditionally~~ TransformReplay::selfReplay replays contiguity Oct 6, 2025

Fix

7003c37

Fix

30a4c19

Fix

cb4a0c6

wujingyue added 3 commits October 6, 2025 13:08

Nit

2ed59d4

Remove a check

a0711a8

Remove a fixme

8479b97

Fix

cbef21b

wujingyue added 4 commits October 6, 2025 19:13

Remove ignore_reductions

c9e9615

Build replay in a separate method

2edd504

Refactor

f65e0b5

Comment

df825a6

wujingyue requested a review from Priya2698 October 7, 2025 04:19

wujingyue commented Oct 7, 2025

View reviewed changes

Priya2698 requested a review from naoyam October 7, 2025 17:52

Priya2698 reviewed Oct 7, 2025

View reviewed changes

csrc/transform_replay.cpp Show resolved Hide resolved

csrc/transform_replay.cpp Show resolved Hide resolved

csrc/transform_replay.cpp Show resolved Hide resolved

Review

974748a

wujingyue requested a review from Priya2698 October 7, 2025 22:35

Merge branch 'main' into wjy/replay

4232e21

wujingyue mentioned this pull request Oct 9, 2025

In-place updating a non-contiguous tensor leads to a wrong result #5356

Open

Minor

80e8dff

Priya2698 mentioned this pull request Oct 9, 2025

Lower stream-parallelized matmul #5302

Merged

naoyam reviewed Oct 17, 2025

View reviewed changes

Merge branch 'main' into wjy/replay

23ee554

Uh oh!

TransformReplay::selfReplay replays contiguity #5316

Are you sure you want to change the base?

TransformReplay::selfReplay replays contiguity #5316

Uh oh!

Conversation

wujingyue commented Oct 3, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

wujingyue commented Oct 3, 2025

Uh oh!

github-actions bot commented Oct 3, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Changes walkthrough 📝

PR Reviewer Guide 🔍

Uh oh!

wujingyue commented Oct 6, 2025

Uh oh!

wujingyue commented Oct 6, 2025

Uh oh!

wujingyue commented Oct 6, 2025

Uh oh!

wujingyue commented Oct 6, 2025

Uh oh!

wujingyue commented Oct 6, 2025

Uh oh!

wujingyue commented Oct 6, 2025

Uh oh!

wujingyue commented Oct 7, 2025

Uh oh!

wujingyue Oct 7, 2025

Choose a reason for hiding this comment

Uh oh!

wujingyue Oct 7, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

wujingyue commented Oct 7, 2025

Uh oh!

wujingyue commented Oct 8, 2025

Uh oh!

naoyam commented Oct 8, 2025

Uh oh!

Priya2698 commented Oct 8, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

wujingyue commented Oct 9, 2025

Uh oh!

wujingyue commented Oct 9, 2025

Uh oh!

naoyam Oct 17, 2025

Choose a reason for hiding this comment

Uh oh!

naoyam Oct 17, 2025

Choose a reason for hiding this comment

Uh oh!

naoyam commented Oct 17, 2025

Uh oh!

naoyam commented Oct 17, 2025

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

wujingyue commented Oct 3, 2025 •

edited

Loading

github-actions bot commented Oct 3, 2025 •

edited

Loading

wujingyue Oct 7, 2025 •

edited

Loading

Priya2698 commented Oct 8, 2025 •

edited

Loading