Exclude IndexSelectOp from inlineMost

[jjsjann123]

This is to restore a performance regression for pointwise kernel containing IndexSelectOp.

IndexSelectOp supports vectorized load. Excluding it from inlineMost allows a codegen change from

array<vector_factor> buffer;
for unroll:
    vec_load(buffer)  // IndexSelectOp
    ... // computation consuming value of buffer

to

array<vector_factor * unroll> buffer;
for unroll:
    vec_load(&buffer[i])  // IndexSelectOp
for unroll:
    ... // computation consuming value of buffer

tasks:

• code comment
• add benchmark numbers in PR

[github-actions]

Review updated until commit 222bc69

Description

• Exclude IndexSelectOp from inlineMost for performance

• Aggregate allocation of manual unroll ID and inner ID

• Add comment explaining the change

---

Changes walkthrough 📝

	Relevant files
Enhancement	pointwise.cpp Exclude IndexSelectOp from inlineMost                                        csrc/scheduler/pointwise.cpp Add logic to exclude IndexSelectOp output from inner_most_tensors Add comment explaining the rationale behind the change +8/-0

---

PR Reviewer Guide 🔍

Here are some key observations to aid the review process:

🧪 No relevant tests

⚡ Recommended focus areas for review Performance Impact The change excludes IndexSelectOp from inlineMost, which may have implications on performance. Ensure that the performance metrics provided in the PR description accurately reflect the intended performance gains and that there are no unintended regressions. // IndexSelectOp reads lookup tv without cache. Because pointwise scheduler // doesn't use ParallelType::Unroll, we need to exclude consumer of fusion // inputs to be inlineMost. This allows us to aggregate the allocation of // manual unroll ID and its inner ID. for (auto idx_sel : ir_utils::getOpsOfType<IndexSelectOp>(fusion)) { inner_most_tensors.erase(idx_sel->output(0)->as<TensorView>()); }

[jjsjann123]

Comparing the performance across the embedding fwd python benchmark. This is the speedup we get from this PR.

[jjsjann123]

!test

[jjsjann123]

!test

[jjsjann123]

!test

[jjsjann123]

@naoyam updated the comment, this one should be good to go.

[naoyam]

LGTM