EnzymeAD / EnzymeAD/Enzyme-JAX

convolution with batch_group_count -> regular convolution

Open
#1,572 0 comments 0 reactions 0 assignees View on GitHub
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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.