Added support for the linalg.broadcast operation and its kernel - #3163
Added support for the linalg.broadcast operation and its kernel#3163Someone117 wants to merge 1 commit into
Conversation
0920ae3 to
42831eb
Compare
…ell as the ability to remove the broadcast operation when it is unnecessary.
2ed228e to
a1a493b
Compare
j2kun
left a comment
j2kun
left a comment
There was a problem hiding this comment.
So overall I think it looks good. I was hoping for a few small changes:
- Can you add the patterns from the two new passes (broadcast-canonicalizations and broadcast-swap) to
linalg-canonicalizationsinstead? (And ideally in a separate PR) - Can you add a simple end-to-end test in tests/Examples/lattigo/ckks that runs against the Lattigo backend?
- As far as I understand, the
implementRotateAndBroadcasthas a limitation on the layouts of its inputs: if the data matrix is laid out among multiple ciphertexts and the broadcast spills over from one ciphertext to another, it will not work with this kernel. Can we guard against this? Imagine for example having a tensor of size 256x256 packed into two ciphertexts of size 32768 in a row-major fashion. reducing along the leading axis and broadcasting would require mapping from one ciphertext to another (I believe I didn't get that backwards...). I think the "right" way to handle this generally is to create a tensor_ext.remap op and write down the broadcast as an explicit permutation (or an ISL mapping string) and then use the existing lowerings to generate a shift network. But that can also be beyond the scope of this PR.
| } | ||
|
|
||
| // if the reduce is immediately broadcasted, the broadcast is a NoOp | ||
| if (checkReductionThenBroadcast(op)) { |
There was a problem hiding this comment.
This block looks like a duplicate of the previous block
| setMaterializedAttr(addBias); | ||
| rewriter.replaceOp(op, addBias); | ||
|
|
||
| // if the reduce is immediately broadcasted, the broadcast is a NoOp |
There was a problem hiding this comment.
Why exactly is it a no-op? I suspect this depends on the chosen layout. For example, we have some cases where we want to do row-summations, but the matrix is laid out in the "bicyclic" layout, which involves a diagonal traversal of the matrix. This led me to define a new, dedicated op for the combined reduce+broadcast op in #3269
- Could we add a check on which kinds of layouts this no-op holds for?
- Would it make sense to instead use the "broadcast swap" in combination with looking for reduce + broadcast to produce the
broadcasted_reduceop as a single op, and then match on it with a dedicated pattern?
There was a problem hiding this comment.
If we do want to merge the reduce + broadcast ops, then the KernelImplementation code in this PR can probably be merged with
heir/lib/Kernel/KernelImplementation.h
Line 517 in cfae7f2
To my eye, your implementation is more general than the linked one (because it supports non-power-of-two and is smarter about its axis analysis).
| return rewriter.notifyMatchFailure( | ||
| op, "missing new layout attribute for input"); | ||
|
|
||
| // TODO: support multi-dimension broadcasts |
There was a problem hiding this comment.
Create a GH issue and tag the issue as // TODO(#dddd): message, which will allow GH to link the line of code to the issue.
3 participants
This patch resolves #2876 by implementing a linalg broadcast kernel, rewerting some linalg.generic operations, and removing unnecessary linalg.broadcast operations.