EnzymeAD / EnzymeAD/Enzyme-JAX
Slicing to conv
- 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.