support dilation and feature/batch group count in convolution reverse #181
Conversation
| // REVERSE-NEXT{LITERAL}: %1 = stablehlo.reshape %0 : (tensor<8x66x66x512xf32>) -> tensor<8x66x66x1x512xf32> | ||
| // REVERSE-NEXT{LITERAL}: %2 = stablehlo.transpose %1, dims = [1, 2, 0, 3, 4] : (tensor<8x66x66x1x512xf32>) -> tensor<66x66x8x1x512xf32> | ||
| // REVERSE-NEXT{LITERAL}: %3 = stablehlo.reshape %2 : (tensor<66x66x8x1x512xf32>) -> tensor<8x66x66x512xf32> |
There was a problem hiding this comment.
We should add an optimization to get rid of this pattern in a follow-up:
(transpose (unsqueeze)) -> (unsqueeze) // unsqueeze is a reshape that introduces 1 sized dims
There was a problem hiding this comment.
this reshape feels sketchy, but I'll trust your math here
There was a problem hiding this comment.
that optimization is only true if the permuted dims are not relative to the original dims right?
like you only permute 1-sized dims with any other dim
There was a problem hiding this comment.
yes that's right it has no effect only when permuting 1 sized dims.
This is to support batch group count which is moved in the feature dimension during the reverse, having a static select tablegen operator would be nice for these things:
(StaticSelect (Cond), $a, $b)
would create either $a or $b based on the condition (would also be useful for #176 IIUC).
There was a problem hiding this comment.
there's this https://llvm.org/docs/TableGen/ProgRef.html#if-select-statements-based-on-a-test but i don't how to use it 😅
There was a problem hiding this comment.
Is it okay to use a regular select here? Constant propagation will immediately get rid of it.
There was a problem hiding this comment.
yeah, actually I think this reshape is wrong here because it's not actually a noop. We actually have the optimization I was referring to already, it's called TransposeIsReshape
There was a problem hiding this comment.
I fixed the problem in the reshape/transpose, the operations are completely removed in the case of batchGroupCount == 1 (which will always be the case from NNlib.conv).
#181) * support dilation and feature group count in convolution reverse * support batch group count * fix dimensions for post conv transpose (batch gruop count)
No description provided.