EnzymeAD / EnzymeAD/Enzyme-JAX
Kernel SROA of arguments
- Dominant language
- MLIR
- Stars
- 131
- Forks
- 53
- Avg merge
- 1d 10h
- Merged PRs (30d)
- 193
Description
cc @ftynse @ivanradanov
```mlir
// /home/wmoses/git/Reactant.jl/cu.jl:22 =# @code_hlo(optimize = false, tuplef2(A)) =
#tbaa_root = #llvm.tbaa_root
#tbaa_type_desc = #llvm.tbaa_type_desc}>
#tbaa_tag = #llvm.tbaa_tag
module {
llvm.func local_unnamed_addr @_Z8tuplef2_5TupleI5Int6413CuTracedArrayIS0_Li0ELi1E2__EE(%arg0: !llvm.struct<(i64, array<1 x ptr<1>>)>) attributes {sym_visibility = "private"} {
%0 = llvm.extractvalue %arg0[0] : !llvm.struct<(i64, array<1 x ptr<1>>)>
%1 = llvm.extractvalue %arg0[1] : !llvm.struct<(i64, array<1 x ptr<1>>)>
%2 = llvm.extractvalue %1[0] : !llvm.array<1 x ptr<1>>
%3 = llvm.bitcast %2 : !llvm.ptr<1> to !llvm.ptr<1>
%4 = llvm.load %3 {alignment = 1 : i64, tbaa = [#tbaa_tag]} : !llvm.ptr<1> -> i64
%5 = llvm.mul %4, %0 : i64
llvm.store %5, %3 {alignment = 1 : i64, tbaa = [#tbaa_tag]} : i64, !llvm.ptr<1>
llvm.return
}
llvm.func ptx_kernelcc @"##call__Z8tuplef2_5TupleI5Int6413CuTracedArrayIS0_Li0ELi1E2__EE#258"(%arg0: !llvm.array<1 x ptr<1>>) attributes {sym_visibility = "private"} {
%0 = llvm.mlir.constant(1 : i64) : i64
%1 = llvm.alloca %0 x !llvm.struct<(i64, array<1 x ptr<1>>)> : (i64) -> !llvm.ptr
%2 = llvm.mlir.constant(dense<[5, 0, 0, 0, 0, 0, 0, 0, 112, 231, 165, 87, 9, 117, 0, 0]> : tensor<16xui8>) : !llvm.array<16 x i8>
llvm.store %2, %1 : !llvm.array<16 x i8>, !llvm.ptr
%3 = llvm.getelementptr %1[8] : (!llvm.ptr) -> !llvm.ptr, ui8
llvm.store %arg0, %3 : !llvm.array<1 x ptr<1>>, !llvm.ptr
%4 = llvm.load %1 : !llvm.ptr -> !llvm.struct<(i64, array<1 x ptr<1>>)>
llvm.call @_Z8tuplef2_5TupleI5Int6413CuTracedArrayIS0_Li0ELi1E2__EE(%4) : (!llvm.struct<(i64, array<1 x ptr<1>>)>) -> ()
llvm.return
}
func.func @main(%arg0: tensor) -> tensor {
%c = stablehlo.constant dense<1> : tensor
%c_0 = stablehlo.constant dense<1> : tensor
%c_1 = stablehlo.constant dense<1> : tensor
%c_2 = stablehlo.constant dense<1> : tensor
%c_3 = stablehlo.constant dense<1> : tensor
%c_4 = stablehlo.constant dense<1> : tensor
%c_5 = stablehlo.constant dense<0> : tensor
%0 = enzymexla.kernel_call @"##call__Z8tuplef2_5TupleI5Int6413CuTracedArrayIS0_Li0ELi1E2__EE#258" blocks in(%c, %c_0, %c_1) threads in(%c_2, %c_3, %c_4) shmem = %c_5 (%arg0) {output_operand_aliases = [#stablehlo.output_operand_alias]} : (tensor) -> tensor
return %0 : tensor
}
}
#= /home/wmoses/git/Reactant.jl/cu.jl:23 =# @code_hlo(tuplef2(A)) = module {
func.func @main(%arg0: tensor) -> tensor {
%c = stablehlo.constant dense<1> : tensor
%c_0 = stablehlo.constant dense<0> : tensor
%0 = stablehlo.custom_call @enzymexla_compile_gpu(%arg0) {api_version = 4 : i32, backend_config = {attr = "P` \05\12u\00\00\00` \05\12u\00\00"}, output_operand_aliases = [#stablehlo.output_operand_alias]} : (tensor) -> tensor
return %0 : tensor
}
}
```
Contributor guide
No contributing guide indexed for this repository
Assessment
This issue has not been assessed yet.