task: assign a PCID to each task page-table root

Process Context IDs let the CPU tag TLB entries with the active
page-table root. Without that, every context switch that reloads CR3
must invalidate the whole TLB, even when the outgoing and incoming
tasks are unrelated and their mappings could safely coexist in the
cache. Giving each live task root its own PCID (1..4095; 0 reserved
for per-CPU transitional tables) is the foundation for cheaper
switches later and keeps stale translations from one dead root from
matching a reused tag.This is also part of my GSoC project to
add support for X86 Process Context Identifiers (PCID) in COCONUT-SVSM.

Each task root gets a TaskPcid from a global pool of 4096 IDs (PCID 0
stays reserved for per-CPU transitional page tables). The allocator keeps a
simple in_use bitmap: allocation scans from 1 upward for the first free
slot and marks it taken. Freeing is tied to task lifetime when the
Task is destroyed, TaskPcid’s Drop runs release(), which first
broadcasts a PCID-targeted TLB shootdown on every CPU
(flush_tlb_pcid_sync, INVPCID on native / INVLPGB+TLBSYNC on SNP) so no
stale translations remain tagged with that ID, then clears the in_use bit
so the number can be reused. That guarantees the invariant that any PCID
loaded into CR3 for a new root has no valid entries under its tag anywhere
in the system.

Signed-off-by: tanish111 <tanishdesai37@gmail.com>

Author: tanish111

kernel/src/cpu/control_regs.rs

