import gpu_cpp as gpu
import numpy as np

ctx = gpu.create_context()

N = 12

input = gpu.create_tensor(ctx, [N], gpu.kf32)
output = gpu.create_tensor(ctx, [N], gpu.kf32)
kernel_code = gpu.create_kernel_code(
    """
    const GELU_SCALING_FACTOR: f32 = 0.7978845608028654; // sqrt(2.0 / PI)
    @group(0) @binding(0) var<storage, read_write> inp: array<{{precision}}>;
    @group(0) @binding(1) var<storage, read_write> out: array<{{precision}}>;
    @group(0) @binding(1) var<storage, read_write> dummy: array<{{precision}}>;
    @compute @workgroup_size({{workgroupSize}})
    fn main(
        @builtin(global_invocation_id) GlobalInvocationID: vec3<u32>) {
        let i: u32 = GlobalInvocationID.x;
        if (i < arrayLength(&inp)) {
            let x: f32 = inp[i];
            out[i] = select(0.5 * x * (1.0 + tanh(GELU_SCALING_FACTOR 
                     * (x + .044715 * x * x * x))), x, x > 10.0);
        }
    }
    """,
    256,
    gpu.kf32
    )

kernel = gpu.create_kernel(ctx, kernel_code, [input, output], [0,0], [12,1,1])

gpu.to_gpu_float(ctx, np.array([1,2,3,4,1,2,3,4,1,2,3,4],np.float32), input)

gpu_async = gpu.dispatch_kernel(ctx, kernel);

gpu.wait(ctx, gpu_async);

print(gpu.to_cpu_float(ctx, output))