[PATCH 4/5] gpu: nova-core: use the TLV heap sizes

[Timur Tabi <ttabi@nvidia.com>]

Instead of calculating the various heap sizes needed by GSP-RM to boot
from constants and HAL functions, use the heap sizes read from gsp.tlv.

Signed-off-by: Timur Tabi <ttabi@nvidia.com>
---
 drivers/gpu/nova-core/fb.rs            |  31 ++++--
 drivers/gpu/nova-core/fb/hal.rs        |   2 +
 drivers/gpu/nova-core/firmware/gsp.rs  |   7 +-
 drivers/gpu/nova-core/gsp/fw.rs        | 143 ++++---------------------
 drivers/gpu/nova-core/gsp/hal/gh100.rs |   2 +-
 5 files changed, 47 insertions(+), 138 deletions(-)

diff --git a/drivers/gpu/nova-core/fb.rs b/drivers/gpu/nova-core/fb.rs
index 9a3314768385..ca1cb4af9ca5 100644
--- a/drivers/gpu/nova-core/fb.rs
+++ b/drivers/gpu/nova-core/fb.rs
@@ -242,7 +242,8 @@ pub(crate) fn new(
             FbRange(fw_image_addr..fw_image_addr + fw_image_size)
         };
 
-        let (vf_partition_count, wpr2_heap_size) = wpr2_heap_params(chipset, vgpu_state, fb.end)?;
+        let (vf_partition_count, wpr2_heap_size) =
+            wpr2_heap_params(&gsp_fw.heap_params, vgpu_state, fb.end)?;
 
         let wpr2_heap = {
             const WPR2_HEAP_DOWN_ALIGN: Alignment = Alignment::new::<SZ_1M>();
@@ -264,7 +265,7 @@ pub(crate) fn new(
         };
 
         let non_wpr_heap = {
-            let non_wpr_heap_size = hal.non_wpr_heap_size();
+            let non_wpr_heap_size = gsp_fw.heap_params.tlv_non_wpr_heap_size;
             FbRange(wpr2.start - non_wpr_heap_size..wpr2.start)
         };
 
@@ -278,7 +279,7 @@ pub(crate) fn new(
             wpr2,
             non_wpr_heap,
             vf_partition_count,
-            pmu_reserved_size: hal.pmu_reserved_size(),
+            pmu_reserved_size: gsp_fw.heap_params.tlv_pmu_reserved_size,
         })
     }
 }
@@ -298,14 +299,16 @@ pub(crate) fn wpr2_range(bar: Bar0<'_>) -> Option<Range<u64>> {
 }
 
 /// Computes the number of VF partitions and the WPR2 heap size from the vGPU state.
