tenflowers-core 0.1.1

Core tensor operations and execution engine for TenfloweRS
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144

/*
 * elementwise_kernels.hip
 * ROCm/HIP kernels for memory-coalesced element-wise operations
 * Optimized for AMD GPU memory hierarchy and wavefront execution
 */

#include <hip/hip_runtime.h>

// Memory-coalesced addition optimized for AMD wavefronts
__global__ void rocm_coalesced_add(float* a, float* b, float* output, int count) {
    int gid = blockIdx.x * blockDim.x + threadIdx.x;
    
    // Process multiple elements per thread for better memory utilization
    const int elements_per_thread = 4;
    int base_idx = gid * elements_per_thread;
    
    // Use vectorized loads when possible (float4 for better memory bandwidth)
    if (base_idx + 4 <= count && (base_idx % 4) == 0) {
        float4 vec_a = reinterpret_cast<float4*>(a)[base_idx / 4];
        float4 vec_b = reinterpret_cast<float4*>(b)[base_idx / 4];
        
        float4 result;
        result.x = vec_a.x + vec_b.x;
        result.y = vec_a.y + vec_b.y;
        result.z = vec_a.z + vec_b.z;
        result.w = vec_a.w + vec_b.w;
        
        reinterpret_cast<float4*>(output)[base_idx / 4] = result;
    } else {
        // Fallback to scalar operations for boundary elements
        #pragma unroll
        for (int i = 0; i < elements_per_thread; i++) {
            int idx = base_idx + i;
            if (idx < count) {
                output[idx] = a[idx] + b[idx];
            }
        }
    }
}

// Memory-coalesced multiplication
__global__ void rocm_coalesced_mul(float* a, float* b, float* output, int count) {
    int gid = blockIdx.x * blockDim.x + threadIdx.x;
    
    const int elements_per_thread = 4;
    int base_idx = gid * elements_per_thread;
    
    if (base_idx + 4 <= count && (base_idx % 4) == 0) {
        float4 vec_a = reinterpret_cast<float4*>(a)[base_idx / 4];
        float4 vec_b = reinterpret_cast<float4*>(b)[base_idx / 4];
        
        float4 result;
        result.x = vec_a.x * vec_b.x;
        result.y = vec_a.y * vec_b.y;
        result.z = vec_a.z * vec_b.z;
        result.w = vec_a.w * vec_b.w;
        
        reinterpret_cast<float4*>(output)[base_idx / 4] = result;
    } else {
        #pragma unroll
        for (int i = 0; i < elements_per_thread; i++) {
            int idx = base_idx + i;
            if (idx < count) {
                output[idx] = a[idx] * b[idx];
            }
        }
    }
}

// Memory-coalesced subtraction
__global__ void rocm_coalesced_sub(float* a, float* b, float* output, int count) {
    int gid = blockIdx.x * blockDim.x + threadIdx.x;
    
    const int elements_per_thread = 4;
    int base_idx = gid * elements_per_thread;
    
    if (base_idx + 4 <= count && (base_idx % 4) == 0) {
        float4 vec_a = reinterpret_cast<float4*>(a)[base_idx / 4];
        float4 vec_b = reinterpret_cast<float4*>(b)[base_idx / 4];
        
        float4 result;
        result.x = vec_a.x - vec_b.x;
        result.y = vec_a.y - vec_b.y;
        result.z = vec_a.z - vec_b.z;
        result.w = vec_a.w - vec_b.w;
        
        reinterpret_cast<float4*>(output)[base_idx / 4] = result;
    } else {
        #pragma unroll
        for (int i = 0; i < elements_per_thread; i++) {
            int idx = base_idx + i;
            if (idx < count) {
                output[idx] = a[idx] - b[idx];
            }
        }
    }
}

// Memory-coalesced division
__global__ void rocm_coalesced_div(float* a, float* b, float* output, int count) {
    int gid = blockIdx.x * blockDim.x + threadIdx.x;
    
    const int elements_per_thread = 4;
    int base_idx = gid * elements_per_thread;
    
    if (base_idx + 4 <= count && (base_idx % 4) == 0) {
        float4 vec_a = reinterpret_cast<float4*>(a)[base_idx / 4];
        float4 vec_b = reinterpret_cast<float4*>(b)[base_idx / 4];
        
        float4 result;
        result.x = vec_a.x / vec_b.x;
        result.y = vec_a.y / vec_b.y;
        result.z = vec_a.z / vec_b.z;
        result.w = vec_a.w / vec_b.w;
        
        reinterpret_cast<float4*>(output)[base_idx / 4] = result;
    } else {
        #pragma unroll
        for (int i = 0; i < elements_per_thread; i++) {
            int idx = base_idx + i;
            if (idx < count) {
                output[idx] = a[idx] / b[idx];
            }
        }
    }
}

// RDNA-optimized element-wise operations using wave intrinsics
__global__ void rocm_rdna_optimized_add(float* a, float* b, float* output, int count) {
    int gid = blockIdx.x * blockDim.x + threadIdx.x;
    int wavefront_id = gid / 64;  // RDNA wavefront size is 64
    int lane_id = gid % 64;
    
    // Use wave-level operations for better RDNA utilization
    if (gid < count) {
        float value_a = a[gid];
        float value_b = b[gid];
        float result = value_a + value_b;
        
        // Use wave shuffle for better memory access patterns
        // This is a simplified example - real implementation would use proper wave intrinsics
        output[gid] = result;
    }
}