EnzymeAD / EnzymeAD/Enzyme-JAX

Slicing to conv

Open
#628 9 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> @code_hlo raise=true simple_stencil!(A, B)
module @reactant_simple_... attributes {mhlo.num_partitions = 1 : i64, mhlo.num_replicas = 1 : i64} {
llvm.module_flags [#llvm.mlir.module_flag, #llvm.mlir.module_flag]
func.func @main(%arg0: tensor<1092xf64> {tf.aliasing_output = 0 : i32}, %arg1: tensor<1095xf64>) -> tensor<1092xf64> {
%0 = stablehlo.slice %arg1 [1:1092] : (tensor<1095xf64>) -> tensor<1091xf64>
%1 = stablehlo.slice %arg1 [0:1091] : (tensor<1095xf64>) -> tensor<1091xf64>
%2 = stablehlo.slice %arg1 [2:1093] : (tensor<1095xf64>) -> tensor<1091xf64>
%3 = stablehlo.add %1, %0 : tensor<1091xf64>
%4 = stablehlo.add %3, %2 : tensor<1091xf64>
%5 = stablehlo.slice %arg0 [0:1] : (tensor<1092xf64>) -> tensor<1xf64>
%6 = stablehlo.concatenate %5, %4, dim = 0 : (tensor<1xf64>, tensor<1091xf64>) -> tensor<1092xf64>
return %6 : tensor<1092xf64>
}
}
```

```
using CUDA

function stencil_kernel!(y, x)
i = threadIdx().x + (blockIdx().x - 1) * blockDim().x
if i <= length(x)
y[i] = x[i-1] + x[i] + x[i+1]
end
return nothing
end

function simple_stencil!(y, x)
@cuda blocks = ((length(x)+255) ÷ 256) threads=256 stencil_kernel!(y, x)
return nothing
end

using Reactant
A = Reactant.to_rarray(ones(1092))
B = Reactant.to_rarray(ones(1092+3))

@code_hlo raise=true simple_stencil!(A, B)
```

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.