EnzymeAD / EnzymeAD/Enzyme-JAX

raise scan-like operations

Open
#1,850 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

```julia
using Reactant

function looped_accumulate(x)
@trace for i in 3:length(x)
x[i] += x[i - 1]
end
return x
end

x = Reactant.to_rarray(rand(Float32, 128));

@code_hlo looped_accumulate(x)

function looped_accumulate(x)
@trace for i in 3:length(x)
x[i] += 2 * x[i - 1]
end
return x
end

x = Reactant.to_rarray(rand(Float32, 128));

@code_hlo looped_accumulate(x)
```

```mlir
module @reactant_looped_... attributes {mhlo.num_partitions = 1 : i64, mhlo.num_replicas = 1 : i64} {
func.func @main(%arg0: tensor<128xf32> {enzymexla.memory_effects = [], tf.aliasing_output = 0 : i32}) -> tensor<128xf32> attributes {enzymexla.memory_effects = []} {
%cst = stablehlo.constant dense<2.000000e+00> : tensor<1xf32>
%c = stablehlo.constant dense<1> : tensor
%c_0 = stablehlo.constant dense<0> : tensor
%c_1 = stablehlo.constant dense<1> : tensor
%c_2 = stablehlo.constant dense<3> : tensor
%c_3 = stablehlo.constant dense<126> : tensor
%0:2 = stablehlo.while(%iterArg = %c_0, %iterArg_4 = %arg0) : tensor, tensor<128xf32> attributes {enzyme.disable_mincut}
cond {
%1 = stablehlo.compare LT, %iterArg, %c_3 : (tensor, tensor) -> tensor
stablehlo.return %1 : tensor
} do {
%1 = stablehlo.add %c_2, %iterArg {enzymexla.bounds = [[3, 128]]} : tensor
%2 = stablehlo.add %iterArg, %c_1 {enzymexla.bounds = [[1, 126]]} : tensor
%3 = stablehlo.convert %1 {enzymexla.bounds = [[3, 128]]} : (tensor) -> tensor
%4 = stablehlo.subtract %3, %c {enzymexla.bounds = [[2, 127]]} : tensor
%5 = stablehlo.dynamic_slice %arg0, %4, sizes = [1] : (tensor<128xf32>, tensor) -> tensor<1xf32>
%6 = stablehlo.subtract %1, %c_1 {enzymexla.bounds = [[2, 127]]} : tensor
%7 = stablehlo.convert %6 {enzymexla.bounds = [[2, 127]]} : (tensor) -> tensor
%8 = stablehlo.subtract %7, %c {enzymexla.bounds = [[1, 126]]} : tensor
%9 = stablehlo.dynamic_slice %iterArg_4, %8, sizes = [1] : (tensor<128xf32>, tensor) -> tensor<1xf32>
%10 = stablehlo.multiply %cst, %9 : tensor<1xf32>
%11 = stablehlo.add %5, %10 : tensor<1xf32>
%12 = stablehlo.dynamic_update_slice %iterArg_4, %11, %4 : (tensor<128xf32>, tensor<1xf32>, tensor) -> tensor<128xf32>
stablehlo.return %2, %12 : tensor, tensor<128xf32>
}
return %0#1 : tensor<128xf32>
}
}
```

Contributor guide

No contributing guide indexed for this repository

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.