Context Switching: The Heart of Multitasking

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

; x86-64 Interrupt Stack Frame (automatically pushed by CPU)
; 
; When an interrupt occurs in user mode, the CPU automatically:
; 1. Switches from user stack to kernel stack (using TSS.RSP0)
; 2. Pushes the following values (growing downward):
;
; Higher addresses
; ┌──────────────────────────────────────┐
; │          SS (user stack segment)     │ +40  (8 bytes, padded)
; ├──────────────────────────────────────┤
; │          RSP (user stack pointer)    │ +32  (8 bytes)
; ├──────────────────────────────────────┤
; │          RFLAGS (flags register)     │ +24  (8 bytes)
; ├──────────────────────────────────────┤
; │          CS (user code segment)      │ +16  (8 bytes, padded)
; ├──────────────────────────────────────┤
; │          RIP (instruction pointer)   │ +8   (8 bytes)
; ├──────────────────────────────────────┤
; │   Error Code (some exceptions only)  │ +0   (8 bytes, optional)
; └──────────────────────────────────────┘
; │          (kernel RSP points here)    │
; Lower addresses
;
; After this automatic push, CPU jumps to interrupt handler.
; Handler must MANUALLY save all general-purpose registers.
 
; The hardware IRET instruction pops this frame to return:
;    iretq   ; Pop RIP, CS, RFLAGS, RSP, SS and return to user space

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

/**
 * struct pt_regs - saved CPU register state
 * 
 * This structure is pushed onto the kernel stack at the start
 * of any interrupt or system call entry. It captures the complete
 * CPU register state at the moment of kernel entry.
 * 
 * Located in: arch/x86/include/asm/ptrace.h
 */
struct pt_regs {
    /* Manually saved by interrupt entry code */
    unsigned long r15;
    unsigned long r14;
    unsigned long r13;
    unsigned long r12;
    unsigned long bp;      /* RBP - base pointer */
    unsigned long bx;      /* RBX */
    unsigned long r11;
    unsigned long r10;
    unsigned long r9;
    unsigned long r8;
    unsigned long ax;      /* RAX - return value / syscall number */
    unsigned long cx;      /* RCX - 4th syscall argument */
    unsigned long dx;      /* RDX - 3rd syscall argument */
    unsigned long si;      /* RSI - 2nd syscall argument */
    unsigned long di;      /* RDI - 1st syscall argument */
    
    /* Identifies the interrupt/exception source */
    unsigned long orig_ax; /* Original RAX (syscall number) */
    
    /* Automatically saved by CPU on interrupt */
    unsigned long ip;      /* RIP - instruction pointer */
    unsigned long cs;      /* Code segment */
    unsigned long flags;   /* RFLAGS */
    unsigned long sp;      /* RSP - user stack pointer */
    unsigned long ss;      /* Stack segment */
};
 
/* Size: 21 * 8 = 168 bytes of register state */

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

/**
 * Simplified interrupt entry macro (Linux style)
 * 
 * This assembly code runs at the very start of interrupt handling,
 * immediately after the CPU pushes SS, RSP, RFLAGS, CS, RIP.
 * 
 * It completes the pt_regs structure by pushing all GPRs.
 */
 
.macro PUSH_REGS
    /* At this point, CPU has already pushed:
     * SS, RSP, RFLAGS, CS, RIP (and maybe error code)
     */
    
    /* Push all general-purpose registers to complete pt_regs */
    pushq   %rdi            /* Save 1st argument */
    pushq   %rsi            /* Save 2nd argument */
    pushq   %rdx            /* Save 3rd argument */
    pushq   %rcx            /* Save 4th argument */
    pushq   %rax            /* Save syscall number / return value */
    pushq   %r8             /* Save 5th argument */
    pushq   %r9             /* Save 6th argument */
    pushq   %r10
    pushq   %r11            /* Scratch registers */
    pushq   %rbx
    pushq   %rbp            /* Frame pointer (callee-saved) */
    pushq   %r12
    pushq   %r13
    pushq   %r14
    pushq   %r15
    
    /* Now RSP points to a complete pt_regs structure */
    /* We can pass this pointer to C interrupt handlers */
