yggdrasil/src/device/nvme/mod.rs

#![allow(missing_docs)]
use core::{mem::size_of, time::Duration};

use abi::error::Error;
use alloc::vec::Vec;
use device_api::Device;
use kernel_util::{sync::IrqSafeSpinlock, util::OneTimeInit};
use tock_registers::{
    interfaces::{ReadWriteable, Readable, Writeable},
    register_bitfields, register_structs,
    registers::{ReadOnly, ReadWrite, WriteOnly},
};

use crate::{
    device::{
        bus::pci::{PciBaseAddress, PciCommandRegister, PciConfigurationSpace},
        nvme::{
            command::IdentifyControllerRequest,
            queue::{CompletionQueueEntry, SubmissionQueueEntry},
        },
    },
    mem::{
        address::{FromRaw, IntoRaw},
        device::DeviceMemoryIo,
        PhysicalAddress,
    },
    task::runtime,
};

use self::{
    command::{CreateIoCompletionQueue, CreateIoSubmissionQueue, SetFeatureRequest},
    error::NvmeError,
    queue::QueuePair,
};

use super::bus::pci::{FromPciBus, PciDeviceInfo};

mod command;
mod error;
mod queue;

register_bitfields! {
    u32,
    CC [
        IOCQES OFFSET(20) NUMBITS(4) [],
        IOSQES OFFSET(16) NUMBITS(4) [],
        AMS OFFSET(11) NUMBITS(3) [],
        MPS OFFSET(7) NUMBITS(4) [],
        CSS OFFSET(4) NUMBITS(3) [
            NvmCommandSet = 0
        ],
        ENABLE OFFSET(0) NUMBITS(1) [],
    ],
    CSTS [
        CFS OFFSET(1) NUMBITS(1) [],
        RDY OFFSET(0) NUMBITS(1) [],
    ],
    AQA [
        /// Admin Completion Queue Size in entries - 1
        ACQS OFFSET(16) NUMBITS(12) [],
        /// Admin Submission Queue Size in entries - 1
        ASQS OFFSET(0) NUMBITS(12) [],
    ]
}

register_bitfields! {
    u64,
    CAP [
        /// Maximum Queue Entries Supported - 1. i.e., 0 means maximum queue len of 1, 1 = 2 etc.
        MQES OFFSET(0) NUMBITS(16) [],
        /// Timeout. Represents the worst-case time the host software should wait for CSTS.RDY to
        /// change its state.
        TO OFFSET(24) NUMBITS(8) [],
        /// Doorbell stride. Stride in bytes = pow(2, 2 + DSTRD).
        DSTRD OFFSET(32) NUMBITS(4) [],
        /// NVM Subsystem Reset Supported (see NVMe BS Section 3.7.1)
        NSSRS OFFSET(36) NUMBITS(1) [],
        /// Controller supports one or more I/O command sets
        CSS_IO_COMMANDS OFFSET(43) NUMBITS(1) [],
        /// Controller only supports admin commands and no I/O commands
        CSS_ADMIN_ONLY OFFSET(44) NUMBITS(1) [],
        /// Memory page size minimum (bytes = pow(2, 12 + MPSMIN))
        MPSMIN OFFSET(48) NUMBITS(4) [],
        /// Memory page size maximum -|-
        MPSMAX OFFSET(52) NUMBITS(4) [],
    ]
}

register_structs! {
    #[allow(non_snake_case)]
    Regs {
        (0x00 => CAP: ReadOnly<u64, CAP::Register>),
        (0x08 => VS: ReadOnly<u32>),
        (0x0C => INTMS: WriteOnly<u32>),
        (0x10 => INTMC: WriteOnly<u32>),
        (0x14 => CC: ReadWrite<u32, CC::Register>),
        (0x18 => _0),
        (0x1C => CSTS: ReadOnly<u32, CSTS::Register>),
        (0x20 => _1),
        (0x24 => AQA: ReadWrite<u32, AQA::Register>),
        (0x28 => ASQ: ReadWrite<u64>),
        (0x30 => ACQ: ReadWrite<u64>),
        (0x38 => _2),
        (0x2000 => @END),
    }
}

pub struct NvmeController {
    regs: IrqSafeSpinlock<DeviceMemoryIo<'static, Regs>>,
    admin_q: OneTimeInit<QueuePair<'static>>,
    ioqs: OneTimeInit<Vec<QueuePair<'static>>>,

    doorbell_shift: usize,
}

impl Regs {
    unsafe fn doorbell_ptr(&self, shift: usize, completion: bool, queue_index: usize) -> *mut u32 {
        let doorbell_base = (self as *const Regs as *mut Regs).addr() + 0x1000;
        let offset = (queue_index << shift) + completion as usize * 4;
        (doorbell_base + offset) as *mut u32
    }
}

