EnzymeAD / EnzymeAD/Enzyme-JAX
Conv with splatted constant LHS/RHS simplification
- Dominant language
- MLIR
- Stars
- 131
- Forks
- 53
- Avg merge
- 1d 10h
- Merged PRs (30d)
- 193
Description
```mlir
module @reactant_enzyme_... attributes {mhlo.num_partitions = 1 : i64, mhlo.num_replicas = 1 : i64} {
func.func @main(%arg0: tensor<8x128x15004xf32> {enzymexla.memory_effects = []}, %arg1: tensor<128x2x160xf32> {enzymexla.memory_effects = []}) -> (tensor<8x128x15004xf32>, tensor<128x2x160xf32>) attributes {enzymexla.memory_effects = []} {
%cst = stablehlo.constant dense<1.000000e+00> : tensor<8x2x480000xf32>
%0 = stablehlo.reverse %arg1, dims = [2] : tensor<128x2x160xf32>
%1 = stablehlo.reverse %0, dims = [2] : tensor<128x2x160xf32>
%2 = stablehlo.convolution(%cst, %1) dim_numbers = [b, f, 0]x[o, i, 0]->[b, f, 0], window = {stride = [32], pad = [[128, 128]], rhs_dilate = [1]} {batch_group_count = 1 : i64, feature_group_count = 1 : i64, precision_config = [#stablehlo, #stablehlo]} : (tensor<8x2x480000xf32>, tensor<128x2x160xf32>) -> tensor<8x128x15004xf32>
%3 = stablehlo.convolution(%arg0, %cst) dim_numbers = [f, b, 0]x[i, o, 0]->[b, f, 0], window = {stride = [1], pad = [[31, 31]], lhs_dilate = [32]} {batch_group_count = 1 : i64, feature_group_count = 1 : i64, precision_config = [#stablehlo, #stablehlo]} : (tensor<8x128x15004xf32>, tensor<8x2x480000xf32>) -> tensor<128x2x160xf32>
%4 = stablehlo.reverse %3, dims = [2] : tensor<128x2x160xf32>
return %2, %4 : tensor<8x128x15004xf32>, tensor<128x2x160xf32>
}
}
```
Contributor guide
No contributing guide indexed for this repository
Research direction
The issue provides only an MLIR reproducer containing two stablehlo.convolution operations and reverse operations; no files or tests are named. Start by locating the convolution simplification entry point and determine the expected handling of the splatted constant on each side. Done should include a regression test for both convolution forms and the intended simplified result.
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
- Needs clarification
- Newbie friendliness
- 25/100