NVIDIA/Fuser

Add more features to the pre-segmentation pass to move take_along_axis ahead of certain unary ops

Open

#4,491 opened on May 21, 2025

View on GitHub
 (0 comments) (0 reactions) (1 assignee)C++ (80 forks)github user discovery
New Opsgood first issue

Repository metrics

Stars
 (392 stars)
PR merge metrics
 (PR metrics pending)

Description

The PR add a pre-segmentation pass that moved a take_along_axis operator ahead of certain unary pointwise operators such as castOp and the squeeze op.

The implementation had certain limitations. We should extend the current implementation to remove these.

  • We only consider fusions with a single take_along_axis op. We should extend our implementation to support the case where there are multiple take_along_axis ops.
  • We move a take_along_axis ahead of a squeeze op. Currently we do so only if a single dimension is being squeezed. We should remove this limitation. This will mean that we have to add multiple unsqueeze ops to the indexTv of the take_along_axis.

Contributor guide

Research direction
Study the existing pre segmentation pass (PR #4399) to understand how 'take along axis' is moved ahead of unary ops. Extend the logic to support multiple 'take along axis' ops by iterating over them. For the squeeze case, remove the single dimension limitation and add multiple 'unsqueeze' ops to the index tensor for each squeezed dimension.
Tech stack
cpp
Domain
backend
Issue type
Feature
DifficultyEstimated implementation difficulty for a new contributor, from 1 for very small changes to 5 for expert-level work.
3
Estimated timeA rough time range for an experienced contributor to investigate, implement, test, and prepare a pull request.
1-2 days
Activity statusHow available the issue appears right now: fresh, active, stale, blocked, or waiting on maintainer input.
Active
ClarityHow clearly the issue explains the expected change, acceptance criteria, and next step.
Clear
Prerequisites
C++
Newbie friendlinessA 1-100 score estimating how approachable this issue is for first-time contributors.
40