RFE: Add batching support (N-D tensors) to the gemm operator

Abierto
  • ¿Algo no está bien?
#929 0 comentarios 0 reacciones 0 asignados Ver en GitHub

Nadie ha tomado este issue todavía.

Evaluación

Dificultad
4/5
Tiempo estimado
3-5 días
Aptitud para principiantes
35/100
Tipo de issue
Nueva funcionalidad
Claridad
Bastante claro
Estado de actividad
Tranquilo
Stack tecnológico
machine-learning
Área
machine-learning

Línea de trabajo

The payload names only the MLGraphBuilder.gemm() and MLGraphBuilder.matmul() operations, so start from their definitions in the WebNN spec and compare how matmul handles leading batch dimensions. Work out the shape rules for c and the broadcasting of the result before editing anything. Done means the gemm spec accepts rank-2 and higher inputs with batch dimensions handled as matmul does, and that the change is agreed in the issue thread first.

Escrito por el modelo de indexación a partir del texto del issue.

Descripción

Problem Statement
Currently, the MLGraphBuilder.gemm() operation is restricted to 2D inputs for the a and b operands. However, MLGraphBuilder.matmul() fully supports batching via N-dimensional tensors.

If a developer wants to perform a batched General Matrix Multiplication — calculating $\alpha A B + \beta C$ across a batch — they cannot use the gemm operator natively.

Current Workaround & Limitations
To achieve a batched GEMM, developers currently have to manually emulate the operation by composing a subgraph of multiple operators:

  1. matmul (to handle the batched multiplication of A and B)
  2. mul (to scale the result by $\alpha$)
  3. mul (to scale C by $\beta$)
  4. add (to combine the results)

This emulation is undesirable for a few reasons:

  • Graph Bloat: It increases the number of nodes in the graph, adding overhead to graph construction and compilation.
  • Performance Penalties: While some backend compilers might successfully fuse these separate matmul, mul, and add nodes back into a single batched GEMM operation, this fusion is not guaranteed. If the fusion step fails or isn't supported by a specific hardware backend, the execution will incur multiple memory read/write passes instead of utilizing a single, highly optimized hardware kernel (e.g., batched GEMM on GPUs or NPUs).

Proposed Solution
Update the gemm operator specification to accept N-dimensional tensors (where rank >= 2). For inputs of rank > 2, the leading dimensions should be treated as batch dimensions exactly like the matmul operation currently does.

Lenguaje dominante
Bikeshed
Estrellas
566
Forks
72
Sin PR fusionados en 30 d

Preparar el entorno

Primeros pasos

  1. Lee el issue completo y luego la guía de contribución del proyecto.
  2. Comenta en el issue que vas a ocuparte — evita que dos personas hagan lo mismo.
  3. Haz un fork del repositorio y trabaja en una rama.
  4. Abre un pull request que haga referencia al número del issue.

Más de webmachinelearning/webnn

Todos los issues de webmachinelearning/webnn

Issues similares

Más issues de Machine Learning

Recibe los nuevos issues en tu correo

Un resumen breve de issues de GitHub para principiantes.