NVIDIA/Fuser

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

Open

#4491 aperta il 21 mag 2025

Vedi su GitHub
 (0 commenti) (0 reazioni) (1 assegnatario)C++ (80 fork)github user discovery
New Opsgood first issue

Metriche repository

Star
 (392 star)
Metriche merge PR
 (Metriche PR in attesa)

Descrizione

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.

Guida contributor

Direzione di ricerca
Studia il passaggio di pre segmentazione esistente (PR #4399) per capire come 'take along axis' viene spostato prima degli operatori unari. Estendi la logica per supportare più operazioni 'take along axis' iterandole. Per il caso squeeze, rimuovi la limitazione di singola dimensione e aggiungi più operazioni 'unsqueeze' al tensore di indice per ogni dimensione compressa.
Tech stack
cpp
Dominio
backend
Tipo issue
Funzionalità
DifficoltàQuanto dovrebbe essere impegnativa per un nuovo contributor.
3
Tempo stimatoTempo stimato per completare e verificare uno scope piccolo.
1-2 giorni
Stato attivitàQuanto è probabile che la issue sia ancora attiva e reviewable.
Attiva
ChiarezzaQuanto chiaramente la issue descrive problema, scope e risultato atteso.
Chiara
Prerequisiti
C++
Adatta ai principiantiQuanto la issue è adatta a contributor alla prima esperienza.
40