GPU Codegen

```bash
./bazel-bin/enzymexlamlir-opt --pass-pipeline="any(inline{default-pipeline=canonicalize max-iterations=4},sroa-wrappers{set_private=false},gpu-launch-recognition,canonicalize,parallel-lower{wrapParallelOps=true},llvm-to-memref-access,polygeist-mem2reg,canonicalize,convert-llvm-to-cf,canonicalize,polygeist-mem2reg,canonicalize,enzyme-lift-cf-to-scf,canonicalize,func.func(canonicalize-loops),canonicalize-scf-for,canonicalize,libdevice-funcs-raise,canonicalize,affine-cfg,canonicalize,func.func(canonicalize-loops),canonicalize,llvm-to-affine-access,canonicalize,delinearize-indexing,canonicalize,simplify-affine-exprs,affine-cfg,canonicalize,llvm-to-affine-access,canonicalize,func.func(affine-loop-invariant-code-motion),canonicalize,sort-memory,raise-affine-to-stablehlo{prefer_while_raising=false dump_failed_lockstep=true},canonicalize,arith-raise{stablehlo=true},symbol-dce,convert-parallel-to-gpu1,gpu-kernel-outlining,canonicalize,convert-parallel-to-gpu2,lower-affine,convert-polygeist-to-llvm)" ./oct.mlir
```

via https://github.com/EnzymeAD/Enzyme-JAX/pull/1022

