GPU version for offset_particle.jl

src/low_level/offset_particle.jl

@@ -1,33 +1,42 @@
+using CUDA, BenchmarkTools, Adapt
 include("structures.jl")
 
-function offset_particle_set!(x_offset, y_offset, tilt, p_lab)
+Adapt.@adapt_structure GPU_Particle; Adapt.@adapt_structure Intermediate_Offset;
+
+function offset_particle_set!(x_offset, y_offset, tilt, p_lab, int)
     """Transform from lab to element coords.
     See Bmad manual (2022-11-06) sections 5.6.1, 15.3.1 and 24.2
     **NOTE**: transverse only as of now.
     """
-
-    (x, px, y, py, z, pz) = p_lab.r
-
-    s = sin(tilt)
-    c = cos(tilt)
-    x_ele_int = p_lab.r[1] - x_offset
-    y_ele_int = p_lab.r[3] - y_offset
-    x_ele = x_ele_int*c + y_ele_int*s
-    y_ele = -x_ele_int*s + y_ele_int*c
 
-    p_lab.r[1], p_lab.r[3] = x_ele, y_ele
+    x_ele_int, y_ele_int, x_ele, y_ele, px_ele, py_ele = int.x_ele_int, 
+    int.y_ele_int, int.x_ele, int.y_ele, int.px_ele, int.py_ele
+    x, px, y, py = p_lab.x, p_lab.px, p_lab.y, p_lab.py
+
+    index = (blockIdx().x - Float32(1)) * blockDim().x + threadIdx().x
+    stride = gridDim().x * blockDim().x
 
-    px_ele = p_lab.r[2]*c + p_lab.r[4]*s
-    py_ele = -p_lab.r[2]*s + p_lab.r[4]*c
+    while i <= length(x)
+
+        @inbounds (s[i] = sin(tilt[i]);
+        c[i] = cos(tilt[i]);
+        x_ele_int[i] = x[i] - x_offset[i];
+        y_ele_int[i] = y[i] - y_offset[i];
+        x_ele[i] = x_ele_int[i]*c[i]+ y_ele_int[i]*s[i];
+        y_ele[i] = -x_ele_int[i]*s[i] + y_ele_int[i]*c[i];
+        px_ele[i] = px[i]*c[i] + py[i]*s[i];
+        py_ele[i] = -px[i]*s[i] + py[i]*c[i];)
 
-    p_ele = p_lab
-    p_ele.r = (x_ele, px_ele, y_ele, py_ele, z, pz)
+        i += stride
+    end
 
-return p_ele
+    GPU_particles = GPU_Particle(x_ele, px_ele, y_ele, py_ele, z, pz, s, p0c, mc2)
 
+    return
 end
 
 
 """Test Value"""
-p_lab = Particle(1.0, 1.0, 1.0, [1.0, 0.5, 1.0, 0.2, 1.0, 0.5])
-offset_particle_set!(0.01, 0.01, 0.01, p_lab)
+x_offset = CUDA.fill(0.01, 50000); y_offset = CUDA.fill(0.01, 50000); tilt = CUDA.fill(0.01, 50000);
+
+p_lab = GPU_Particle(x, px, y, py, z, pz, s, p0c, mc2)

src/low_level/structures.jl

@@ -27,7 +27,7 @@ struct Particle{T}
 
 end
 
-struct Intermediate{K}
+struct Intermediate_Drift{K}
     """Intermediate helper struct"""
     P::K
     Px::K
@@ -42,7 +42,16 @@ struct GPU_Drift{K}
     L::K
 end
 
-struct Drift{T}
-    L::T
+struct Intermediate_Offset{K}
+    x_ele_int::K
+    y_ele_int::K
+    x_ele::K
+    y_ele::K
+    px_ele::K
+    py_ele::K
+    s::K
+    c::K
 end
 
+
+