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
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175

//! This is a wrapper crate around all other context switch implementation crates
//! that helps to manage the complex configuration options involving SIMD and personalities.

#![no_std]
#![feature(naked_functions)]

pub use context_switch_regular::read_first_register;

// If `simd_personality` is enabled, all of the `context_switch*` implementation crates are simultaneously enabled,
// in order to allow choosing one of them based on the configuration options of each Task (SIMD, regular, etc).
// If `simd_personality` is NOT enabled, then we use the context_switch routine that matches the actual build target. 
cfg_if::cfg_if! {
    if #[cfg(simd_personality)] {
        use core::arch::asm;
        pub use context_switch_sse::*;
        pub use context_switch_regular::*;
        pub use context_switch_avx::*;

        /// Switches context from a regular Task to an SSE Task.
        /// 
        /// # Arguments
        /// * First argument  (in `RDI`): mutable pointer to the previous task's stack pointer
        /// * Second argument (in `RSI`): the value of the next task's stack pointer
        /// 
        /// # Safety
        /// This function is unsafe because it changes the content on both task's stacks. 
        #[naked]
        pub unsafe extern "C" fn context_switch_regular_to_sse(_prev_stack_pointer: *mut usize, _next_stack_pointer_value: usize) {
            // Since this is a naked function that expects its arguments in two registers,
            // you CANNOT place any log statements or other instructions here
            // before, in between, or after anything below.
            asm!(
                save_registers_regular!(),
                switch_stacks!(),
                restore_registers_sse!(),
                restore_registers_regular!(),
                options(noreturn)
            );
        }


        /// Switches context from an SSE Task to a regular Task.
        ///
        /// # Arguments
        /// * First argument  (in `RDI`): mutable pointer to the previous task's stack pointer
        /// * Second argument (in `RSI`): the value of the next task's stack pointer
        ///
        /// # Safety
        /// This function is unsafe because it changes the content on both task's stacks.
        #[naked]
        pub unsafe extern "C" fn context_switch_sse_to_regular(_prev_stack_pointer: *mut usize, _next_stack_pointer_value: usize) {
            // Since this is a naked function that expects its arguments in two registers,
            // you CANNOT place any log statements or other instructions here
            // before, in between, or after anything below.
            asm!(
                save_registers_regular!(),
                save_registers_sse!(),
                switch_stacks!(),
                restore_registers_regular!(),
                options(noreturn)
            );
        }

        /// Switches context from a regular Task to an AVX Task.
        /// 
        /// # Arguments
        /// * First argument  (in `RDI`): mutable pointer to the previous task's stack pointer
        /// * Second argument (in `RSI`): the value of the next task's stack pointer
        /// 
        /// # Safety
        /// This function is unsafe because it changes the content on both task's stacks. 
        #[naked]
        pub unsafe extern "C" fn context_switch_regular_to_avx(_prev_stack_pointer: *mut usize, _next_stack_pointer_value: usize) {
            // Since this is a naked function that expects its arguments in two registers,
            // you CANNOT place any log statements or other instructions here
            // before, in between, or after anything below.
            asm!(
                save_registers_regular!(),
                switch_stacks!(),
                restore_registers_avx!(),
                restore_registers_regular!(),
                options(noreturn)
            );
        }

        /// Switches context from an SSE Task to an AVX regular Task.
        /// 
        /// # Arguments
        /// * First argument  (in `RDI`): mutable pointer to the previous task's stack pointer
        /// * Second argument (in `RSI`): the value of the next task's stack pointer
        /// 
        /// # Safety
        /// This function is unsafe because it changes the content on both task's stacks. 
        #[naked]
        pub unsafe extern "C" fn context_switch_sse_to_avx(_prev_stack_pointer: *mut usize, _next_stack_pointer_value: usize) {
            // Since this is a naked function that expects its arguments in two registers,
            // you CANNOT place any log statements or other instructions here
            // before, in between, or after anything below.
            asm!(
                save_registers_regular!(),
                save_registers_sse!(),
                switch_stacks!(),
                restore_registers_avx!(),
                restore_registers_regular!(),
                options(noreturn)
            );
        }

        /// Switches context from an AVX Task to a regular Task.
        /// 
        /// # Arguments
        /// * First argument  (in `RDI`): mutable pointer to the previous task's stack pointer
        /// * Second argument (in `RSI`): the value of the next task's stack pointer
        /// 
        /// # Safety
        /// This function is unsafe because it changes the content on both task's stacks. 
        #[naked]
        pub unsafe extern "C" fn context_switch_avx_to_regular(_prev_stack_pointer: *mut usize, _next_stack_pointer_value: usize) {
            // Since this is a naked function that expects its arguments in two registers,
            // you CANNOT place any log statements or other instructions here
            // before, in between, or after anything below.
            asm!(
                save_registers_regular!(),
                save_registers_avx!(),
                switch_stacks!(),
                restore_registers_regular!(),
                options(noreturn)
            );
        }

        /// Switches context from an AVX Task to an SSE Task.
        /// 
        /// # Arguments
        /// * First argument  (in `RDI`): mutable pointer to the previous task's stack pointer
        /// * Second argument (in `RSI`): the value of the next task's stack pointer
        /// 
        /// # Safety
        /// This function is unsafe because it changes the content on both task's stacks. 
        #[naked]
        pub unsafe extern "C" fn context_switch_avx_to_sse(_prev_stack_pointer: *mut usize, _next_stack_pointer_value: usize) {
            // Since this is a naked function that expects its arguments in two registers,
            // you CANNOT place any log statements or other instructions here
            // before, in between, or after anything below.
            asm!(
                save_registers_regular!(),
                save_registers_avx!(),
                switch_stacks!(),
                restore_registers_sse!(),
                restore_registers_regular!(),
                options(noreturn)
            );
        }
    }

    // BELOW HERE: simd_personality is disabled, and we just need to re-export a single Context & context_switch.
    // Because "else-if" statements are executed top-down, we need to put newer standards (supersets) before older standards (subsets).
    // For example, AVX512 is above AVX2, which is above AVX, which is above SSE, which is above regular (non-SIMD).

    else if #[cfg(target_feature = "avx")] {
        pub use context_switch_avx::ContextAVX as Context;
        pub use context_switch_avx::context_switch_avx as context_switch;
    }

    else if #[cfg(target_feature = "sse2")] {
        // this crate covers SSE, SSE2, SSE3, SSE4, but we're only currently using it for SSE2
        pub use context_switch_sse::ContextSSE as Context;
        pub use context_switch_sse::context_switch_sse as context_switch;
    }

    else {
        // this covers only the default x86_64 registers
        pub use context_switch_regular::ContextRegular as Context;
        pub use context_switch_regular::context_switch_regular as context_switch;
    }
}