.endm
 
/* Example interrupt entry point */
ENTRY(interrupt_entry)
    /* CPU already pushed SS, RSP, RFLAGS, CS, RIP */
    PUSH_REGS               /* Push remaining registers */
    
    movq    %rsp, %rdi      /* pt_regs pointer as 1st arg to C */
    call    do_IRQ          /* Call C interrupt handler */
    
    /* After handler returns, restore all registers */
    POP_REGS               /* Reverse of PUSH_REGS */
    iretq                  /* Return from interrupt, pop CPU frame */
ENDPROC(interrupt_entry)
 
.macro POP_REGS
    popq    %r15
    popq    %r14
    popq    %r13
    popq    %r12
    popq    %rbp
    popq    %rbx
    popq    %r11
    popq    %r10
    popq    %r9
    popq    %r8
    popq    %rax
    popq    %rcx
    popq    %rdx
    popq    %rsi
    popq    %rdi
.endm

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

/**
 * struct thread_struct - Per-thread CPU state (x86-64)
 * 
 * This structure is embedded in task_struct and holds the
 * saved CPU context when the thread is not running.
 * 
 * Key fields shown (simplified from actual Linux):
 */
struct thread_struct {
    /* Segment descriptor caches - for segment reload on switch */
    unsigned short          es;
    unsigned short          ds;
    unsigned short          fsindex;
    unsigned short          gsindex;
    
    /* Thread-local storage base addresses */
    unsigned long           fsbase;    /* FS segment base (TLS) */
    unsigned long           gsbase;    /* GS segment base (per-CPU data) */
    
    /* Saved stack pointer - THE critical value for switch */
    unsigned long           sp;        /* Kernel stack pointer (RSP) */
    
    /* Debug registers */
    unsigned long           debugreg[8];
    
    /* Floating-point state (not directly here in modern kernels) */
    struct fpu              fpu;       /* FPU/SSE/AVX state container */
    
    /* I/O permission bitmap */
    unsigned long           *io_bitmap_ptr;
    unsigned long           io_bitmap_max;
    
    /*
     * Note: General registers (RAX, RBX, etc.) and RIP are NOT here!
     * They're stored in the pt_regs on the kernel stack.
     * The 'sp' field points to where pt_regs is located.
     */
};
 
/**
 * During context switch (simplified):
 * 
 * 1. Save current RSP into current->thread.sp
 * 2. Load next->thread.sp into RSP
 * 3. (RSP now points to next's kernel stack with its pt_regs)
 * 4. Return "into" the next process
 */

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

/**
 * Lazy FPU context switching (traditional approach)
 * 
 * Modern Linux (since ~4.x) uses "eager" FPU switching for
 * security reasons, but lazy switching illustrates the concept.
 */
 
/* The CR0.TS (Task Switched) bit controls FPU traps */
#define CR0_TS_BIT  (1 << 3)
 
/**
 * Called during context switch - DON'T save FPU, just set trap
 */
void switch_fpu_lazy(struct task_struct *prev, struct task_struct *next)
{
    /* 
     * Set CR0.TS bit - next FPU instruction will trap.
     * We haven't spent time saving prev's FPU state yet.
     */
    unsigned long cr0 = read_cr0();
    write_cr0(cr0 | CR0_TS_BIT);
    
    /* Remember who "owns" the FPU right now */
    this_cpu_write(fpu_owner, prev);
}
 
/**
 * Device Not Available exception handler (#NM, vector 7)
 * 
 * Called when process tries to use FPU but TS bit is set.
 */