impl NvmeController {
    async fn late_init(&'static self) -> Result<(), NvmeError> {
        runtime::spawn(self.poll_task()).expect("Couldn't spawn NVMe poll task");
        let admin_q = self.admin_q.get();

        // Request a CQ/SQ pair for I/O
        admin_q
            .request_no_data(SetFeatureRequest::NumberOfQueues(1, 1))
            .await?;

        // Allocate the queue
        let (sq_doorbell, cq_doorbell) = unsafe { self.doorbell_pair(1) };
        let io_q = QueuePair::new(32, sq_doorbell, cq_doorbell).map_err(NvmeError::MemoryError)?;

        // Identify the controller
        let identify = admin_q
            .request(IdentifyControllerRequest { nsid: 0 })?
            .await?;

        // Create the queue on the device side
        admin_q
            .request_no_data(CreateIoCompletionQueue {
                id: 1,
                size: 32,
                data: io_q.cq_physical_pointer(),
            })
            .await?;
        admin_q
            .request_no_data(CreateIoSubmissionQueue {
                id: 1,
                cq_id: 1,
                size: 32,
                data: io_q.sq_physical_pointer(),
            })
            .await?;

        loop {}
    }

    // TODO MSI(-X) or IRQ (ACPI currently broken) support for PCIe-based NVMe
    async fn poll_task(&'static self) {
        loop {
            self.admin_q.get().process_completions();
            runtime::sleep(Duration::from_millis(100)).await;
        }
    }

    unsafe fn doorbell_pair(&self, idx: usize) -> (*mut u32, *mut u32) {
        let regs = self.regs.lock();
        let sq_ptr = regs.doorbell_ptr(self.doorbell_shift, false, idx);
        let cq_ptr = regs.doorbell_ptr(self.doorbell_shift, true, idx);
        (sq_ptr, cq_ptr)
    }
}

impl Device for NvmeController {
    unsafe fn init(&'static self) -> Result<(), Error> {
        let regs = self.regs.lock();

        let min_page_size = 1usize << (12 + regs.CAP.read(CAP::MPSMIN));
        let max_page_size = 1usize << (12 + regs.CAP.read(CAP::MPSMAX));

        if min_page_size > 4096 {
            panic!();
        }

        let timeout = Duration::from_millis(regs.CAP.read(CAP::TO) * 500);
        debugln!("Worst-case timeout: {:?}", timeout);

        while regs.CSTS.matches_any(CSTS::RDY::SET) {
            core::hint::spin_loop();
        }

        let queue_slots = 32;
        if queue_slots > regs.CAP.read(CAP::MQES) + 1 {
            todo!(
                "queue_slots too big, max = {}",
                regs.CAP.read(CAP::MQES) + 1
            );
        }

        // Setup the admin queue (index 0)
        let admin_sq_doorbell = unsafe { regs.doorbell_ptr(self.doorbell_shift, false, 0) };
        let admin_cq_doorbell = unsafe { regs.doorbell_ptr(self.doorbell_shift, true, 0) };
        let admin_q =
            QueuePair::new(queue_slots as usize, admin_sq_doorbell, admin_cq_doorbell).unwrap();

        regs.AQA
            .modify(AQA::ASQS.val(queue_slots as u32 - 1) + AQA::ACQS.val(queue_slots as u32 - 1));
        regs.ASQ.set(admin_q.sq_physical_pointer().into_raw());
        regs.ACQ.set(admin_q.cq_physical_pointer().into_raw());

        // Configure the controller
        const IOSQES: u32 = size_of::<SubmissionQueueEntry>().ilog2();
        const IOCQES: u32 = size_of::<CompletionQueueEntry>().ilog2();

        regs.CC.modify(
            CC::IOCQES.val(IOCQES)
                + CC::IOSQES.val(IOSQES)
                + CC::MPS.val(0)
                + CC::CSS::NvmCommandSet,
        );

        // Enable the controller
        regs.CC.modify(CC::ENABLE::SET);

        debugln!("Reset the controller");

        while !regs.CSTS.matches_any(CSTS::RDY::SET + CSTS::CFS::SET) {
            core::hint::spin_loop();
        }

        if regs.CSTS.matches_any(CSTS::CFS::SET) {
            todo!("CFS set after reset!");
        }

        self.admin_q.init(admin_q);

        // Schedule late_init task
        runtime::spawn(self.late_init())?;

        Ok(())
    }

    fn display_name(&self) -> &'static str {
        "NVM Express Controller"
    }
}

impl FromPciBus for NvmeController {
    fn from_pci_bus(info: &PciDeviceInfo) -> Result<Self, Error> {
        let PciBaseAddress::Memory(bar0) = info.config_space.bar(0).unwrap() else {
            panic!();
        };

        let mut cmd = PciCommandRegister::from_bits_retain(info.config_space.command());
        cmd &= !(PciCommandRegister::DISABLE_INTERRUPTS | PciCommandRegister::ENABLE_IO);
        cmd |= PciCommandRegister::ENABLE_MEMORY | PciCommandRegister::BUS_MASTER;
        info.config_space.set_command(cmd.bits());

        let regs = unsafe { DeviceMemoryIo::<Regs>::map(PhysicalAddress::from_raw(bar0)) }?;

        // Disable the controller
        regs.CC.modify(CC::ENABLE::CLEAR);

        let doorbell_shift = regs.CAP.read(CAP::DSTRD) as usize + 2;

        Ok(Self {
            regs: IrqSafeSpinlock::new(regs),
            admin_q: OneTimeInit::new(),
            ioqs: OneTimeInit::new(),
            doorbell_shift,
        })
    }
}