Skip to main content

nvme_driver/
nvme.rs

1use alloc::vec::Vec;
2use core::ptr::NonNull;
3
4use dma_api::{CoherentArray, DeviceDma, DmaDirection, DmaOp};
5use log::{debug, info};
6use mmio_api::{Mmio, MmioAddr, MmioOp};
7
8use crate::{
9    command::{
10        self, ControllerInfo, Feature, Identify, IdentifyActiveNamespaceList, IdentifyController,
11        IdentifyNamespaceDataStructure,
12    },
13    err::*,
14    queue::{CommandSet, NvmeQueue},
15    registers::NvmeReg,
16};
17
18pub struct Nvme {
19    bar: NonNull<NvmeReg>,
20    _mmio: Option<Mmio>,
21    dma: DeviceDma,
22    admin_queue: NvmeQueue,
23    io_queues: Vec<Option<NvmeQueue>>,
24    num_ns: usize,
25    sqes: u32,
26    cqes: u32,
27    page_size: usize,
28    max_transfer_bytes: Option<usize>,
29    io_queue_interrupts: bool,
30    msix_interrupts: bool,
31    interrupt_vectors: Vec<u16>,
32}
33
34#[derive(Debug, Clone)]
35pub struct Config {
36    pub page_size: usize,
37    pub io_queue_pair_count: usize,
38    pub io_queue_interrupts: bool,
39    pub interrupt_vector: u32,
40    pub msix_interrupts: bool,
41    pub interrupt_vectors: Vec<u16>,
42}
43
44impl Config {
45    pub const fn new(page_size: usize, io_queue_pair_count: usize) -> Self {
46        Self {
47            page_size,
48            io_queue_pair_count,
49            io_queue_interrupts: false,
50            interrupt_vector: 0,
51            msix_interrupts: false,
52            interrupt_vectors: Vec::new(),
53        }
54    }
55
56    pub fn with_intx_irq(mut self) -> Self {
57        self.io_queue_interrupts = true;
58        self.interrupt_vector = 0;
59        self.msix_interrupts = false;
60        self.interrupt_vectors = Vec::from([0]);
61        self
62    }
63
64    pub fn with_msix_vectors(mut self, vectors: impl Into<Vec<u16>>) -> Self {
65        self.interrupt_vectors = vectors.into();
66        self.io_queue_interrupts = !self.interrupt_vectors.is_empty();
67        self.msix_interrupts = self.io_queue_interrupts;
68        self.interrupt_vector = self
69            .interrupt_vectors
70            .first()
71            .copied()
72            .map(u32::from)
73            .unwrap_or(0);
74        self
75    }
76
77    fn interrupt_vector_for_queue(&self, queue_index: usize) -> u32 {
78        self.interrupt_vectors
79            .get(queue_index)
80            .copied()
81            .map(u32::from)
82            .unwrap_or(self.interrupt_vector)
83    }
84}
85
86impl Nvme {
87    pub fn new(
88        bar_addr: impl Into<MmioAddr>,
89        bar_size: usize,
90        dma_mask: u64,
91        dma_op: &'static dyn DmaOp,
92        mmio_op: &'static dyn MmioOp,
93        config: Config,
94    ) -> Result<Self> {
95        mmio_api::init(mmio_op);
96        let mmio = mmio_api::ioremap(bar_addr.into(), bar_size)?;
97        let dma = DeviceDma::new_legacy(dma_mask, dma_op);
98        Self::new_mmio(mmio, dma, config)
99    }
100
101    fn new_mmio(mmio: Mmio, dma: DeviceDma, config: Config) -> Result<Self> {
102        let bar = NonNull::new(mmio.as_ptr()).expect("mmio mapping must not be null");
103        Self::new_with_bar(bar.cast(), Some(mmio), dma, config)
104    }
105
106    fn new_with_bar(
107        bar: NonNull<NvmeReg>,
108        mmio: Option<Mmio>,
109        dma: DeviceDma,
110        config: Config,
111    ) -> Result<Self> {
112        let admin_queue = NvmeQueue::new(0, bar, &dma, config.page_size, 64, 64)?;
113
114        assert!(config.io_queue_pair_count > 0);
115
116        let mut s = Self {
117            bar,
118            _mmio: mmio,
119            dma,
120            admin_queue,
121            io_queues: Vec::new(),
122            num_ns: 0,
123            sqes: 6,
124            cqes: 4,
125            page_size: config.page_size,
126            max_transfer_bytes: None,
127            io_queue_interrupts: config.io_queue_interrupts,
128            msix_interrupts: config.msix_interrupts,
129            interrupt_vectors: config.interrupt_vectors.clone(),
130        };
131
132        let version = s.version();
133
134        info!(
135            "NVME @{bar:?} init begin, version: {}.{}.{} ",
136            version.0, version.1, version.2
137        );
138
139        s.init(config)?;
140
141        Ok(s)
142    }
143
144    pub fn dma_mask(&self) -> u64 {
145        self.dma.dma_mask()
146    }
147
148    fn reset(&mut self) {
149        self.reg().reset();
150    }
151
152    fn reset_and_setup_controller_info(&mut self) -> Result<ControllerInfo> {
153        self.reset();
154        self.nvme_configure_admin_queue();
155        self.reg().ready_for_read_controller_info();
156
157        self.get_identfy(IdentifyController::new())
158    }
159
160    fn init(&mut self, config: Config) -> Result {
161        let controller = self.reset_and_setup_controller_info()?;
162
163        debug!("Controller: {:?}", controller);
164
165        self.sqes = controller.sqes_min as _;
166        self.cqes = controller.cqes_min as _;
167        self.reset();
168        self.nvme_configure_admin_queue();
169        self.reg().setup_cc(self.sqes, self.cqes);
170        let controller = self.get_identfy(IdentifyController::new())?;
171
172        debug!("Controller: {:?}", controller);
173
174        self.num_ns = controller.number_of_namespaces as _;
175        self.max_transfer_bytes = controller_max_transfer_bytes(config.page_size, controller.mdts);
176        if config.io_queue_interrupts {
177            for vector in &config.interrupt_vectors {
178                self.mask_interrupt_vector(u32::from(*vector));
179            }
180        }
181        self.config_io_queue(config)?;
182
183        debug!("IO queue ok.");
184        loop {
185            let ns = self.get_identfy(IdentifyNamespaceDataStructure::new(1))?;
186            if let Some(ns) = ns {
187                debug!("Namespace: {:?}", ns);
188                break;
189            }
190        }
191        debug!("Namespace ok.");
192        Ok(())
193    }
194
195    pub fn namespace_list(&mut self) -> Result<Vec<Namespace>> {
196        let id_list = self.get_identfy(IdentifyActiveNamespaceList::new())?;
197        let mut out = Vec::new();
198
199        for id in id_list {
200            let ns = self
201                .get_identfy(IdentifyNamespaceDataStructure::new(id))?
202                .unwrap();
203
204            out.push(Namespace {
205                id,
206                lba_size: ns.lba_size as _,
207                lba_count: ns.namespace_size as _,
208                metadata_size: ns.metadata_size as _,
209            });
210        }
211
212        Ok(out)
213    }
214
215    // config admin queue
216    // 1. set admin queue(cq && sq) size
217    // 2. set admin queue(cq && sq) dma address
218    // 3. enable ctrl
219    fn nvme_configure_admin_queue(&mut self) {
220        self.reg().set_admin_submission_and_completion_queue_size(
221            self.admin_queue.sq_len(),
222            self.admin_queue.cq_len(),
223        );
224
225        self.reg()
226            .set_admin_submission_queue_base_address(self.admin_queue.sq_bus_addr());
227
228        self.reg()
229            .set_admin_completion_queue_base_address(self.admin_queue.cq_bus_addr());
230    }
231
232    fn config_io_queue(&mut self, config: Config) -> Result {
233        let num = config.io_queue_pair_count;
234        // 设置 io queue 数量
235        let cmd = CommandSet::set_features(Feature::NumberOfQueues {
236            nsq: num as u32 - 1,
237            ncq: num as u32 - 1,
238        });
239        self.admin_queue.command_sync(cmd)?;
240
241        for i in 0..num {
242            let id = (i + 1) as u32;
243            let io_queue = NvmeQueue::new(
244                id,
245                self.bar,
246                &self.dma,
247                config.page_size,
248                2usize.pow(self.sqes as _),
249                2usize.pow(self.cqes as _),
250            )?;
251
252            let data = CommandSet::create_io_completion_queue(
253                io_queue.qid,
254                io_queue.cq_len() as _,
255                io_queue.cq_bus_addr(),
256                true,
257                config.io_queue_interrupts,
258                config.interrupt_vector_for_queue(i),
259            );
260            self.admin_queue.command_sync(data)?;
261
262            let data = CommandSet::create_io_submission_queue(
263                io_queue.qid,
264                io_queue.sq_len() as _,
265                io_queue.sq_bus_addr(),
266                true,
267                0,
268                io_queue.qid,
269                0,
270            );
271
272            self.admin_queue.command_sync(data)?;
273
274            self.io_queues.push(Some(io_queue));
275        }
276
277        Ok(())
278    }
279
280    pub fn io_queue_count(&self) -> usize {
281        self.io_queues.len()
282    }
283
284    pub fn page_size(&self) -> usize {
285        self.page_size
286    }
287
288    pub(crate) const fn max_transfer_bytes(&self) -> Option<usize> {
289        self.max_transfer_bytes
290    }
291
292    pub fn io_queue_interrupts_enabled(&self) -> bool {
293        self.io_queue_interrupts
294    }
295
296    pub fn interrupt_vector(&self) -> u32 {
297        self.interrupt_vectors
298            .first()
299            .copied()
300            .map(u32::from)
301            .unwrap_or(0)
302    }
303
304    pub fn msix_interrupts_enabled(&self) -> bool {
305        self.io_queue_interrupts && self.msix_interrupts
306    }
307
308    pub fn interrupt_vectors(&self) -> &[u16] {
309        &self.interrupt_vectors
310    }
311
312    pub fn mask_interrupt_vector(&mut self, vector: u32) {
313        self.reg().mask_interrupt_vector(vector);
314    }
315
316    pub fn unmask_interrupt_vector(&mut self, vector: u32) {
317        self.reg().unmask_interrupt_vector(vector);
318    }
319
320    pub(crate) fn take_io_queue(&mut self, index: usize) -> Option<NvmeQueue> {
321        self.io_queues.get_mut(index)?.take()
322    }
323
324    pub(crate) fn alloc_prp_list(&self) -> Result<CoherentArray<u64>> {
325        self.dma
326            .coherent_array_zero_with_align(
327                self.page_size / core::mem::size_of::<u64>(),
328                self.page_size,
329            )
330            .map_err(Into::into)
331    }
332
333    pub fn get_identfy<T: Identify>(&mut self, mut want: T) -> Result<T::Output> {
334        let cmd = want.command_set_mut();
335
336        cmd.cdw0 = CommandSet::cdw0_from_opcode(command::Opcode::IDENTIFY);
337        cmd.cdw10 = T::CNS;
338
339        let buff = self.dma.contiguous_array_zero_with_align::<u8>(
340            0x1000,
341            0x1000,
342            DmaDirection::FromDevice,
343        )?;
344        cmd.prp1 = buff.dma_addr().as_u64();
345
346        self.admin_queue.command_sync(*cmd)?;
347
348        let data = buff.read_from_device(buff.len(), |data| data.to_vec());
349        let res = want.parse(&data);
350        Ok(res)
351    }
352
353    pub fn block_write_sync(
354        &mut self,
355        ns: &Namespace,
356        block_start: u64,
357        buff: &[u8],
358    ) -> Result<()> {
359        assert!(
360            buff.len().is_multiple_of(ns.lba_size),
361            "buffer size must be multiple of lba size"
362        );
363
364        let mut dma_buff = self.dma.contiguous_array_zero_with_align::<u8>(
365            buff.len(),
366            ns.lba_size,
367            DmaDirection::ToDevice,
368        )?;
369        dma_buff.copy_to_device_from_slice(buff);
370
371        let blk_num = dma_buff.len() / ns.lba_size;
372
373        let cmd = CommandSet::nvm_cmd_write(
374            ns.id,
375            dma_buff.dma_addr().as_u64(),
376            block_start,
377            blk_num as _,
378        );
379
380        self.io_queues
381            .get_mut(0)
382            .and_then(Option::as_mut)
383            .ok_or(Error::Unknown("missing IO queue"))?
384            .command_sync(cmd)?;
385
386        Ok(())
387    }
388
389    pub fn block_read_sync(
390        &mut self,
391        ns: &Namespace,
392        block_start: u64,
393        buff: &mut [u8],
394    ) -> Result<()> {
395        assert!(
396            buff.len().is_multiple_of(ns.lba_size),
397            "buffer size must be multiple of lba size"
398        );
399
400        let dma_buff = self.dma.contiguous_array_zero_with_align::<u8>(
401            buff.len(),
402            ns.lba_size,
403            DmaDirection::FromDevice,
404        )?;
405
406        let blk_num = dma_buff.len() / ns.lba_size;
407
408        let cmd = CommandSet::nvm_cmd_read(
409            ns.id,
410            dma_buff.dma_addr().as_u64(),
411            block_start,
412            blk_num as _,
413        );
414
415        self.io_queues
416            .get_mut(0)
417            .and_then(Option::as_mut)
418            .ok_or(Error::Unknown("missing IO queue"))?
419            .command_sync(cmd)?;
420        dma_buff.copy_from_device_to_slice(buff);
421        Ok(())
422    }
423
424    pub fn version(&self) -> (usize, usize, usize) {
425        self.reg().version()
426    }
427
428    fn reg(&self) -> &NvmeReg {
429        unsafe { self.bar.as_ref() }
430    }
431}
432
433unsafe impl Send for Nvme {}
434
435fn controller_max_transfer_bytes(page_size: usize, mdts: u8) -> Option<usize> {
436    if mdts == 0 {
437        None
438    } else {
439        Some(page_size.checked_shl(u32::from(mdts)).unwrap_or(usize::MAX))
440    }
441}
442
443#[derive(Debug, Clone, Copy)]
444pub struct Namespace {
445    pub id: u32,
446    pub lba_size: usize,
447    pub lba_count: usize,
448    pub metadata_size: usize,
449}
450
451#[cfg(test)]
452mod tests {
453    use super::{Config, controller_max_transfer_bytes};
454
455    #[test]
456    fn config_defaults_to_polling_and_can_enable_intx() {
457        let config = Config::new(4096, 1);
458        assert!(!config.io_queue_interrupts);
459        assert_eq!(config.interrupt_vector, 0);
460        assert!(!config.msix_interrupts);
461        assert!(config.interrupt_vectors.is_empty());
462
463        let irq_config = config.with_intx_irq();
464        assert!(irq_config.io_queue_interrupts);
465        assert_eq!(irq_config.interrupt_vector, 0);
466        assert!(!irq_config.msix_interrupts);
467        assert_eq!(irq_config.interrupt_vectors, [0]);
468    }
469
470    #[test]
471    fn config_can_enable_msix_per_queue_vectors() {
472        let config = Config::new(4096, 2).with_msix_vectors([4, 5]);
473
474        assert!(config.io_queue_interrupts);
475        assert!(config.msix_interrupts);
476        assert_eq!(config.interrupt_vector, 4);
477        assert_eq!(config.interrupt_vector_for_queue(0), 4);
478        assert_eq!(config.interrupt_vector_for_queue(1), 5);
479    }
480
481    #[test]
482    fn controller_mdts_zero_means_unrestricted_transfer_size() {
483        assert_eq!(controller_max_transfer_bytes(4096, 0), None);
484    }
485
486    #[test]
487    fn controller_mdts_scales_with_controller_page_size() {
488        assert_eq!(controller_max_transfer_bytes(4096, 7), Some(512 * 1024));
489    }
490}