feat(hpu): backend can now read/write HPU registers using mmapped BAR0 segment

Author: pgardratzama

backends/tfhe-hpu-backend/Cargo.toml

@@ -53,7 +53,7 @@ ipc-channel = "0.18.3"
 num-traits = { version = "0.2", optional = true }
 clap = { version = "4.4.4", features = ["derive"], optional = true }
 clap-num = { version = "1.1.1", optional = true }
-nix = { version = "0.29.0", features = ["ioctl", "uio", "fs"] }
+nix = { version = "0.29.0", features = ["mman", "ioctl", "uio", "fs"] }
 
 # Dependencies used for rtl_graph features
 dot2 = { version = "1.0", optional = true }

backends/tfhe-hpu-backend/src/ffi/mod.rs

@@ -228,6 +228,11 @@ impl HpuHw {
     pub fn iop_ack_rd(&mut self) -> u32 {
         self.0.ami.iop_ackq_rd()
     }
+
+    #[cfg(feature = "hw-v80")]
+    pub fn map_bar_reg(&mut self) -> Result<(), Box<dyn std::error::Error>> {
+        self.0.ami.map_bar_reg()
+    }
 }
 
 pub struct MemZone(

backends/tfhe-hpu-backend/src/ffi/v80/ami.rs

@@ -3,14 +3,16 @@
 //! AMI driver is used to issue gcq command to the RPU
 //! Those command are used for configuration and register R/W
 use lazy_static::lazy_static;
-use libc;
 use std::error::Error;
 use std::fs::{File, OpenOptions};
 use std::io::{BufRead, BufReader, Read};
 use std::os::fd::AsRawFd;
 use std::ptr::NonNull;
 use std::sync::atomic::{AtomicU32, Ordering};
 use std::time::Duration;
+use nix::sys::mman::{mmap, munmap, MapFlags, ProtFlags};
+use std::num::NonZero;
+use std::os::unix::fs::OpenOptionsExt;
 
 const AMI_VERSION_FILE: &str = "/sys/module/ami/version";
 const AMI_VERSION_PATTERN: &str = r"3\.2\.\d+-zama";
@@ -81,6 +83,7 @@ impl AmiInfo {
 
 pub struct AmiDriver {
     ami_dev: File,
+    bar_reg_ptr: Option<NonNull<u8>>,
     iop_ack_atomic_ptr: NonNull<AtomicU32>,
     retry_rate: Duration,
 }
@@ -100,6 +103,7 @@ impl AmiDriver {
             .read(true)
             .write(true)
             .create(false)
+            .custom_flags(libc::O_SYNC)
             .open(ami_path)?;
 
         let ami_proc_path = format!("/proc/ami_iop_ack_{}", ami_info.devn);
@@ -109,43 +113,56 @@ impl AmiDriver {
             .create(false)
             .open(&ami_proc_path)
             .unwrap();
-        unsafe {
-            let addr = libc::mmap(
-                std::ptr::null_mut(),
-                4096,
-                libc::PROT_READ | libc::PROT_WRITE,
-                libc::MAP_SHARED,
-                ami_proc.as_raw_fd(),
+
+        let addr = unsafe {
+            mmap(
+                None,
+                NonZero::new(4096 as usize).unwrap(),
+                ProtFlags::PROT_READ | ProtFlags::PROT_WRITE,
+                MapFlags::MAP_SHARED,
+                &ami_proc,
                 0,
-            );
+            )?
+        };
 
-            if addr == libc::MAP_FAILED {
-                return Err(format!("mmap on ami_iop_ack_{} failed", ami_info.devn).into());
-            }
+        let iop_ack_atomic_ptr: NonNull<AtomicU32> = addr.cast();
 
-            Ok(Self {
-                ami_dev,
-                iop_ack_atomic_ptr: NonNull::new_unchecked(addr as *mut AtomicU32),
-                retry_rate,
-            })
-        }
+        Ok(Self {
+            ami_dev,
+            bar_reg_ptr: None,
+            iop_ack_atomic_ptr,
+            retry_rate,
+        })
+    }
+
+    pub fn map_bar_reg(&mut self) -> Result<(), Box<dyn std::error::Error>> {
+        let length: usize = 0x140000;
+
+        let map_addr = unsafe {
+            mmap(
+                None,
+                NonZero::new(length).unwrap(),
+                ProtFlags::PROT_READ | ProtFlags::PROT_WRITE, // Read & Write
+                MapFlags::MAP_SHARED,
+                &self.ami_dev,
+                0, // Offset in BAR0
+            )?
+        };
+        tracing::info!("mapping HPU BAR0 at address -> {:p}", map_addr);
+
+        let bar_addr: NonNull<u8> = map_addr.cast();
+        self.bar_reg_ptr = Some(bar_addr);
+
+        Ok(())
     }
 
+
     pub fn munmap_cnt(&self) -> Result<(), Box<dyn Error>> {
-        let cnt_addr = self.iop_ack_atomic_ptr.as_ptr() as *mut libc::c_void;
+        let cnt_addr = self.iop_ack_atomic_ptr.cast();
         unsafe {
-            let retc = libc::munmap(cnt_addr, 4096);
-
-            if retc == 0 {
-                Ok(())
-            } else {
-                Err(format!(
-                    "Could not unmap the atomic shared cnt (mumap returned {})",
-                    retc
-                )
-                .into())
-            }
+            munmap(cnt_addr, 4096)?;
         }
+        Ok(())
     }
 
     /// Read currently loaded UUID in BAR
@@ -276,24 +293,32 @@ impl AmiDriver {
         let data = Box::<u32>::new(0xdeadc0de);
         let data_ptr = Box::into_raw(data);
 
-        // Populate payload
-        let payload = AmiPeakPokePayload {
-            data_ptr,
-            len: 0x1,
-            offset: addr as u32,
-        };
-
-        tracing::trace!("AMI: Read request with following payload {payload:x?}");
-        loop {
-            let ret = unsafe { ami_peak(ami_fd, &payload) };
-            match ret {
-                Err(err) => {
-                    tracing::debug!("AMI: Read failed -> {err:?}");
-                    std::thread::sleep(self.retry_rate);
-                }
-                Ok(val) => {
-                    tracing::trace!("AMI: Read ack received {payload:x?} -> {val:?}");
-                    break;
+        if let Some(base) = self.bar_reg_ptr {
+            unsafe {
+                let raw_base = base.as_ptr();
+                let reg_ptr = raw_base.add((addr + 0x100000).try_into().unwrap()) as *const u32;
+                *data_ptr = std::ptr::read_volatile(reg_ptr);
+            }
+        } else {
+            // Populate payload
+            let payload = AmiPeakPokePayload {
+                data_ptr,
+                len: 0x1,
+                offset: addr as u32,
+            };
+
+            tracing::trace!("AMI: Read request with following payload {payload:x?}");
+            loop {
+                let ret = unsafe { ami_peak(ami_fd, &payload) };
+                match ret {
+                    Err(err) => {
+                        tracing::debug!("AMI: Read failed -> {err:?}");
+                        std::thread::sleep(self.retry_rate);
+                    }
+                    Ok(val) => {
+                        tracing::trace!("AMI: Read ack received {payload:x?} -> {val:?}");
+                        break;
+                    }
                 }
             }
         }
@@ -307,24 +332,32 @@ impl AmiDriver {
         let data = Box::<u32>::new(value);
         let data_ptr = Box::into_raw(data);
 
-        // Populate payload
-        let payload = AmiPeakPokePayload {
-            data_ptr,
-            len: 0x1,
-            offset: addr as u32,
-        };
-
-        tracing::trace!("AMI: Write request with following payload {payload:x?}");
-        loop {
-            let ret = unsafe { ami_poke(ami_fd, &payload) };
-            match ret {
-                Err(err) => {
-                    tracing::debug!("AMI: Write failed -> {err:?}");
-                    std::thread::sleep(self.retry_rate);
-                }
-                Ok(val) => {
-                    tracing::trace!("AMI: Write ack received {payload:x?} -> {val:?}");
-                    break;
+        if let Some(base) = self.bar_reg_ptr {
+            unsafe {
+                let raw_base = base.as_ptr();
+                let reg_ptr = raw_base.add((addr + 0x100000).try_into().unwrap()) as *mut u32;
+                std::ptr::write_volatile(reg_ptr, value);
+            }
+        } else {
+            // Populate payload
+            let payload = AmiPeakPokePayload {
+                data_ptr,
+                len: 0x1,
+                offset: addr as u32,
+            };
+
+            tracing::trace!("AMI: Write request with following payload {payload:x?}");
+            loop {
+                let ret = unsafe { ami_poke(ami_fd, &payload) };
+                match ret {
+                    Err(err) => {
+                        tracing::debug!("AMI: Write failed -> {err:?}");
+                        std::thread::sleep(self.retry_rate);
+                    }
+                    Ok(val) => {
+                        tracing::trace!("AMI: Write ack received {payload:x?} -> {val:?}");
+                        break;
+                    }
                 }
             }
         }

backends/tfhe-hpu-backend/src/interface/backend.rs

@@ -24,6 +24,10 @@ pub struct HpuBackend {
 
     // Extracted parameters
     pub(crate) params: HpuParameters,
+    #[cfg(feature = "hw-v80")]
+    hpu_version_major: u32,
+    #[cfg(feature = "hw-v80")]
+    hpu_version_minor: u32,
     // Prevent to parse regmap at each polling iteration
     #[cfg(not(feature = "hw-v80"))]
     workq_addr: u64,
@@ -109,6 +113,18 @@ impl HpuBackend {
         let regmap = hw_regmap::FlatRegmap::from_file(&regmap_str);
         let mut params = HpuParameters::from_rtl(&mut hpu_hw, &regmap);
 
+        #[cfg(feature = "hw-v80")]
+        let (hpu_version_major, hpu_version_minor) = {
+            let version_reg = regmap
+                              .register()
+                              .get("info::version")
+                              .expect("Unknown register, check regmap definition");
+            let hpu_version_val = hpu_hw.read_reg(*version_reg.offset() as u64);
+            let hpu_version_fields = version_reg.as_field(hpu_version_val);
+            (*hpu_version_fields.get("major").expect("Unknown field"),
+             *hpu_version_fields.get("minor").expect("Unknown field"))
+        };
+
         // In case this is not filled by from_rtl()
         if params.ntt_params.min_pbs_nb.is_none() {
             params.ntt_params.min_pbs_nb = Some(config.firmware.min_batch_size);
@@ -282,6 +298,10 @@ impl HpuBackend {
                 hpu_hw,
                 regmap,
                 params,
+                #[cfg(feature = "hw-v80")]
+                hpu_version_major,
+                #[cfg(feature = "hw-v80")]
+                hpu_version_minor,
                 #[cfg(not(feature = "hw-v80"))]
                 workq_addr,
                 #[cfg(not(feature = "hw-v80"))]
@@ -916,6 +936,16 @@ impl HpuBackend {
         while self.poll_ack_q()? {}
         Ok(())
     }
+
+    #[cfg(feature = "hw-v80")]
+    pub(crate) fn get_hpu_version(&self) -> (u32, u32) {
+        (self.hpu_version_major, self.hpu_version_minor)
+    }
+
+    #[cfg(feature = "hw-v80")]
+    pub(crate) fn map_bar_reg(&mut self) -> Result<(), Box<dyn std::error::Error>> {
+        self.hpu_hw.map_bar_reg()
+    }
 }
 
 impl Drop for HpuBackend {

backends/tfhe-hpu-backend/src/interface/device.rs

@@ -94,6 +94,17 @@ impl HpuDevice {
     ) where
         F: Fn(HpuParameters, &crate::asm::Pbs) -> HpuGlweLookuptableOwned<u64>,
     {
+        // print HPU version
+        {
+            let mut backend = self.backend.lock().unwrap();
+            let (major, minor) = backend.get_hpu_version();
+            tracing::info!("HPU version -> {}.{}", major, minor);
+
+            if major >= 2 && minor >= 3 {
+                backend.map_bar_reg().unwrap();
+            }
+        }
+
         // Properly reset keys
         self.bsk_unset();
         self.ksk_unset();