EnzymeAD / EnzymeAD/Enzyme-JAX

New optimization pattern

Open
#2,622 3 comments 0 reactions 0 assignees View on GitHub
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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.