tree 3d04d12094cb0ad90f35bfc625decb3b6c9f1fe9
parent 6caa91743beb38c369d45f5ccf10f6bd5fc2fd57
author Roberto Laudani <laudani@roofline.ai> 1783602163 +0200
committer GitHub <noreply@github.com> 1783602163 +0200
gpgsig -----BEGIN PGP SIGNATURE-----
 
 wsFcBAABCAAQBQJqT5vzCRC1aQ7uu5UhlAAAOwsQAGVe661ZkOHIogbTAE5q3JaG
 9DWQlAyFUYlLMlIrmLQ6tninwXPWys4EPIg4JCH8mtiPt3C6gBEi7DZZ3u3+euMK
 17bLXhEn8W+I7okRnfUnZh9nuJABOPutSY99QaccaEL+1p2sQEFPONyYasFebgve
 u5TTujeD1wN8HSW4AQ7OxRpU5fkrMp16ix4B4bdbQ/DkLAuA0rI74znnTLRiHeXs
 WGgTgUuRsmrrhtbqFocu6GG0Q6hbBVWpwsTXJVnfmDnbLanqZktzzm2LcT/iemtc
 8ampQajXNiF2rMWlN7ul3+5z46uWtBPjAiuhW30aKtnzTKVl1rO8lV1nIg3IHTMg
 oKXcX6EynqDCHu2o+i/Dks+hR20SZZRIHOp7jNcW7ekYSyBuqzFbRL5Nyh6z9OVK
 n228iei2XaUVdEpM8y9v4nvJqyLqfyA5saWWAtNqt5oyWicJBhRN2ESJ5tNJwG0d
 GQb6juqbFCN8aa2THyZbpOcRE4q5cESZAIqFrtoVV+PUl452l+RsYP2gbbfxEQr/
 5A+aAsW2Uzk53UYFBJhTzbNgOnYEu9gPu2HV1ajPsTuuwoCrVuy17c/0aj8CJowb
 CwHrZonYc1R5AgTlRadY9kuFkuB2Ugv5KbWfycVpRs0JrqM+WzPNrEAzF7RZEOh4
 h4NgMnrv16nS6yrl0hku
 =VIGD
 -----END PGP SIGNATURE-----
 

[GlobalOpt] Add pass to convert broadcast batch_matmul to matmul (#24670)

When a `linalg.batch_matmul` multiplies a batched input by a weight that
is the same for every batch, every batch slice runs the same matmul
against the same weight. This pass rewrites that case into a single
`linalg.matmul`.

For an input of shape `[B, M, K]` and a weight `[K, N]` broadcast to
`[B, K, N]`, the pass:

  1. collapses the input to `[B*M, K]` using tensor.collapse_shape
  2. multiplies it by the `[K, N]` weight using linalg.matmul
  3. expands the result back to `[B, M, N]` using tensor.expand_shape

This removes the broadcast and turns the batched contraction into a
plain
matmul, which downstream passes can handle better.

---------

Signed-off-by: Roberto Laudani <laudani@roofline.ai>
Co-authored-by: ziereis <ziereis@roofline.ai>