void do_device_not_available(struct pt_regs *regs)
{
    struct task_struct *current = get_current();
    struct task_struct *prev_owner = this_cpu_read(fpu_owner);
    
    /*
     * NOW we actually save the FPU state - but only if someone
     * else had been using it (not just saving empty state)
     */
    if (prev_owner && prev_owner != current) {
        /* Save prev_owner's FPU state to their fpu struct */
        fxsave(&prev_owner->thread.fpu.state);
    }
    
    /*
     * Restore current process's FPU state
     */
    if (current->thread.fpu.initialized) {
        fxrstor(&current->thread.fpu.state);
    } else {
        /* First time this process uses FPU - initialize to default */
        fninit();
        current->thread.fpu.initialized = true;
    }
    
    /*
     * Clear TS bit - subsequent FPU instructions won't trap
     * (until next context switch sets it again)
     */
    clts();  /* Clear Task-Switched flag */
    
    /* Mark current process as FPU owner */
    this_cpu_write(fpu_owner, current);
    
    /* Exception handler returns, instruction retries, works now */
}

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

/**
 * switch_mm() - Switch memory management context
 * 
 * Called during context switch when switching between processes
 * (not needed when switching between threads of same process).
 */
void switch_mm(struct mm_struct *prev_mm, 
               struct mm_struct *next_mm,
               struct task_struct *next)
{
    unsigned long cr3;
    
    /* If same address space (threads of same process), skip */
    if (likely(prev_mm == next_mm)) {
        return;
    }
    
    /* 
     * Get physical address of next process's PML4 page table.
     * The pgd (Page Global Directory) is the top-level table.
     */
    cr3 = __pa(next_mm->pgd);
    
    /*
     * Load new CR3 - this switches the entire address space.
     * 
     * CRITICAL SIDE EFFECT: Loading CR3 flushes the TLB!
     * All cached virtual->physical translations are invalidated.
     * This is expensive and a major source of context switch overhead.
     */
    write_cr3(cr3);
    
    /* Update per-CPU tracking of current mm */
    this_cpu_write(cpu_current_mm, next_mm);
    
    /* 
     * PCID optimization (if available):
     * Process Context IDentifiers allow keeping TLB entries
     * from multiple address spaces simultaneously, tagged by PCID.
     * This avoids full TLB flush on context switch.
     */
    if (cpu_has_pcid) {
        /* CR3 includes PCID in lower bits */
        cr3 = build_cr3(next_mm->pgd, next_mm->context.ctx_id);
        write_cr3_pcid(cr3);  /* May not flush TLB */
    }
}
 
/**
 * TLB flush cost example:
 * 
 * - TLB has ~1000+ entries caching page translations
 * - After flush, EVERY memory access causes page table walk
 * - Page table walk: 4-5 memory accesses per translation
 * - Until TLB warms up, process runs much slower
 * 
 * This is why process context switches are more expensive
 * than thread context switches (threads share address space).
 */

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

/**
 * How Linux tracks the current task on x86-64
 * 
 * The 'current' macro is used everywhere in the kernel to get
 * the currently executing task. It must be updated during context switch.
 */
 
/* Per-CPU variable holding pointer to current task */
DEFINE_PER_CPU(struct task_struct *, current_task);
 
/**
 * The 'current' macro uses the GS segment base to access per-CPU data.
 * GS base points to this CPU's per-CPU area.
 */
#define current get_current()
 
static inline struct task_struct *get_current(void)
{
    return this_cpu_read_stable(current_task);
}
 
/**
 * During context switch, update current task pointer
 */
void update_current_task(struct task_struct *next)
{
    /* Update per-CPU current_task to point to next */
    this_cpu_write(current_task, next);
    
    /*
     * From this point on, current == next.
     * Any kernel code executing on this CPU will see
     * 'next' as the current task.
     */
}
 
/**
 * Also update TSS for next interrupt's kernel stack
 */
void update_tss_sp(struct task_struct *next)
{
    struct tss_struct *tss = this_cpu_ptr(&cpu_tss_rw);
    
    /*
     * TSS.SP0 holds the kernel stack top for ring 0.
     * When next interrupt occurs in user mode, CPU loads
     * RSP from here. Must point to next's kernel stack.
     */
    tss->x86_tss.sp0 = (unsigned long)next->stack + THREAD_SIZE;
}