NVIDIA/Fuser

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

Open

#4 491 ouverte le 21 mai 2025

Voir sur GitHub
 (0 commentaires) (0 réactions) (1 assigné)C++ (80 forks)github user discovery
New Opsgood first issue

Métriques du dépôt

Stars
 (392 stars)
Métriques de merge PR
 (Métriques PR en attente)

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.

Guide contributeur

Direction de recherche
Étudiez le passage de pré segmentation existant (PR #4399) pour comprendre comment 'take along axis' est déplacé avant les opérations unaires. Étendez la logique pour prendre en charge plusieurs opérations 'take along axis' en les parcourant. Pour le cas squeeze, supprimez la limitation de dimension unique et ajoutez plusieurs opérations 'unsqueeze' au tenseur d'index pour chaque dimension compressée.
Stack technique
cpp
Domaine
backend
Type d'issue
Fonctionnalité
DifficultéDifficulté estimée pour un nouveau contributeur, de 1 pour un très petit changement à 5 pour un travail expert.
3
Temps estiméFourchette de temps approximative pour investiguer, implémenter, tester et préparer une pull request.
1-2 jours
Statut d'activitéDisponibilité apparente de l'issue : fraîche, active, ancienne, bloquée ou en attente d'un mainteneur.
Active
ClartéClarté avec laquelle l'issue explique le changement attendu, les critères d'acceptation et la prochaine étape.
Claire
Prérequis
C++
Accessibilité débutantScore de 1 à 100 estimant l'accessibilité de cette issue pour un premier contributeur.
40