EnzymeAD / EnzymeAD/Enzyme-JAX
New optimization pattern
- Dominant language
- MLIR
- Stars
- 131
- Forks
- 53
- Avg merge
- 1d 10h
- Merged PRs (30d)
- 193
Description
@wsmoses
From running `-enzyme-diff-batch` on https://github.com/EnzymeAD/Reactant.jl/blob/vim/brusselator/benchmark/brusselator/workload.jl
```
// RUN: enzymexlamlir-opt %s --pass-pipeline="builtin.module(enzyme-hlo-generate-td{patterns=slice_elementwise<1>;slice_concat<1>;noop_slice<16>},transform-interpreter,enzyme-hlo-remove-transform,canonicalize,cse)" | FileCheck %s
module {
func.func @full_coverage_slice_elementwise(
%lane0: tensor<1x4x4xf32>,
%lane1: tensor<1x4x4xf32>)
-> (tensor<1x4x4xf32>, tensor<1x4x4xf32>) {
%packed = stablehlo.concatenate %lane0, %lane1, dim = 0
: (tensor<1x4x4xf32>, tensor<1x4x4xf32>) -> tensor<2x4x4xf32>
%sum = stablehlo.add %packed, %packed : tensor<2x4x4xf32>
%result0 = stablehlo.slice %sum [0:1, 0:4, 0:4]
: (tensor<2x4x4xf32>) -> tensor<1x4x4xf32>
%result1 = stablehlo.slice %sum [1:2, 0:4, 0:4]
: (tensor<2x4x4xf32>) -> tensor<1x4x4xf32>
return %result0, %result1
: tensor<1x4x4xf32>, tensor<1x4x4xf32>
}
}
// CHECK-LABEL: func.func @full_coverage_slice_elementwise
// CHECK-NOT: stablehlo.concatenate
// CHECK-NOT: stablehlo.slice
// CHECK: %[[R0:.*]] = stablehlo.add %arg0, %arg0 : tensor<1x4x4xf32>
// CHECK: %[[R1:.*]] = stablehlo.add %arg1, %arg1 : tensor<1x4x4xf32>
// CHECK: return %[[R0]], %[[R1]]
```
Contributor guide
No contributing guide indexed for this repository
Research direction
Start with benchmark/brusselator/workload.jl and run the shown -enzyme-diff-batch pass pipeline using enzyme-hlo-generate-td with the slice_elementwise, slice_concat, and noop_slice patterns. Check the FileCheck expectations for full_coverage_slice_elementwise: the transformed output should contain two stablehlo.add operations and no stablehlo.concatenate or stablehlo.slice operations.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- julia
- Domain
- compilers
- Issue type
- Feature
- Difficulty
- 3/5
- Estimated time
- 1-2 days
- Activity status
- Quiet
- Clarity
- Mostly clear
- Newbie friendliness
- 65/100