-fn wpr2_heap_params(chipset: Chipset, vgpu_state: VgpuState, fb_size: u64) -> Result<(u8, u64)> {
-    let heap_params = gsp::GspFwHeapParams::from_chipset(chipset);
-
+fn wpr2_heap_params(
+    heap_params: &gsp::GspFwHeapParams,
+    vgpu_state: VgpuState,
+    fb_size: u64,
+) -> Result<(u8, u64)> {
     Ok(match vgpu_state {
         VgpuState::Disabled => (0, heap_params.wpr_heap_size(fb_size)?),
         VgpuState::Enabled { total_vfs } => (
             u8::try_from(total_vfs.get()).map_err(|_| EINVAL)?,
-            heap_params.vgpu_wpr_heap_size(),
+            heap_params.vgpu_wpr_heap_size(total_vfs.get())?,
         ),
     })
 }
@@ -327,16 +330,22 @@ pub(crate) struct FbSizes {
 
 impl FbSizes {
     /// Computes the framebuffer region sizes for GSP-FMC boot.
-    pub(crate) fn new(chipset: Chipset, bar: Bar0<'_>, vgpu_state: VgpuState) -> Result<Self> {
+    pub(crate) fn new(
+        chipset: Chipset,
+        bar: Bar0<'_>,
+        gsp_fw: &GspFirmware<'_>,
+        vgpu_state: VgpuState,
+    ) -> Result<Self> {
         let hal = hal::fb_hal(chipset);
         let fb_size = hal.vidmem_size(bar);
-        let (vf_partition_count, wpr2_heap_size) = wpr2_heap_params(chipset, vgpu_state, fb_size)?;
+        let (vf_partition_count, wpr2_heap_size) =
+            wpr2_heap_params(&gsp_fw.heap_params, vgpu_state, fb_size)?;
 
         Ok(Self {
             frts_size: hal.frts_size(),
             wpr2_heap_size,
-            non_wpr_heap_size: hal.non_wpr_heap_size(),
-            pmu_reserved_size: hal.pmu_reserved_size(),
+            non_wpr_heap_size: gsp_fw.heap_params.tlv_non_wpr_heap_size,
+            pmu_reserved_size: gsp_fw.heap_params.tlv_pmu_reserved_size,
             vf_partition_count,
         })
     }
diff --git a/drivers/gpu/nova-core/fb/hal.rs b/drivers/gpu/nova-core/fb/hal.rs
index 99363cfb116b..e30eb7de2a1c 100644
--- a/drivers/gpu/nova-core/fb/hal.rs
+++ b/drivers/gpu/nova-core/fb/hal.rs
@@ -34,9 +34,11 @@ pub(crate) trait FbHal {
     fn vidmem_size(&self, bar: Bar0<'_>) -> u64;
 
     /// Returns the amount of VRAM to reserve for the PMU.
+    #[expect(unused)]
     fn pmu_reserved_size(&self) -> u32;
 
     /// Returns the non-WPR heap size for this chipset, in bytes.
+    #[expect(unused)]
     fn non_wpr_heap_size(&self) -> u64;
 
     /// Returns the FRTS size, in bytes.
diff --git a/drivers/gpu/nova-core/firmware/gsp.rs b/drivers/gpu/nova-core/firmware/gsp.rs
index a2db7b6ba131..7fd518af9dba 100644
--- a/drivers/gpu/nova-core/firmware/gsp.rs
+++ b/drivers/gpu/nova-core/firmware/gsp.rs
@@ -18,7 +18,8 @@
             Tlv,
         },
     },
-    gpu::Chipset, //
+    gpu::Chipset,
+    gsp::GspFwHeapParams, //
 };
 
 /// The GSP firmware image, its signatures, and the GSP bootloader.
@@ -31,6 +32,8 @@ pub(crate) struct GspFirmware<'a> {
     pub(crate) signatures: Coherent<'a, [u8]>,
     /// GSP bootloader, verifies the GSP firmware before loading and running it.
     pub(crate) bootloader: RiscvFirmware<'a>,
+    /// Heap sizing parameters reported by the firmware metadata.
+    pub(crate) heap_params: GspFwHeapParams,
 }
 
 impl<'a> GspFirmware<'a> {
@@ -48,6 +51,7 @@ pub(crate) fn new(
             let fw_vvec = tlv.load_file(dev, chipset)?;
 
             let signatures = Coherent::from_slice(dev, tlv.get_bytes(b"SIGN")?, GFP_KERNEL)?;
+            let heap_params = GspFwHeapParams::from_tlv(&tlv)?;
 
             Ok(try_pin_init!(Self {
                 radix3 <- Radix3::new(dev, fw_vvec),
@@ -57,6 +61,7 @@ pub(crate) fn new(
 
                     RiscvFirmware::new(dev, &bl)?
                 },
+                heap_params,
             }))
         })
     }
diff --git a/drivers/gpu/nova-core/gsp/fw.rs b/drivers/gpu/nova-core/gsp/fw.rs
index f0edae9656d8..8873282ac4cd 100644
--- a/drivers/gpu/nova-core/gsp/fw.rs
+++ b/drivers/gpu/nova-core/gsp/fw.rs
@@ -7,8 +7,6 @@
 // Alias to avoid repeating the version number with every use.
 use r000_00 as bindings;
 
-use core::ops::Range;
-
 use kernel::{
     bitfield,
     dma::Coherent,
@@ -40,10 +38,6 @@
         gsp::GspFirmware,
         tlv::Tlv, //
     },
-    gpu::{
-        Architecture,
-        Chipset, //
-    },
     gsp::{
         cmdq::Cmdq,
         GSP_PAGE_SHIFT,
@@ -68,110 +62,23 @@
 
 /// Heap memory requirements and constraints for booting GSP-RM.
 pub(crate) struct GspFwHeapParams {
-    /// The amount of heap memory used during GSP-RM boot and initialization (up to and including
-    /// the first client subdevice allocation), in bytes.
-    base_rm_size: u64,
-    /// The base amount of heap required by the GSP operating system, in bytes.
-    carveout_size: u64,
-    /// The minimum and maximum sizes allowed for the GSP FW heap, in bytes.
-    allowed_heap_size: Range<u64>,
-
-    #[expect(unused)]
     tlv_gsp_fw_heap_base_size: u64,
-    #[expect(unused)]
     tlv_gsp_hw_heap_size_per_gb: u64,
-    #[expect(unused)]
     tlv_gsp_fw_heap_base_size_vgpu: u64,
-    #[expect(unused)]
     tlv_gsp_fw_heap_size_per_vm: u64,
-    #[expect(unused)]
-    tlv_pmu_reserved_size: u32,
-    #[expect(unused)]
-    tlv_non_wpr_heap_size: u64,
+    pub(crate) tlv_pmu_reserved_size: u32,
+    pub(crate) tlv_non_wpr_heap_size: u64,
 }
 
 /// Minimum required alignment for the GSP heap.
 const GSP_HEAP_ALIGNMENT: Alignment = Alignment::new::<{ 1 << 20 }>();
 
 impl GspFwHeapParams {
-    /// Version 2 of the GSP LIBOS (Turing and GA100)
-    fn libos2() -> Self {
-        Self {
-            base_rm_size: cv!(bindings::GSP_FW_HEAP_PARAM_BASE_RM_SIZE_TU10X),
-            carveout_size: cv!(bindings::GSP_FW_HEAP_PARAM_OS_SIZE_LIBOS2),
-            allowed_heap_size: cv!(bindings::GSP_FW_HEAP_SIZE_OVERRIDE_LIBOS2_MIN_MB => u64)
-                * u64::SZ_1M
-                ..cv!(bindings::GSP_FW_HEAP_SIZE_OVERRIDE_LIBOS2_MAX_MB => u64) * u64::SZ_1M,
-
-            tlv_gsp_fw_heap_base_size: 0,
-            tlv_gsp_hw_heap_size_per_gb: 0,
-            tlv_gsp_fw_heap_base_size_vgpu: 0,
-            tlv_gsp_fw_heap_size_per_vm: 0,
-            tlv_pmu_reserved_size: 0,
-            tlv_non_wpr_heap_size: 0,
-        }
-    }
-
-    /// Version 3 of the GSP LIBOS (GA102 - AD10x)
-    fn libos3() -> Self {
-        Self {
-            base_rm_size: cv!(bindings::GSP_FW_HEAP_PARAM_BASE_RM_SIZE_TU10X),
-            carveout_size: cv!(bindings::GSP_FW_HEAP_PARAM_OS_SIZE_LIBOS3_BAREMETAL),
-            allowed_heap_size: cv!(
-                bindings::GSP_FW_HEAP_SIZE_OVERRIDE_LIBOS3_BAREMETAL_MIN_MB => u64)
-                * u64::SZ_1M
-                ..cv!(bindings::GSP_FW_HEAP_SIZE_OVERRIDE_LIBOS3_BAREMETAL_MAX_MB => u64)
-                    * u64::SZ_1M,
-
-            tlv_gsp_fw_heap_base_size: 0,
-            tlv_gsp_hw_heap_size_per_gb: 0,
-            tlv_gsp_fw_heap_base_size_vgpu: 0,
-            tlv_gsp_fw_heap_size_per_vm: 0,
-            tlv_pmu_reserved_size: 0,
-            tlv_non_wpr_heap_size: 0,
-        }
-    }
-
-    /// Version 3 of the GSP LIBOS, (GH100+)
-    fn libos3_gh100() -> Self {
-        let base = Self::libos3();
-
-        Self {
-            base_rm_size: cv!(bindings::GSP_FW_HEAP_PARAM_BASE_RM_SIZE_GH100),
-            carveout_size: base.carveout_size,
-            allowed_heap_size: base.allowed_heap_size,
-
-            tlv_gsp_fw_heap_base_size: 0,
-            tlv_gsp_hw_heap_size_per_gb: 0,
-            tlv_gsp_fw_heap_base_size_vgpu: 0,
-            tlv_gsp_fw_heap_size_per_vm: 0,
-            tlv_pmu_reserved_size: 0,
-            tlv_non_wpr_heap_size: 0,
-        }
-    }
-
     /// Returns the heap parameters corresponding to `chipset`.
     ///
     /// The parameters reported by GSP-RM are initialized to zero.
-    pub(crate) fn from_chipset(chipset: Chipset) -> Self {
-        match chipset.arch() {
-            Architecture::Turing => Self::libos2(),
-            Architecture::Ampere if chipset < Chipset::GA102 => Self::libos2(),
-            Architecture::Ampere | Architecture::Ada => Self::libos3(),
-            _ => Self::libos3_gh100(),
-        }
-    }
-
-    /// Returns the heap parameters from the TLV
-    #[expect(unused)]
     pub(crate) fn from_tlv(tlv: &Tlv<'_>) -> Result<Self> {
         Ok(Self {
-            // These three fields are not used for TLV-based implementations,
-            // so assign dummy values to them for now.
-            base_rm_size: 0,
-            carveout_size: 0,
-            allowed_heap_size: u64::MIN..u64::MAX,
-
             tlv_gsp_fw_heap_base_size: tlv.get_u64(b"HBSZ")?,
             tlv_gsp_hw_heap_size_per_gb: tlv.get_u64(b"HSPG")?,
             tlv_gsp_fw_heap_base_size_vgpu: tlv.get_u64(b"HBSV")?,
@@ -182,44 +89,30 @@ pub(crate) fn from_tlv(tlv: &Tlv<'_>) -> Result<Self> {
     }
 
     /// Returns the WPR heap size to reserve when vGPU is enabled.
-    pub(crate) fn vgpu_wpr_heap_size(&self) -> u64 {
-        u64::from(bindings::GSP_FW_HEAP_SIZE_VGPU_DEFAULT)
-    }
-
-    /// Returns the amount of heap memory required to support a single channel allocation.
-    fn client_alloc_size() -> u64 {
-        u64::from(bindings::GSP_FW_HEAP_PARAM_CLIENT_ALLOC_SIZE)
+    pub(crate) fn vgpu_wpr_heap_size(&self, total_vfs: u16) -> Result<u64> {
+        //u64::from(bindings::GSP_FW_HEAP_SIZE_VGPU_DEFAULT)
+        let vm_size = self
+            .tlv_gsp_fw_heap_size_per_vm
+            .saturating_mul(u64::from(total_vfs))
             .align_up(GSP_HEAP_ALIGNMENT)
-            .unwrap_or(u64::MAX)
-    }
+            .ok_or(EINVAL)?;
 
-    /// Returns the amount of memory to reserve for management purposes for a framebuffer of size
-    /// `fb_size`.
-    fn management_overhead(fb_size: u64) -> Result<u64> {
-        let fb_size_gb = fb_size.div_ceil(u64::SZ_1G);
-
-        u64::from(bindings::GSP_FW_HEAP_PARAM_SIZE_PER_GB)
-            .checked_mul(fb_size_gb)
-            .ok_or(EINVAL)?
-            .align_up(GSP_HEAP_ALIGNMENT)
-            .ok_or(EINVAL)
+        Ok(self.tlv_gsp_fw_heap_base_size_vgpu.saturating_add(vm_size))
     }
 
     /// Returns the amount of memory (in bytes) to allocate for the WPR heap for a framebuffer size
     /// of `fb_size` (in bytes).
     pub(crate) fn wpr_heap_size(&self, fb_size: u64) -> Result<u64> {
-        // The WPR heap will contain the following:
-        // LIBOS carveout,
+        let fb_size_gb = fb_size.div_ceil(u64::SZ_1G);
+        let management_overhead: u64 = self
+            .tlv_gsp_hw_heap_size_per_gb
+            .saturating_mul(fb_size_gb)
+            .align_up(GSP_HEAP_ALIGNMENT)
+            .ok_or(EINVAL)?;
+
         Ok(self
-            .carveout_size
-            // RM boot working memory,
-            .saturating_add(self.base_rm_size)
-            // One RM client,
-            .saturating_add(Self::client_alloc_size())
-            // Overhead for memory management.
-            .saturating_add(Self::management_overhead(fb_size)?)
-            // Clamp to the supported heap sizes.
-            .clamp(self.allowed_heap_size.start, self.allowed_heap_size.end - 1))
+            .tlv_gsp_fw_heap_base_size
+            .saturating_add(management_overhead))
     }
 }
 
diff --git a/drivers/gpu/nova-core/gsp/hal/gh100.rs b/drivers/gpu/nova-core/gsp/hal/gh100.rs
index 91201b51030e..f43fbe3dd798 100644
--- a/drivers/gpu/nova-core/gsp/hal/gh100.rs
+++ b/drivers/gpu/nova-core/gsp/hal/gh100.rs
@@ -151,7 +151,7 @@ fn boot<'gpu>(
         let chipset = ctx.chipset;
         let gsp_falcon = ctx.gsp_falcon;
 
-        let fb_sizes = FbSizes::new(chipset, ctx.bar, ctx.vgpu.state())?;
+        let fb_sizes = FbSizes::new(chipset, ctx.bar, gsp_fw, ctx.vgpu.state())?;
         dev_dbg!(dev, "{:#x?}\n", fb_sizes);
 
         let wpr_meta =
-- 
2.54.0