EnzymeAD / EnzymeAD/Enzyme-JAX

Conv with splatted constant LHS/RHS simplification

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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.