```mlir
// oct.mlir
module attributes {dlti.dl_spec = #dlti.dl_spec = dense<32> : vector<4xi64>, !llvm.ptr<271> = dense<32> : vector<4xi64>, !llvm.ptr<272> = dense<64> : vector<4xi64>, i64 = dense<64> : vector<2xi64>, i128 = dense<128> : vector<2xi64>, f80 = dense<128> : vector<2xi64>, !llvm.ptr = dense<64> : vector<4xi64>, i1 = dense<8> : vector<2xi64>, i8 = dense<8> : vector<2xi64>, i16 = dense<16> : vector<2xi64>, i32 = dense<32> : vector<2xi64>, f16 = dense<16> : vector<2xi64>, f64 = dense<64> : vector<2xi64>, f128 = dense<128> : vector<2xi64>, "dlti.endianness" = "little", "dlti.mangling_mode" = "e", "dlti.legal_int_widths" = array, "dlti.stack_alignment" = 128 : i64>, llvm.target_triple = "x86_64-unknown-linux-gnu"} {
llvm.mlir.global private unnamed_addr constant @".str"("%f %f\0A\00") {addr_space = 0 : i32, alignment = 1 : i64, dso_local}
llvm.mlir.global external local_unnamed_addr @enzyme_dup(0 : i32) {addr_space = 1 : i32, alignment = 4 : i64, dso_local} : i32
llvm.mlir.global external local_unnamed_addr @enzyme_dupnoneed(0 : i32) {addr_space = 1 : i32, alignment = 4 : i64, dso_local} : i32
llvm.mlir.global external local_unnamed_addr @enzyme_out(0 : i32) {addr_space = 1 : i32, alignment = 4 : i64, dso_local} : i32
llvm.mlir.global external local_unnamed_addr @enzyme_const(0 : i32) {addr_space = 1 : i32, alignment = 4 : i64, dso_local} : i32
llvm.module_flags [#llvm.mlir.module_flag, #llvm.mlir.module_flag, #llvm.mlir.module_flag, #llvm.mlir.module_flag, #llvm.mlir.module_flag, #llvm.mlir.module_flag]
llvm.func local_unnamed_addr @main() -> (i32 {llvm.noundef}) attributes {dso_local, frame_pointer = #llvm.framePointerKind, passthrough = ["mustprogress", "norecurse", ["min-legal-vector-width", "0"], ["no-trapping-math", "true"], ["stack-protector-buffer-size", "8"], ["target-cpu", "x86-64"]], target_cpu = "x86-64", target_features = #llvm.target_features<["+cmov", "+cx8", "+fxsr", "+mmx", "+sse", "+sse2", "+x87"]>, tune_cpu = "generic", uwtable_kind = #llvm.uwtableKind} {
%0 = llvm.mlir.constant(1 : i32) : i32
%1 = llvm.mlir.constant(8 : i64) : i64
%2 = llvm.mlir.constant(1.400000e+00 : f64) : f64
%3 = llvm.mlir.constant(0.000000e+00 : f64) : f64
%4 = llvm.mlir.constant(1.000000e+00 : f64) : f64
%5 = llvm.mlir.constant(4294967297 : i64) : i64
%6 = llvm.mlir.constant(0 : i64) : i64
%7 = llvm.mlir.zero : !llvm.ptr
%8 = llvm.mlir.constant(0 : i32) : i32
%9 = llvm.mlir.constant(16 : i64) : i64
%10 = llvm.mlir.constant(24 : i64) : i64
%11 = llvm.mlir.constant(1 : i64) : i64
%12 = llvm.mlir.constant(2 : i64) : i64
%13 = llvm.mlir.constant(3 : i64) : i64
%14 = llvm.mlir.constant(32 : i64) : i64
%15 = llvm.mlir.addressof @_Z26__device_stub__square_gradPdS_S_S_ : !llvm.ptr
%16 = llvm.mlir.constant(2 : i32) : i32
%17 = llvm.mlir.addressof @".str" : !llvm.ptr
%18 = llvm.alloca %0 x !llvm.ptr {alignment = 8 : i64} : (i32) -> !llvm.ptr
%19 = llvm.alloca %0 x !llvm.ptr {alignment = 8 : i64} : (i32) -> !llvm.ptr
%20 = llvm.alloca %0 x !llvm.ptr {alignment = 8 : i64} : (i32) -> !llvm.ptr
%21 = llvm.alloca %0 x !llvm.ptr {alignment = 8 : i64} : (i32) -> !llvm.ptr
%22 = llvm.alloca %0 x !llvm.struct<"struct.dim3", (i32, i32, i32)> {alignment = 8 : i64} : (i32) -> !llvm.ptr
%23 = llvm.alloca %0 x !llvm.struct<"struct.dim3", (i32, i32, i32)> {alignment = 8 : i64} : (i32) -> !llvm.ptr
%24 = llvm.alloca %0 x i64 {alignment = 8 : i64} : (i32) -> !llvm.ptr
%25 = llvm.alloca %0 x !llvm.ptr {alignment = 8 : i64} : (i32) -> !llvm.ptr
%26 = llvm.alloca %0 x !llvm.array<4 x ptr> {alignment = 16 : i64} : (i32) -> !llvm.ptr
%27 = llvm.alloca %0 x !llvm.ptr {alignment = 8 : i64} : (i32) -> !llvm.ptr
%28 = llvm.alloca %0 x !llvm.ptr {alignment = 8 : i64} : (i32) -> !llvm.ptr
%29 = llvm.alloca %0 x !llvm.ptr {alignment = 8 : i64} : (i32) -> !llvm.ptr
%30 = llvm.alloca %0 x !llvm.ptr {alignment = 8 : i64} : (i32) -> !llvm.ptr
%31 = llvm.alloca %0 x f64 {alignment = 8 : i64} : (i32) -> !llvm.ptr
%32 = llvm.alloca %0 x f64 {alignment = 8 : i64} : (i32) -> !llvm.ptr
%33 = llvm.alloca %0 x f64 {alignment = 8 : i64} : (i32) -> !llvm.ptr
%34 = llvm.alloca %0 x f64 {alignment = 8 : i64} : (i32) -> !llvm.ptr
%35 = llvm.alloca %0 x i64 {alignment = 8 : i64} : (i32) -> !llvm.ptr
%36 = llvm.alloca %0 x i32 {alignment = 4 : i64} : (i32) -> !llvm.ptr
%37 = llvm.alloca %0 x i64 {alignment = 8 : i64} : (i32) -> !llvm.ptr
%38 = llvm.alloca %0 x i32 {alignment = 4 : i64} : (i32) -> !llvm.ptr
%39 = llvm.alloca %0 x i64 {alignment = 8 : i64} : (i32) -> !llvm.ptr
%40 = llvm.alloca %0 x !llvm.ptr {alignment = 8 : i64} : (i32) -> !llvm.ptr
llvm.intr.lifetime.start 8, %27 : !llvm.ptr
llvm.intr.lifetime.start 8, %28 : !llvm.ptr
llvm.intr.lifetime.start 8, %29 : !llvm.ptr
llvm.intr.lifetime.start 8, %30 : !llvm.ptr
%41 = llvm.call @cudaMalloc(%27, %1) : (!llvm.ptr {llvm.nonnull, llvm.noundef}, i64 {llvm.noundef}) -> (i32 {llvm.noundef})
%42 = llvm.call @cudaMalloc(%28, %1) : (!llvm.ptr {llvm.nonnull, llvm.noundef}, i64 {llvm.noundef}) -> (i32 {llvm.noundef})
%43 = llvm.call @cudaMalloc(%29, %1) : (!llvm.ptr {llvm.nonnull, llvm.noundef}, i64 {llvm.noundef}) -> (i32 {llvm.noundef})
%44 = llvm.call @cudaMalloc(%30, %1) : (!llvm.ptr {llvm.nonnull, llvm.noundef}, i64 {llvm.noundef}) -> (i32 {llvm.noundef})
llvm.intr.lifetime.start 8, %31 : !llvm.ptr
llvm.store %2, %31 {alignment = 8 : i64, tbaa = [#llvm.tbaa_tag, 0>}>, 0>}>, access_type = , 0>}>, 0>}>, offset = 0>]} : f64, !llvm.ptr
llvm.intr.lifetime.start 8, %32 : !llvm.ptr
llvm.store %3, %32 {alignment = 8 : i64, tbaa = [#llvm.tbaa_tag, 0>}>, 0>}>, access_type = , 0>}>, 0>}>, offset = 0>]} : f64, !llvm.ptr
llvm.intr.lifetime.start 8, %33 : !llvm.ptr
llvm.store %4, %33 {alignment = 8 : i64, tbaa = [#llvm.tbaa_tag, 0>}>, 0>}>, access_type = , 0>}>, 0>}>, offset = 0>]} : f64, !llvm.ptr
llvm.intr.lifetime.start 8, %34 : !llvm.ptr
llvm.store %4, %34 {alignment = 8 : i64, tbaa = [#llvm.tbaa_tag, 0>}>, 0>}>, access_type = , 0>}>, 0>}>, offset = 0>]} : f64, !llvm.ptr
%45 = llvm.load %27 {alignment = 8 : i64, tbaa = [#llvm.tbaa_tag, 0>}>, 0>}>, 0>}>, access_type = , 0>}>, 0>}>, 0>}>, offset = 0>]} : !llvm.ptr -> !llvm.ptr
%46 = llvm.call @cudaMemcpy(%45, %31, %1, %0) : (!llvm.ptr {llvm.noundef}, !llvm.ptr {llvm.nonnull, llvm.noundef}, i64 {llvm.noundef}, i32 {llvm.noundef}) -> i32
%47 = llvm.load %28 {alignment = 8 : i64, tbaa = [#llvm.tbaa_tag, 0>}>, 0>}>, 0>}>, access_type = , 0>}>, 0>}>, 0>}>, offset = 0>]} : !llvm.ptr -> !llvm.ptr
%48 = llvm.call @cudaMemcpy(%47, %32, %1, %0) : (!llvm.ptr {llvm.noundef}, !llvm.ptr {llvm.nonnull, llvm.noundef}, i64 {llvm.noundef}, i32 {llvm.noundef}) -> i32
%49 = llvm.load %29 {alignment = 8 : i64, tbaa = [#llvm.tbaa_tag, 0>}>, 0>}>, 0>}>, access_type = , 0>}>, 0>}>, 0>}>, offset = 0>]} : !llvm.ptr -> !llvm.ptr
%50 = llvm.call @cudaMemcpy(%49, %33, %1, %0) : (!llvm.ptr {llvm.noundef}, !llvm.ptr {llvm.nonnull, llvm.noundef}, i64 {llvm.noundef}, i32 {llvm.noundef}) -> i32
%51 = llvm.load %30 {alignment = 8 : i64, tbaa = [#llvm.tbaa_tag, 0>}>, 0>}>, 0>}>, access_type = , 0>}>, 0>}>, 0>}>, offset = 0>]} : !llvm.ptr -> !llvm.ptr
%52 = llvm.call @cudaMemcpy(%51, %34, %1, %0) : (!llvm.ptr {llvm.noundef}, !llvm.ptr {llvm.nonnull, llvm.noundef}, i64 {llvm.noundef}, i32 {llvm.noundef}) -> i32
llvm.store %5, %35 {alignment = 8 : i64} : i64, !llvm.ptr
llvm.store %0, %36 {alignment = 4 : i64} : i32, !llvm.ptr
llvm.store %5, %37 {alignment = 8 : i64} : i64, !llvm.ptr
llvm.store %0, %38 {alignment = 4 : i64} : i32, !llvm.ptr
llvm.store %6, %39 {alignment = 8 : i64} : i64, !llvm.ptr
llvm.store %7, %40 {alignment = 8 : i64} : !llvm.ptr, !llvm.ptr
%53 = llvm.icmp "eq" %8, %8 : i32
llvm.cond_br %53, ^bb1, ^bb2
^bb1: // pred: ^bb0
%54 = llvm.load %27 {alignment = 8 : i64, tbaa = [#llvm.tbaa_tag, 0>}>, 0>}>, 0>}>, access_type = , 0>}>, 0>}>, 0>}>, offset = 0>]} : !llvm.ptr -> !llvm.ptr
%55 = llvm.load %28 {alignment = 8 : i64, tbaa = [#llvm.tbaa_tag, 0>}>, 0>}>, 0>}>, access_type = , 0>}>, 0>}>, 0>}>, offset = 0>]} : !llvm.ptr -> !llvm.ptr
%56 = llvm.load %29 {alignment = 8 : i64, tbaa = [#llvm.tbaa_tag, 0>}>, 0>}>, 0>}>, access_type = , 0>}>, 0>}>, 0>}>, offset = 0>]} : !llvm.ptr -> !llvm.ptr
%57 = llvm.load %30 {alignment = 8 : i64, tbaa = [#llvm.tbaa_tag, 0>}>, 0>}>, 0>}>, access_type = , 0>}>, 0>}>, 0>}>, offset = 0>]} : !llvm.ptr -> !llvm.ptr
llvm.intr.lifetime.start 8, %18 : !llvm.ptr
llvm.intr.lifetime.start 8, %19 : !llvm.ptr
llvm.intr.lifetime.start 8, %20 : !llvm.ptr
llvm.intr.lifetime.start 8, %21 : !llvm.ptr
llvm.intr.lifetime.start 12, %22 : !llvm.ptr
llvm.intr.lifetime.start 12, %23 : !llvm.ptr
llvm.intr.lifetime.start 8, %24 : !llvm.ptr
llvm.intr.lifetime.start 8, %25 : !llvm.ptr
llvm.intr.lifetime.start 32, %26 : !llvm.ptr
llvm.store %54, %18 {alignment = 8 : i64, tbaa = [#llvm.tbaa_tag, 0>}>, 0>}>, 0>}>, access_type = , 0>}>, 0>}>, 0>}>, offset = 0>]} : !llvm.ptr, !llvm.ptr
llvm.store %55, %19 {alignment = 8 : i64, tbaa = [#llvm.tbaa_tag, 0>}>, 0>}>, 0>}>, access_type = , 0>}>, 0>}>, 0>}>, offset = 0>]} : !llvm.ptr, !llvm.ptr
llvm.store %56, %20 {alignment = 8 : i64, tbaa = [#llvm.tbaa_tag, 0>}>, 0>}>, 0>}>, access_type = , 0>}>, 0>}>, 0>}>, offset = 0>]} : !llvm.ptr, !llvm.ptr
llvm.store %57, %21 {alignment = 8 : i64, tbaa = [#llvm.tbaa_tag, 0>}>, 0>}>, 0>}>, access_type = , 0>}>, 0>}>, 0>}>, offset = 0>]} : !llvm.ptr, !llvm.ptr
llvm.store %18, %26 {alignment = 8 : i64} : !llvm.ptr, !llvm.ptr
%58 = llvm.getelementptr inbounds|nuw %26[%1] : (!llvm.ptr, i64) -> !llvm.ptr, i8
llvm.store %19, %58 {alignment = 8 : i64} : !llvm.ptr, !llvm.ptr
%59 = llvm.getelementptr inbounds|nuw %26[%9] : (!llvm.ptr, i64) -> !llvm.ptr, i8
llvm.store %20, %59 {alignment = 8 : i64} : !llvm.ptr, !llvm.ptr
%60 = llvm.getelementptr inbounds|nuw %26[%10] : (!llvm.ptr, i64) -> !llvm.ptr, i8
llvm.store %21, %60 {alignment = 8 : i64} : !llvm.ptr, !llvm.ptr
%61 = llvm.load %35 {alignment = 8 : i64} : !llvm.ptr -> i64
%62 = llvm.load %36 {alignment = 4 : i64} : !llvm.ptr -> i32
%63 = llvm.load %37 {alignment = 8 : i64} : !llvm.ptr -> i64
%64 = llvm.load %38 {alignment = 4 : i64} : !llvm.ptr -> i32
%65 = llvm.load %39 {alignment = 8 : i64} : !llvm.ptr -> i64
%66 = llvm.load %40 {alignment = 8 : i64} : !llvm.ptr -> !llvm.ptr
%67 = llvm.load %24 {alignment = 8 : i64} : !llvm.ptr -> i64
%68 = llvm.load %25 {alignment = 8 : i64} : !llvm.ptr -> !llvm.ptr
%69 = llvm.load %22 {alignment = 8 : i64} : !llvm.ptr -> i64
%70 = llvm.getelementptr inbounds|nuw %22[%1] : (!llvm.ptr, i64) -> !llvm.ptr, i8
%71 = llvm.load %70 {alignment = 8 : i64} : !llvm.ptr -> i32
%72 = llvm.load %23 {alignment = 8 : i64} : !llvm.ptr -> i64
%73 = llvm.getelementptr inbounds|nuw %23[%1] : (!llvm.ptr, i64) -> !llvm.ptr, i8
%74 = llvm.load %73 {alignment = 8 : i64} : !llvm.ptr -> i32
%75 = llvm.getelementptr inbounds %26[%6] : (!llvm.ptr, i64) -> !llvm.ptr, !llvm.ptr
%76 = llvm.load %75 {alignment = 8 : i64} : !llvm.ptr -> !llvm.ptr
%77 = llvm.load %76 {alignment = 8 : i64} : !llvm.ptr -> !llvm.ptr
%78 = llvm.getelementptr inbounds %26[%11] : (!llvm.ptr, i64) -> !llvm.ptr, !llvm.ptr
%79 = llvm.load %78 {alignment = 8 : i64} : !llvm.ptr -> !llvm.ptr
%80 = llvm.load %79 {alignment = 8 : i64} : !llvm.ptr -> !llvm.ptr
%81 = llvm.getelementptr inbounds %26[%12] : (!llvm.ptr, i64) -> !llvm.ptr, !llvm.ptr
%82 = llvm.load %81 {alignment = 8 : i64} : !llvm.ptr -> !llvm.ptr
%83 = llvm.load %82 {alignment = 8 : i64} : !llvm.ptr -> !llvm.ptr
%84 = llvm.getelementptr inbounds %26[%13] : (!llvm.ptr, i64) -> !llvm.ptr, !llvm.ptr
%85 = llvm.load %84 {alignment = 8 : i64} : !llvm.ptr -> !llvm.ptr
%86 = llvm.load %85 {alignment = 8 : i64} : !llvm.ptr -> !llvm.ptr
%87 = llvm.trunc %61 : i64 to i32
%88 = llvm.lshr %61, %14 : i64
%89 = llvm.trunc %88 : i64 to i32
%90 = llvm.trunc %63 : i64 to i32
%91 = llvm.lshr %63, %14 : i64
%92 = llvm.trunc %91 : i64 to i32
llvm.call @__mlir_launch_kernel__Z26__device_stub__square_gradPdS_S_S_(%15, %87, %89, %62, %90, %92, %64, %65, %66, %77, %80, %83, %86) : (!llvm.ptr, i32, i32, i32, i32, i32, i32, i64, !llvm.ptr, !llvm.ptr, !llvm.ptr, !llvm.ptr, !llvm.ptr) -> ()
llvm.intr.lifetime.end 8, %18 : !llvm.ptr
llvm.intr.lifetime.end 8, %19 : !llvm.ptr
llvm.intr.lifetime.end 8, %20 : !llvm.ptr
llvm.intr.lifetime.end 8, %21 : !llvm.ptr
llvm.intr.lifetime.end 12, %22 : !llvm.ptr
llvm.intr.lifetime.end 12, %23 : !llvm.ptr
llvm.intr.lifetime.end 8, %24 : !llvm.ptr
llvm.intr.lifetime.end 8, %25 : !llvm.ptr
llvm.intr.lifetime.end 32, %26 : !llvm.ptr
llvm.br ^bb2
^bb2: // 2 preds: ^bb0, ^bb1
%93 = llvm.call @cudaDeviceSynchronize() : () -> i32
%94 = llvm.load %27 {alignment = 8 : i64, tbaa = [#llvm.tbaa_tag, 0>}>, 0>}>, 0>}>, access_type = , 0>}>, 0>}>, 0>}>, offset = 0>]} : !llvm.ptr -> !llvm.ptr
%95 = llvm.call @cudaMemcpy(%31, %94, %1, %16) : (!llvm.ptr {llvm.nonnull, llvm.noundef}, !llvm.ptr {llvm.noundef}, i64 {llvm.noundef}, i32 {llvm.noundef}) -> i32
%96 = llvm.load %28 {alignment = 8 : i64, tbaa = [#llvm.tbaa_tag, 0>}>, 0>}>, 0>}>, access_type = , 0>}>, 0>}>, 0>}>, offset = 0>]} : !llvm.ptr -> !llvm.ptr
%97 = llvm.call @cudaMemcpy(%32, %96, %1, %16) : (!llvm.ptr {llvm.nonnull, llvm.noundef}, !llvm.ptr {llvm.noundef}, i64 {llvm.noundef}, i32 {llvm.noundef}) -> i32
%98 = llvm.load %29 {alignment = 8 : i64, tbaa = [#llvm.tbaa_tag, 0>}>, 0>}>, 0>}>, access_type = , 0>}>, 0>}>, 0>}>, offset = 0>]} : !llvm.ptr -> !llvm.ptr
%99 = llvm.call @cudaMemcpy(%33, %98, %1, %16) : (!llvm.ptr {llvm.nonnull, llvm.noundef}, !llvm.ptr {llvm.noundef}, i64 {llvm.noundef}, i32 {llvm.noundef}) -> i32
%100 = llvm.load %30 {alignment = 8 : i64, tbaa = [#llvm.tbaa_tag, 0>}>, 0>}>, 0>}>, access_type = , 0>}>, 0>}>, 0>}>, offset = 0>]} : !llvm.ptr -> !llvm.ptr
%101 = llvm.call @cudaMemcpy(%34, %100, %1, %16) : (!llvm.ptr {llvm.nonnull, llvm.noundef}, !llvm.ptr {llvm.noundef}, i64 {llvm.noundef}, i32 {llvm.noundef}) -> i32
%102 = llvm.load %31 {alignment = 8 : i64, tbaa = [#llvm.tbaa_tag, 0>}>, 0>}>, access_type = , 0>}>, 0>}>, offset = 0>]} : !llvm.ptr -> f64
%103 = llvm.load %33 {alignment = 8 : i64, tbaa = [#llvm.tbaa_tag, 0>}>, 0>}>, access_type = , 0>}>, 0>}>, offset = 0>]} : !llvm.ptr -> f64
%104 = llvm.call @printf(%17, %102, %103) vararg(!llvm.func) : (!llvm.ptr {llvm.dereferenceable = 1 : i64, llvm.nonnull, llvm.noundef}, f64 {llvm.noundef}, f64 {llvm.noundef}) -> i32
%105 = llvm.load %32 {alignment = 8 : i64, tbaa = [#llvm.tbaa_tag, 0>}>, 0>}>, access_type = , 0>}>, 0>}>, offset = 0>]} : !llvm.ptr -> f64
%106 = llvm.load %34 {alignment = 8 : i64, tbaa = [#llvm.tbaa_tag, 0>}>, 0>}>, access_type = , 0>}>, 0>}>, offset = 0>]} : !llvm.ptr -> f64
%107 = llvm.call @printf(%17, %105, %106) vararg(!llvm.func) : (!llvm.ptr {llvm.dereferenceable = 1 : i64, llvm.nonnull, llvm.noundef}, f64 {llvm.noundef}, f64 {llvm.noundef}) -> i32
llvm.intr.lifetime.end 8, %34 : !llvm.ptr
llvm.intr.lifetime.end 8, %33 : !llvm.ptr
llvm.intr.lifetime.end 8, %32 : !llvm.ptr
llvm.intr.lifetime.end 8, %31 : !llvm.ptr
llvm.intr.lifetime.end 8, %30 : !llvm.ptr
llvm.intr.lifetime.end 8, %29 : !llvm.ptr
llvm.intr.lifetime.end 8, %28 : !llvm.ptr
llvm.intr.lifetime.end 8, %27 : !llvm.ptr
llvm.return %8 : i32
}
llvm.func local_unnamed_addr @cudaMemcpy(!llvm.ptr {llvm.noundef}, !llvm.ptr {llvm.noundef}, i64 {llvm.noundef}, i32 {llvm.noundef}) -> i32 attributes {frame_pointer = #llvm.framePointerKind, passthrough = [["no-trapping-math", "true"], ["stack-protector-buffer-size", "8"], ["target-cpu", "x86-64"]], target_cpu = "x86-64", target_features = #llvm.target_features<["+cmov", "+cx8", "+fxsr", "+mmx", "+sse", "+sse2", "+x87"]>, tune_cpu = "generic"}
llvm.func local_unnamed_addr @cudaDeviceSynchronize() -> i32 attributes {frame_pointer = #llvm.framePointerKind, passthrough = [["no-trapping-math", "true"], ["stack-protector-buffer-size", "8"], ["target-cpu", "x86-64"]], target_cpu = "x86-64", target_features = #llvm.target_features<["+cmov", "+cx8", "+fxsr", "+mmx", "+sse", "+sse2", "+x87"]>, tune_cpu = "generic"}
llvm.func local_unnamed_addr @printf(!llvm.ptr {llvm.nocapture, llvm.noundef, llvm.readonly}, ...) -> (i32 {llvm.noundef}) attributes {frame_pointer = #llvm.framePointerKind, no_unwind, passthrough = ["nofree", ["no-trapping-math", "true"], ["stack-protector-buffer-size", "8"], ["target-cpu", "x86-64"]], target_cpu = "x86-64", target_features = #llvm.target_features<["+cmov", "+cx8", "+fxsr", "+mmx", "+sse", "+sse2", "+x87"]>, tune_cpu = "generic"}
llvm.func local_unnamed_addr @cudaMalloc(!llvm.ptr {llvm.noundef}, i64 {llvm.noundef}) -> i32 attributes {frame_pointer = #llvm.framePointerKind, passthrough = [["no-trapping-math", "true"], ["stack-protector-buffer-size", "8"], ["target-cpu", "x86-64"]], target_cpu = "x86-64", target_features = #llvm.target_features<["+cmov", "+cx8", "+fxsr", "+mmx", "+sse", "+sse2", "+x87"]>, tune_cpu = "generic"}
llvm.func local_unnamed_addr @__mlir_launch_kernel__Z26__device_stub__square_gradPdS_S_S_(!llvm.ptr, i32, i32, i32, i32, i32, i32, i64, !llvm.ptr, !llvm.ptr, !llvm.ptr, !llvm.ptr, !llvm.ptr)
llvm.func internal @_Z26__device_stub__square_gradPdS_S_S_(%arg0: !llvm.ptr {llvm.nocapture, llvm.noundef, llvm.readonly}, %arg1: !llvm.ptr {llvm.nocapture, llvm.noundef, llvm.readnone}, %arg2: !llvm.ptr {llvm.nocapture, llvm.noundef}, %arg3: !llvm.ptr {llvm.nocapture, llvm.noundef, llvm.readnone}) attributes {dso_local, frame_pointer = #llvm.framePointerKind, memory_effects = #llvm.memory_effects, no_unwind, passthrough = ["mustprogress", "nofree", "norecurse", "nosync", ["no-trapping-math", "true"], ["stack-protector-buffer-size", "8"], ["target-cpu", "sm_52"], ["uniform-work-group-size", "true"]], target_cpu = "sm_52", target_features = #llvm.target_features<["+ptx87", "+sm_52"]>, will_return} {
%0 = llvm.load %arg0 {alignment = 8 : i64, tbaa = [#llvm.tbaa_tag, 0>}>, 0>}>, access_type = , 0>}>, 0>}>, offset = 0>]} : !llvm.ptr -> f64
%1 = llvm.load %arg2 {alignment = 8 : i64, tbaa = [#llvm.tbaa_tag, 0>}>, 0>}>, access_type = , 0>}>, 0>}>, offset = 0>]} : !llvm.ptr -> f64
%2 = llvm.fadd %0, %1 {fastmathFlags = #llvm.fastmath} : f64
llvm.store %2, %arg2 {alignment = 8 : i64, tbaa = [#llvm.tbaa_tag, 0>}>, 0>}>, access_type = , 0>}>, 0>}>, offset = 0>]} : f64, !llvm.ptr
llvm.return
}
}
```