Skip to main content

scirs2_core/gpu/kernels/
elementwise.rs

1//! Element-wise operation kernels for GPU
2//!
3//! These kernels implement basic element-wise operations that are fundamental
4//! to tensor computations: addition, multiplication, division, subtraction.
5
6use std::collections::HashMap;
7
8use crate::gpu::kernels::{
9    BaseKernel, DataType, GpuKernel, KernelMetadata, KernelParams, OperationType,
10};
11use crate::gpu::{GpuBackend, GpuError};
12
13/// Element-wise addition kernel (a + b)
14pub struct ElementwiseAddKernel {
15    base: BaseKernel,
16}
17
18impl Default for ElementwiseAddKernel {
19    fn default() -> Self {
20        Self::new()
21    }
22}
23
24impl ElementwiseAddKernel {
25    /// Create a new element-wise addition kernel
26    pub fn new() -> Self {
27        let metadata = KernelMetadata {
28            workgroup_size: [256, 1, 1],
29            local_memory_usage: 0,
30            supports_tensor_cores: false,
31            operationtype: OperationType::MemoryIntensive,
32            backend_metadata: HashMap::new(),
33        };
34
35        let cuda_source = r#"
36extern "C" __global__ void elementwise_add(
37    const float* __restrict__ a,
38    const float* __restrict__ b,
39    float* __restrict__ result,
40    int n
41) {
42    int i = blockIdx.x * blockDim.x + threadIdx.x;
43    if (i < n) {
44        result[i] = a[i] + b[i];
45    }
46}
47"#
48        .to_string();
49
50        let rocm_source = cuda_source.clone();
51
52        let wgpu_source = r#"
53struct Uniforms {
54    n: u32,
55};
56
57@group(0) @binding(0) var<uniform> uniforms: Uniforms;
58@group(0) @binding(1) var<storage, read> a: array<f32>;
59@group(0) @binding(2) var<storage, read> b: array<f32>;
60@group(0) @binding(3) var<storage, write> result: array<f32>;
61
62@compute @workgroup_size(256)
63fn elementwise_add(@builtin(global_invocation_id) global_id: vec3<u32>) {
64    let i = global_id.x;
65
66    if (i < uniforms.n) {
67        result[i] = a[i] + b[i];
68    }
69}
70"#
71        .to_string();
72
73        let metal_source = r#"
74#include <metal_stdlib>
75using namespace metal;
76
77kernel void elementwise_add(
78    const device float* a [[buffer(0)]],
79    const device float* b [[buffer(1)]],
80    device float* result [[buffer(2)]],
81    constant uint& n [[buffer(3)]],
82    uint gid [[thread_position_in_grid]])
83{
84    if (gid < n) {
85        result[gid] = a[gid] + b[gid];
86    }
87}
88"#
89        .to_string();
90
91        let opencl_source = r#"
92__kernel void elementwise_add(
93    __global const float* a,
94    __global const float* b,
95    __global float* result,
96    const int n)
97{
98    int i = get_global_id(0);
99    if (i < n) {
100        result[i] = a[i] + b[i];
101    }
102}
103"#
104        .to_string();
105
106        Self {
107            base: BaseKernel::new(
108                "elementwise_add",
109                &cuda_source,
110                &rocm_source,
111                &wgpu_source,
112                &metal_source,
113                &opencl_source,
114                metadata,
115            ),
116        }
117    }
118}
119
120impl GpuKernel for ElementwiseAddKernel {
121    fn name(&self) -> &str {
122        self.base.name()
123    }
124
125    fn source_for_backend(&self, backend: GpuBackend) -> Result<String, GpuError> {
126        self.base.source_for_backend(backend)
127    }
128
129    fn metadata(&self) -> KernelMetadata {
130        self.base.metadata()
131    }
132
133    fn can_specialize(&self, params: &KernelParams) -> bool {
134        matches!(
135            params.datatype,
136            DataType::Float32 | DataType::Float64 | DataType::Int32
137        )
138    }
139
140    fn specialize(&self, params: &KernelParams) -> Result<Box<dyn GpuKernel>, GpuError> {
141        if !self.can_specialize(params) {
142            return Err(GpuError::SpecializationNotSupported);
143        }
144
145        // For now, return the same kernel (no type specialization implemented yet)
146        Ok(Box::new(Self::new()))
147    }
148}
149
150/// Element-wise multiplication kernel (a * b)
151pub struct ElementwiseMulKernel {
152    base: BaseKernel,
153}
154
155impl Default for ElementwiseMulKernel {
156    fn default() -> Self {
157        Self::new()
158    }
159}
160
161impl ElementwiseMulKernel {
162    /// Create a new element-wise multiplication kernel
163    pub fn new() -> Self {
164        let metadata = KernelMetadata {
165            workgroup_size: [256, 1, 1],
166            local_memory_usage: 0,
167            supports_tensor_cores: false,
168            operationtype: OperationType::MemoryIntensive,
169            backend_metadata: HashMap::new(),
170        };
171
172        let cuda_source = r#"
173extern "C" __global__ void elementwise_mul(
174    const float* __restrict__ a,
175    const float* __restrict__ b,
176    float* __restrict__ result,
177    int n
178) {
179    int i = blockIdx.x * blockDim.x + threadIdx.x;
180    if (i < n) {
181        result[i] = a[i] * b[i];
182    }
183}
184"#
185        .to_string();
186
187        let rocm_source = cuda_source.clone();
188
189        let wgpu_source = r#"
190struct Uniforms {
191    n: u32,
192};
193
194@group(0) @binding(0) var<uniform> uniforms: Uniforms;
195@group(0) @binding(1) var<storage, read> a: array<f32>;
196@group(0) @binding(2) var<storage, read> b: array<f32>;
197@group(0) @binding(3) var<storage, write> result: array<f32>;
198
199@compute @workgroup_size(256)
200fn elementwise_mul(@builtin(global_invocation_id) global_id: vec3<u32>) {
201    let i = global_id.x;
202
203    if (i < uniforms.n) {
204        result[i] = a[i] * b[i];
205    }
206}
207"#
208        .to_string();
209
210        let metal_source = r#"
211#include <metal_stdlib>
212using namespace metal;
213
214kernel void elementwise_mul(
215    const device float* a [[buffer(0)]],
216    const device float* b [[buffer(1)]],
217    device float* result [[buffer(2)]],
218    constant uint& n [[buffer(3)]],
219    uint gid [[thread_position_in_grid]])
220{
221    if (gid < n) {
222        result[gid] = a[gid] * b[gid];
223    }
224}
225"#
226        .to_string();
227
228        let opencl_source = r#"
229__kernel void elementwise_mul(
230    __global const float* a,
231    __global const float* b,
232    __global float* result,
233    const int n)
234{
235    int i = get_global_id(0);
236    if (i < n) {
237        result[i] = a[i] * b[i];
238    }
239}
240"#
241        .to_string();
242
243        Self {
244            base: BaseKernel::new(
245                "elementwise_mul",
246                &cuda_source,
247                &rocm_source,
248                &wgpu_source,
249                &metal_source,
250                &opencl_source,
251                metadata,
252            ),
253        }
254    }
255}
256
257impl GpuKernel for ElementwiseMulKernel {
258    fn name(&self) -> &str {
259        self.base.name()
260    }
261
262    fn source_for_backend(&self, backend: GpuBackend) -> Result<String, GpuError> {
263        self.base.source_for_backend(backend)
264    }
265
266    fn metadata(&self) -> KernelMetadata {
267        self.base.metadata()
268    }
269
270    fn can_specialize(&self, params: &KernelParams) -> bool {
271        matches!(
272            params.datatype,
273            DataType::Float32 | DataType::Float64 | DataType::Int32
274        )
275    }
276
277    fn specialize(&self, params: &KernelParams) -> Result<Box<dyn GpuKernel>, GpuError> {
278        if !self.can_specialize(params) {
279            return Err(GpuError::SpecializationNotSupported);
280        }
281
282        Ok(Box::new(Self::new()))
283    }
284}
285
286/// Element-wise division kernel (a / b)
287pub struct ElementwiseDivKernel {
288    base: BaseKernel,
289}
290
291impl Default for ElementwiseDivKernel {
292    fn default() -> Self {
293        Self::new()
294    }
295}
296
297impl ElementwiseDivKernel {
298    /// Create a new element-wise division kernel
299    pub fn new() -> Self {
300        let metadata = KernelMetadata {
301            workgroup_size: [256, 1, 1],
302            local_memory_usage: 0,
303            supports_tensor_cores: false,
304            operationtype: OperationType::ComputeIntensive, // Division is more compute-intensive
305            backend_metadata: HashMap::new(),
306        };
307
308        let cuda_source = r#"
309extern "C" __global__ void elementwise_div(
310    const float* __restrict__ a,
311    const float* __restrict__ b,
312    float* __restrict__ result,
313    int n
314) {
315    int i = blockIdx.x * blockDim.x + threadIdx.x;
316    if (i < n) {
317        result[i] = a[i] / b[i];
318    }
319}
320"#
321        .to_string();
322
323        let rocm_source = cuda_source.clone();
324
325        let wgpu_source = r#"
326struct Uniforms {
327    n: u32,
328};
329
330@group(0) @binding(0) var<uniform> uniforms: Uniforms;
331@group(0) @binding(1) var<storage, read> a: array<f32>;
332@group(0) @binding(2) var<storage, read> b: array<f32>;
333@group(0) @binding(3) var<storage, write> result: array<f32>;
334
335@compute @workgroup_size(256)
336fn elementwise_div(@builtin(global_invocation_id) global_id: vec3<u32>) {
337    let i = global_id.x;
338
339    if (i < uniforms.n) {
340        result[i] = a[i] / b[i];
341    }
342}
343"#
344        .to_string();
345
346        let metal_source = r#"
347#include <metal_stdlib>
348using namespace metal;
349
350kernel void elementwise_div(
351    const device float* a [[buffer(0)]],
352    const device float* b [[buffer(1)]],
353    device float* result [[buffer(2)]],
354    constant uint& n [[buffer(3)]],
355    uint gid [[thread_position_in_grid]])
356{
357    if (gid < n) {
358        result[gid] = a[gid] / b[gid];
359    }
360}
361"#
362        .to_string();
363
364        let opencl_source = r#"
365__kernel void elementwise_div(
366    __global const float* a,
367    __global const float* b,
368    __global float* result,
369    const int n)
370{
371    int i = get_global_id(0);
372    if (i < n) {
373        result[i] = a[i] / b[i];
374    }
375}
376"#
377        .to_string();
378
379        Self {
380            base: BaseKernel::new(
381                "elementwise_div",
382                &cuda_source,
383                &rocm_source,
384                &wgpu_source,
385                &metal_source,
386                &opencl_source,
387                metadata,
388            ),
389        }
390    }
391}
392
393impl GpuKernel for ElementwiseDivKernel {
394    fn name(&self) -> &str {
395        self.base.name()
396    }
397
398    fn source_for_backend(&self, backend: GpuBackend) -> Result<String, GpuError> {
399        self.base.source_for_backend(backend)
400    }
401
402    fn metadata(&self) -> KernelMetadata {
403        self.base.metadata()
404    }
405
406    fn can_specialize(&self, params: &KernelParams) -> bool {
407        matches!(params.datatype, DataType::Float32 | DataType::Float64)
408    }
409
410    fn specialize(&self, params: &KernelParams) -> Result<Box<dyn GpuKernel>, GpuError> {
411        if !self.can_specialize(params) {
412            return Err(GpuError::SpecializationNotSupported);
413        }
414
415        Ok(Box::new(Self::new()))
416    }
417}
418
419/// Scalar multiplication kernel (a * scalar)
420pub struct ScalarMulKernel {
421    base: BaseKernel,
422}
423
424impl Default for ScalarMulKernel {
425    fn default() -> Self {
426        Self::new()
427    }
428}
429
430impl ScalarMulKernel {
431    /// Create a new scalar multiplication kernel
432    pub fn new() -> Self {
433        let metadata = KernelMetadata {
434            workgroup_size: [256, 1, 1],
435            local_memory_usage: 0,
436            supports_tensor_cores: false,
437            operationtype: OperationType::MemoryIntensive,
438            backend_metadata: HashMap::new(),
439        };
440
441        let cuda_source = r#"
442extern "C" __global__ void scalar_mul(
443    const float* __restrict__ input,
444    float* __restrict__ output,
445    float scalar,
446    int n
447) {
448    int i = blockIdx.x * blockDim.x + threadIdx.x;
449    if (i < n) {
450        output[i] = input[i] * scalar;
451    }
452}
453"#
454        .to_string();
455
456        let rocm_source = cuda_source.clone();
457
458        let wgpu_source = r#"
459struct Uniforms {
460    n: u32,
461    scalar: f32,
462};
463
464@group(0) @binding(0) var<uniform> uniforms: Uniforms;
465@group(0) @binding(1) var<storage, read> input: array<f32>;
466@group(0) @binding(2) var<storage, write> output: array<f32>;
467
468@compute @workgroup_size(256)
469fn scalar_mul(@builtin(global_invocation_id) global_id: vec3<u32>) {
470    let i = global_id.x;
471
472    if (i < uniforms.n) {
473        output[i] = input[i] * uniforms.scalar;
474    }
475}
476"#
477        .to_string();
478
479        let metal_source = r#"
480#include <metal_stdlib>
481using namespace metal;
482
483kernel void scalar_mul(
484    const device float* input [[buffer(0)]],
485    device float* output [[buffer(1)]],
486    constant float& scalar [[buffer(2)]],
487    constant uint& n [[buffer(3)]],
488    uint gid [[thread_position_in_grid]])
489{
490    if (gid < n) {
491        output[gid] = input[gid] * scalar;
492    }
493}
494"#
495        .to_string();
496
497        let opencl_source = r#"
498__kernel void scalar_mul(
499    __global const float* input,
500    __global float* output,
501    const float scalar,
502    const int n)
503{
504    int i = get_global_id(0);
505    if (i < n) {
506        output[i] = input[i] * scalar;
507    }
508}
509"#
510        .to_string();
511
512        Self {
513            base: BaseKernel::new(
514                "scalar_mul",
515                &cuda_source,
516                &rocm_source,
517                &wgpu_source,
518                &metal_source,
519                &opencl_source,
520                metadata,
521            ),
522        }
523    }
524}
525
526impl GpuKernel for ScalarMulKernel {
527    fn name(&self) -> &str {
528        self.base.name()
529    }
530
531    fn source_for_backend(&self, backend: GpuBackend) -> Result<String, GpuError> {
532        self.base.source_for_backend(backend)
533    }
534
535    fn metadata(&self) -> KernelMetadata {
536        self.base.metadata()
537    }
538
539    fn can_specialize(&self, params: &KernelParams) -> bool {
540        matches!(
541            params.datatype,
542            DataType::Float32 | DataType::Float64 | DataType::Int32
543        )
544    }
545
546    fn specialize(&self, params: &KernelParams) -> Result<Box<dyn GpuKernel>, GpuError> {
547        if !self.can_specialize(params) {
548            return Err(GpuError::SpecializationNotSupported);
549        }
550
551        Ok(Box::new(Self::new()))
552    }
553}
554/// Element-wise subtraction kernel (a - b)
555pub struct ElementwiseSubKernel {
556    base: BaseKernel,
557}
558
559impl Default for ElementwiseSubKernel {
560    fn default() -> Self {
561        Self::new()
562    }
563}
564
565impl ElementwiseSubKernel {
566    pub fn new() -> Self {
567        let metadata = KernelMetadata {
568            workgroup_size: [256, 1, 1],
569            local_memory_usage: 0,
570            supports_tensor_cores: false,
571            operationtype: OperationType::MemoryIntensive,
572            backend_metadata: HashMap::new(),
573        };
574
575        let cuda_source = r#"
576extern \"C\" __global__ void elementwise_sub(
577    const float* __restrict__ a,
578    const float* __restrict__ b,
579    float* __restrict__ result,
580    int n
581) {
582    int i = blockIdx.x * blockDim.x + threadIdx.x;
583    if (i < n) {
584        result[i] = a[i] - b[i];
585    }
586}
587"#
588        .to_string();
589
590        let rocm_source = cuda_source.clone();
591
592        let wgpu_source = r#"
593struct Uniforms {
594    n: u32,
595};
596
597@group(0) @binding(0) var<uniform> uniforms: Uniforms;
598@group(0) @binding(1) var<storage, read> a: array<f32>;
599@group(0) @binding(2) var<storage, read> b: array<f32>;
600@group(0) @binding(3) var<storage, read_write> result: array<f32>;
601
602@compute @workgroup_size(256)
603fn main(@builtin(global_invocation_id) global_id: vec3<u32>) {
604    let i = global_id.x;
605
606    if (i < uniforms.n) {
607        result[i] = a[i] - b[i];
608    }
609}
610"#
611        .to_string();
612
613        let metal_source = String::new();
614        let opencl_source = String::new();
615
616        Self {
617            base: BaseKernel::new(
618                "elementwise_sub",
619                &cuda_source,
620                &rocm_source,
621                &wgpu_source,
622                &metal_source,
623                &opencl_source,
624                metadata,
625            ),
626        }
627    }
628}
629
630impl GpuKernel for ElementwiseSubKernel {
631    fn name(&self) -> &str {
632        self.base.name()
633    }
634
635    fn source_for_backend(&self, backend: GpuBackend) -> Result<String, GpuError> {
636        self.base.source_for_backend(backend)
637    }
638
639    fn metadata(&self) -> KernelMetadata {
640        self.base.metadata()
641    }
642
643    fn can_specialize(&self, params: &KernelParams) -> bool {
644        matches!(
645            params.datatype,
646            DataType::Float32 | DataType::Float64 | DataType::Int32
647        )
648    }
649
650    fn specialize(&self, params: &KernelParams) -> Result<Box<dyn GpuKernel>, GpuError> {
651        if !self.can_specialize(params) {
652            return Err(GpuError::SpecializationNotSupported);
653        }
654        Ok(Box::new(Self::new()))
655    }
656}
657/// Element-wise power kernel (pow(a, b))
658pub struct ElementwisePowKernel {
659    base: BaseKernel,
660}
661
662impl Default for ElementwisePowKernel {
663    fn default() -> Self {
664        Self::new()
665    }
666}
667
668impl ElementwisePowKernel {
669    pub fn new() -> Self {
670        let metadata = KernelMetadata {
671            workgroup_size: [256, 1, 1],
672            local_memory_usage: 0,
673            supports_tensor_cores: false,
674            operationtype: OperationType::ComputeIntensive,
675            backend_metadata: HashMap::new(),
676        };
677
678        let cuda_source = r#"
679extern "C" __global__ void elementwise_pow(
680    const float* __restrict__ a,
681    const float* __restrict__ b,
682    float* __restrict__ result,
683    int n
684) {
685    int i = blockIdx.x * blockDim.x + threadIdx.x;
686    if (i < n) {
687        result[i] = powf(a[i], b[i]);
688    }
689}
690"#
691        .to_string();
692
693        let rocm_source = cuda_source.clone();
694
695        let wgpu_source = r#"
696struct Uniforms {
697    n: u32,
698};
699
700@group(0) @binding(0) var<uniform> uniforms: Uniforms;
701@group(0) @binding(1) var<storage, read> a: array<f32>;
702@group(0) @binding(2) var<storage, read> b: array<f32>;
703@group(0) @binding(3) var<storage, read_write> result: array<f32>;
704
705@compute @workgroup_size(256)
706fn main(@builtin(global_invocation_id) global_id: vec3<u32>) {
707    let i = global_id.x;
708
709    if (i < uniforms.n) {
710        result[i] = pow(a[i], b[i]);
711    }
712}
713"#
714        .to_string();
715
716        let metal_source = String::new();
717        let opencl_source = String::new();
718
719        Self {
720            base: BaseKernel::new(
721                "elementwise_pow",
722                &cuda_source,
723                &rocm_source,
724                &wgpu_source,
725                &metal_source,
726                &opencl_source,
727                metadata,
728            ),
729        }
730    }
731}
732
733impl GpuKernel for ElementwisePowKernel {
734    fn name(&self) -> &str {
735        self.base.name()
736    }
737
738    fn source_for_backend(&self, backend: GpuBackend) -> Result<String, GpuError> {
739        self.base.source_for_backend(backend)
740    }
741
742    fn metadata(&self) -> KernelMetadata {
743        self.base.metadata()
744    }
745
746    fn can_specialize(&self, params: &KernelParams) -> bool {
747        matches!(params.datatype, DataType::Float32 | DataType::Float64)
748    }
749
750    fn specialize(&self, params: &KernelParams) -> Result<Box<dyn GpuKernel>, GpuError> {
751        if !self.can_specialize(params) {
752            return Err(GpuError::SpecializationNotSupported);
753        }
754        Ok(Box::new(Self::new()))
755    }
756}
757
758/// Element-wise square root kernel (sqrt(a))
759pub struct ElementwiseSqrtKernel {
760    base: BaseKernel,
761}
762
763impl Default for ElementwiseSqrtKernel {
764    fn default() -> Self {
765        Self::new()
766    }
767}
768
769impl ElementwiseSqrtKernel {
770    pub fn new() -> Self {
771        let metadata = KernelMetadata {
772            workgroup_size: [256, 1, 1],
773            local_memory_usage: 0,
774            supports_tensor_cores: false,
775            operationtype: OperationType::ComputeIntensive,
776            backend_metadata: HashMap::new(),
777        };
778
779        let cuda_source = r#"
780extern "C" __global__ void elementwise_sqrt(
781    const float* __restrict__ a,
782    float* __restrict__ result,
783    int n
784) {
785    int i = blockIdx.x * blockDim.x + threadIdx.x;
786    if (i < n) {
787        result[i] = sqrtf(a[i]);
788    }
789}
790"#
791        .to_string();
792
793        let rocm_source = cuda_source.clone();
794
795        let wgpu_source = r#"
796struct Uniforms {
797    n: u32,
798};
799
800@group(0) @binding(0) var<uniform> uniforms: Uniforms;
801@group(0) @binding(1) var<storage, read> input: array<f32>;
802@group(0) @binding(2) var<storage, read_write> result: array<f32>;
803
804@compute @workgroup_size(256)
805fn main(@builtin(global_invocation_id) global_id: vec3<u32>) {
806    let i = global_id.x;
807
808    if (i < uniforms.n) {
809        result[i] = sqrt(input[i]);
810    }
811}
812"#
813        .to_string();
814
815        let metal_source = String::new();
816        let opencl_source = String::new();
817
818        Self {
819            base: BaseKernel::new(
820                "elementwise_sqrt",
821                &cuda_source,
822                &rocm_source,
823                &wgpu_source,
824                &metal_source,
825                &opencl_source,
826                metadata,
827            ),
828        }
829    }
830}
831
832impl GpuKernel for ElementwiseSqrtKernel {
833    fn name(&self) -> &str {
834        self.base.name()
835    }
836
837    fn source_for_backend(&self, backend: GpuBackend) -> Result<String, GpuError> {
838        self.base.source_for_backend(backend)
839    }
840
841    fn metadata(&self) -> KernelMetadata {
842        self.base.metadata()
843    }
844
845    fn can_specialize(&self, params: &KernelParams) -> bool {
846        matches!(params.datatype, DataType::Float32 | DataType::Float64)
847    }
848
849    fn specialize(&self, params: &KernelParams) -> Result<Box<dyn GpuKernel>, GpuError> {
850        if !self.can_specialize(params) {
851            return Err(GpuError::SpecializationNotSupported);
852        }
853        Ok(Box::new(Self::new()))
854    }
855}
856
857/// Element-wise exponential kernel (exp(a))
858pub struct ElementwiseExpKernel {
859    base: BaseKernel,
860}
861
862impl Default for ElementwiseExpKernel {
863    fn default() -> Self {
864        Self::new()
865    }
866}
867
868impl ElementwiseExpKernel {
869    pub fn new() -> Self {
870        let metadata = KernelMetadata {
871            workgroup_size: [256, 1, 1],
872            local_memory_usage: 0,
873            supports_tensor_cores: false,
874            operationtype: OperationType::ComputeIntensive,
875            backend_metadata: HashMap::new(),
876        };
877
878        let cuda_source = r#"
879extern "C" __global__ void elementwise_exp(
880    const float* __restrict__ a,
881    float* __restrict__ result,
882    int n
883) {
884    int i = blockIdx.x * blockDim.x + threadIdx.x;
885    if (i < n) {
886        result[i] = expf(a[i]);
887    }
888}
889"#
890        .to_string();
891
892        let rocm_source = cuda_source.clone();
893
894        let wgpu_source = r#"
895struct Uniforms {
896    n: u32,
897};
898
899@group(0) @binding(0) var<uniform> uniforms: Uniforms;
900@group(0) @binding(1) var<storage, read> input: array<f32>;
901@group(0) @binding(2) var<storage, read_write> result: array<f32>;
902
903@compute @workgroup_size(256)
904fn main(@builtin(global_invocation_id) global_id: vec3<u32>) {
905    let i = global_id.x;
906
907    if (i < uniforms.n) {
908        result[i] = exp(input[i]);
909    }
910}
911"#
912        .to_string();
913
914        let metal_source = String::new();
915        let opencl_source = String::new();
916
917        Self {
918            base: BaseKernel::new(
919                "elementwise_exp",
920                &cuda_source,
921                &rocm_source,
922                &wgpu_source,
923                &metal_source,
924                &opencl_source,
925                metadata,
926            ),
927        }
928    }
929}
930
931impl GpuKernel for ElementwiseExpKernel {
932    fn name(&self) -> &str {
933        self.base.name()
934    }
935
936    fn source_for_backend(&self, backend: GpuBackend) -> Result<String, GpuError> {
937        self.base.source_for_backend(backend)
938    }
939
940    fn metadata(&self) -> KernelMetadata {
941        self.base.metadata()
942    }
943
944    fn can_specialize(&self, params: &KernelParams) -> bool {
945        matches!(params.datatype, DataType::Float32 | DataType::Float64)
946    }
947
948    fn specialize(&self, params: &KernelParams) -> Result<Box<dyn GpuKernel>, GpuError> {
949        if !self.can_specialize(params) {
950            return Err(GpuError::SpecializationNotSupported);
951        }
952        Ok(Box::new(Self::new()))
953    }
954}
955
956/// Element-wise logarithm kernel (log(a))
957pub struct ElementwiseLogKernel {
958    base: BaseKernel,
959}
960
961impl Default for ElementwiseLogKernel {
962    fn default() -> Self {
963        Self::new()
964    }
965}
966
967impl ElementwiseLogKernel {
968    pub fn new() -> Self {
969        let metadata = KernelMetadata {
970            workgroup_size: [256, 1, 1],
971            local_memory_usage: 0,
972            supports_tensor_cores: false,
973            operationtype: OperationType::ComputeIntensive,
974            backend_metadata: HashMap::new(),
975        };
976
977        let cuda_source = r#"
978extern "C" __global__ void elementwise_log(
979    const float* __restrict__ a,
980    float* __restrict__ result,
981    int n
982) {
983    int i = blockIdx.x * blockDim.x + threadIdx.x;
984    if (i < n) {
985        result[i] = logf(a[i]);
986    }
987}
988"#
989        .to_string();
990
991        let rocm_source = cuda_source.clone();
992
993        let wgpu_source = r#"
994struct Uniforms {
995    n: u32,
996};
997
998@group(0) @binding(0) var<uniform> uniforms: Uniforms;
999@group(0) @binding(1) var<storage, read> input: array<f32>;
1000@group(0) @binding(2) var<storage, read_write> result: array<f32>;
1001
1002@compute @workgroup_size(256)
1003fn main(@builtin(global_invocation_id) global_id: vec3<u32>) {
1004    let i = global_id.x;
1005
1006    if (i < uniforms.n) {
1007        result[i] = log(input[i]);
1008    }
1009}
1010"#
1011        .to_string();
1012
1013        let metal_source = String::new();
1014        let opencl_source = String::new();
1015
1016        Self {
1017            base: BaseKernel::new(
1018                "elementwise_log",
1019                &cuda_source,
1020                &rocm_source,
1021                &wgpu_source,
1022                &metal_source,
1023                &opencl_source,
1024                metadata,
1025            ),
1026        }
1027    }
1028}
1029
1030impl GpuKernel for ElementwiseLogKernel {
1031    fn name(&self) -> &str {
1032        self.base.name()
1033    }
1034
1035    fn source_for_backend(&self, backend: GpuBackend) -> Result<String, GpuError> {
1036        self.base.source_for_backend(backend)
1037    }
1038
1039    fn metadata(&self) -> KernelMetadata {
1040        self.base.metadata()
1041    }
1042
1043    fn can_specialize(&self, params: &KernelParams) -> bool {
1044        matches!(params.datatype, DataType::Float32 | DataType::Float64)
1045    }
1046
1047    fn specialize(&self, params: &KernelParams) -> Result<Box<dyn GpuKernel>, GpuError> {
1048        if !self.can_specialize(params) {
1049            return Err(GpuError::SpecializationNotSupported);
1050        }
1051        Ok(Box::new(Self::new()))
1052    }
1053}