EnzymeAD / EnzymeAD/Enzyme-JAX
convolution with batch_group_count -> regular convolution
- Dominant language
- MLIR
- Stars
- 133
- Forks
- 53
- Avg merge
- 1d 10h
- Merged PRs (30d)
- 193
Description
If the kernel is repeated along the kernel_output_dim, then we can drop the batch group count and permute the output to get the desired output
```
module {
func.func @main(%arg0: tensor<6x3x2x5x4x4xf32>) -> tensor<6x3x7x5x2x2xf32> {
%cst = stablehlo.constant dense<3.000000e+00> : tensor<3x3x2x210xf32>
%0 = stablehlo.transpose %arg0, dims = [5, 4, 2, 1, 3, 0] : (tensor<6x3x2x5x4x4xf32>) -> tensor<4x4x2x3x5x6xf32>
%1 = stablehlo.reshape %0 : (tensor<4x4x2x3x5x6xf32>) -> tensor<4x4x2x90xf32>
%2 = stablehlo.convolution(%1, %cst) dim_numbers = [0, 1, f, b]x[0, 1, i, o]->[0, 1, f, b], window = {stride = [1, 1], pad = [[0, 0], [0, 0]], rhs_dilate = [1, 1]} {batch_group_count = 30 : i64, feature_group_count = 1 : i64, precision_config = [#stablehlo, #stablehlo]} : (tensor<4x4x2x90xf32>, tensor<3x3x2x210xf32>) -> tensor<2x2x210x3xf32>
%3 = stablehlo.reshape %2 : (tensor<2x2x210x3xf32>) -> tensor<2x2x7x5x6x3xf32>
%4 = stablehlo.transpose %3, dims = [4, 5, 2, 3, 1, 0] : (tensor<2x2x7x5x6x3xf32>) -> tensor<6x3x7x5x2x2xf32>
return %4 : tensor<6x3x7x5x2x2xf32>
}
}
```
Contributor guide
No contributing guide indexed for this repository
Research direction
No source files, entry points, or tests are named. Start by locating the compiler handling for StableHLO convolution operations with batch_group_count, then use the supplied MLIR module as the reproducer; done means the equivalent regular-convolution form preserves the shown output shape and semantics.
Written by the indexing model from the issue text.
Assessment
- Domain
- compilers
- Issue type
- Feature
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Stale
- Clarity
- Mostly clear
- Newbie friendliness
- 35/100