@@ -37,6 +37,10 @@ pub fn cr4_init() {
 
     cr4.insert(CR4Flags::PGE); // Enable Global Pages
 
+    if cpu_has_feat(Feature::Pcid) {
+        cr4.insert(CR4Flags::PCIDE);
+    }
+
     if !cfg!(feature = "nosmep") {
         assert!(cpu_has_feat(Feature::Smep), "CPU does not support SMEP");
         cr4.insert(CR4Flags::SMEP);

kernel/src/cpu/features.rs

@@ -170,5 +170,8 @@ define_cpu_feats! {
     HyperV => CpuFeat::new_u32(0x40000001, CpuidReg::Eax, 0x31237648),
     PhysAddrSizes => CpuFeat::new_u32(0x80000008, CpuidReg::Eax, 0),
     InvlpgbMax => CpuFeat::new(0x80000008, CpuidReg::Edx, 0, u16::BITS as u8, 0),
+    Invlpgb => CpuFeat::new_bit(0x80000008, CpuidReg::Ebx, 3),
     Cbit => CpuFeat::new(0x8000001f, CpuidReg::Ebx, 0, 6, 0),
+    Pcid => CpuFeat::new_bit(0x0000_0001, CpuidReg::Ecx, 17),
+    Invpcid => CpuFeat::new_bit(0x0000_0007, CpuidReg::Ebx, 24),
 }

kernel/src/cpu/tlb.rs

@@ -6,6 +6,7 @@
 
 use crate::address::{Address, VirtAddr};
 use crate::cpu::control_regs::{read_cr3, read_cr4, write_cr3, write_cr4};
+use crate::cpu::features::{Feature, cpu_has_feat};
 use crate::cpu::ipi::{IpiMessage, IpiTarget, send_multicast_ipi};
 use crate::platform::SVSM_PLATFORM;
 use crate::types::PageSize;
@@ -21,6 +22,13 @@ static FLUSH_SMP: AtomicBool = AtomicBool::new(false);
 /// be performed instead.
 const TLB_FLUSH_ALL_THRESHOLD: usize = 256;
 
+/// INVPCID descriptor (Intel SDM Vol. 2A, `INVPCID`).
+#[repr(C, align(16))]
+struct InvpcidDesc {
+    pcid: u64,
+    reserved: u64,
+}
+
 /// Defines the scope of a TLB flush.
 #[derive(Copy, Clone, Debug)]
 pub enum TlbFlushRange {
@@ -40,6 +48,9 @@ pub struct TlbFlushScope {
     pub global: bool,
     /// Indicates the range of virtual addresses that should be flushed.
     pub range: TlbFlushRange,
+    /// When set, flush all non-global entries for this PCID (INVPCID or
+    /// INVLPGB, depending on CPU features).
+    pub pcid: Option<u16>,
 }
 
 impl TlbFlushScope {
@@ -48,6 +59,16 @@ impl TlbFlushScope {
         Self {
             global: true,
             range: TlbFlushRange::All,
+            pcid: None,
+        }
+    }
+
+    /// Flush all non-global TLB entries tagged with the given PCID.
+    pub const fn pcid(pcid: u16) -> Self {
+        Self {
+            global: true,
+            range: TlbFlushRange::All,
+            pcid: Some(pcid),
         }
     }
 
@@ -63,6 +84,7 @@ impl TlbFlushScope {
         Self {
             global: true,
             range: TlbFlushRange::Range { region, pgsize },
+            pcid: None,
         }
     }
 
@@ -94,6 +116,10 @@ impl TlbFlushScope {
 
     /// Flushes all the entries in the TLB for the current CPU.
     fn flush_percpu_all(&self) {
+        if let Some(pcid) = self.pcid {
+            flush_pcid_percpu(pcid);
+            return;
+        }
         match self.global {
             true => __flush_tlb_global_percpu(),
             false => __flush_tlb_percpu(),
@@ -130,6 +156,10 @@ pub fn set_tlb_flush_smp() {
     FLUSH_SMP.store(true, Ordering::Relaxed);
 }
 
+pub fn flush_tlb_pcid_sync(pcid: u16) {
+    TlbFlushScope::pcid(pcid).flush_all_cpus();
+}
+
 pub fn flush_tlb_global_sync() {
     TlbFlushScope::all().with_global(true).flush_all_cpus();
 }
@@ -166,6 +196,24 @@ pub fn flush_tlb_percpu() {
     TlbFlushScope::all().with_global(false).flush_percpu();
 }
 
+fn flush_pcid_percpu(pcid: u16) {
+    if cpu_has_feat(Feature::Invpcid) {
+        let desc = InvpcidDesc {
+            pcid: u64::from(pcid),
+            reserved: 0,
+        };
+        // SAFETY: INVPCID type 1 — type in any GPR, descriptor in memory.
+        unsafe {
+            asm!("invpcid ({0}), {1}",
+                 in(reg) &raw const desc,
+                 in(reg) 1u64,
+                 options(att_syntax));
+        }
+    } else {
+        __flush_tlb_global_percpu();
+    }
+}
+
 fn __flush_tlb_global_percpu() {
     let cr4 = read_cr4();
 

kernel/src/sev/tlb.rs

@@ -113,6 +113,15 @@ pub fn flush_address_sync(va: VirtAddr) {
 }
 
 pub fn flush_tlb_scope(flush_scope: &TlbFlushScope) {
+    if let Some(pcid) = flush_scope.pcid {
+        let rax = InvlpgbRax::new()
+            .with_valid_asid(true)
+            .with_valid_pcid(true);
+        // EDX: PCID in bits 31:16, ASID in bits 15:0 (Linux: pcid << 16 | asid).
+        do_invlpgb(rax.into_bits(), 0, u64::from(pcid) << 16);
+        do_tlbsync();
+        return;
+    }
     match flush_scope.range {
         TlbFlushRange::All => flush_tlb_sync(flush_scope.global),
         TlbFlushRange::Range { region, pgsize } => {

kernel/src/task/mod.rs

@@ -5,6 +5,7 @@
 // Author: Roy Hopkins <rhopkins@suse.de>
 
 mod exec;
+mod pcid;
 mod schedule;
 mod task_mm;
 mod tasks;

kernel/src/task/pcid.rs

@@ -0,0 +1,76 @@
+// SPDX-License-Identifier: MIT
+//
+// Copyright (c) 2026 Tanish Desai
+//
+// Author: Tanish Desai
+
+use crate::cpu::features::{Feature, cpu_has_feat};
+use crate::cpu::tlb::flush_tlb_pcid_sync;
+use crate::error::SvsmError;
+use crate::locking::SpinLock;
+use crate::mm::alloc::AllocError;
+
+const PCID_COUNT: usize = 4096;
+
+static PCID_ALLOC: SpinLock<PcidAllocator> = SpinLock::new(PcidAllocator::new());
+
+/// True when PCIDs can be used for task page-table roots.
+///
+/// Requires `Pcid` (CR4.PCIDE) plus a targeted flush path: INVPCID (`Invpcid`)
+/// or AMD INVLPGB (`Invlpgb`).
+pub fn pcid_supported() -> bool {
+    cpu_has_feat(Feature::Pcid)
+        && (cpu_has_feat(Feature::Invpcid) || cpu_has_feat(Feature::Invlpgb))
+}
+
+#[derive(Debug)]
+struct PcidAllocator {
+    in_use: [bool; PCID_COUNT],
+}
+
+impl PcidAllocator {
+    const fn new() -> Self {
+        Self {
+            in_use: [false; PCID_COUNT],
+        }
+    }
+
+    fn alloc(&mut self) -> Result<u16, SvsmError> {
+        for pcid in 1..PCID_COUNT as u16 {
+            let idx = usize::from(pcid);
+            if !self.in_use[idx] {
+                self.in_use[idx] = true;
+                return Ok(pcid);
+            }
+        }
+        Err(SvsmError::Alloc(AllocError::OutOfMemory))
+    }
+
+    fn release(&mut self, pcid: u16) {
+        flush_tlb_pcid_sync(pcid);
+        self.in_use[usize::from(pcid)] = false;
+    }
+}
+
+/// Owned PCID for a task page-table root.
+#[derive(Debug)]
+pub struct TaskPcid {
+    pcid: u16,
+}
+
+impl TaskPcid {
+    pub fn new() -> Result<Self, SvsmError> {
+        let pcid = PCID_ALLOC.lock().alloc()?;
+        Ok(Self { pcid })
+    }
+
+    pub fn pcid(&self) -> u16 {
+        self.pcid
+    }
+}
+
+impl Drop for TaskPcid {
+    fn drop(&mut self) {
+        PCID_ALLOC.lock().release(self.pcid);
+    }
+}

kernel/src/task/schedule.rs

@@ -526,7 +526,13 @@ unsafe fn switch_to(prev_task: Option<TaskPointer>, next_task: TaskPointer) -> O
     // the page table and stack information in those tasks are correct and
     // can be used to switch to the correct page table and execution stack.
     unsafe {
-        let cr3 = (*next).page_table.lock().cr3_value().bits();
+        let cr3 = {
+            let mut cr3 = (*next).page_table.lock().cr3_value().bits();
+            if let Some(pcid) = (*next).pcid() {
+                cr3 |= usize::from(pcid);
+            }
+            cr3
+        };
 
         // Switch to new task
         let new_prev = switch_context(

kernel/src/task/tasks.rs

@@ -51,6 +51,7 @@ use crate::utils::{MemoryRegion, is_aligned};
 use intrusive_collections::{LinkedListAtomicLink, intrusive_adapter};
 
 use super::WaitQueue;
+use super::pcid::{TaskPcid, pcid_supported};
 use super::schedule::complete_task_switch;
 use super::schedule::terminate;
 use super::task_mm::{TaskKernelMapping, TaskMM};
@@ -289,6 +290,9 @@ pub struct Task {
     /// Page table that is loaded when the task is scheduled
     pub page_table: SpinLock<PageBox<PageTable>>,
 
+    /// PCID tagging this task's page-table root
+    pcid: Option<TaskPcid>,
+
     /// Task kernel stack mapping
     _kernel_stack: TaskKernelMapping,
 
@@ -396,6 +400,14 @@ impl Task {
 
         cpu.populate_page_table(&mut pgtable);
 
+        // Allocate a PCID for this root before the task can be registered on
+        // TASKLIST (done by the schedule.rs wrappers after we return).
+        let pcid = if pcid_supported() {
+            Some(TaskPcid::new()?)
+        } else {
+            None
+        };
+
         let (task_mm, objtree) = {
             if let Some(parent_thread) = args.thread_of {
                 (parent_thread.mm.clone(), parent_thread.objs.clone())
@@ -484,6 +496,7 @@ impl Task {
             stack_bounds: bounds,
             shadow_stack_base,
             page_table: SpinLock::new(pgtable),
+            pcid,
             _kernel_stack: kernel_stack_mapping,
             _shadow_stack: shadow_stack_mapping,
             mm: task_mm,
@@ -581,6 +594,10 @@ impl Task {
         self.rootdir.clone()
     }
 
+    pub fn pcid(&self) -> Option<u16> {
+        self.pcid.as_ref().map(TaskPcid::pcid)
+    }
+
     pub fn set_task_active(&self) {
         self.sched_state.set_active();
     }