[PATCH v2 7/7] gpu: nova-core: reserve larger WPR2 heap for vGPU

[Zhi Wang]

From: Zhi Wang <zhi.wang.ai@xxxxxxxx>

Reserve a larger GSP WPR2 heap when vGPU is enabled, using the VF count
reported by VgpuManager to select the heap size and to populate the VF
partition count in the GSP WPR metadata.

Use the baremetal heap sizing path when vGPU is disabled. For vGPU,
reserve 581 MiB for up to 32 VFs and 1370 MiB for larger VF counts.

Signed-off-by: Zhi Wang <zhiw@xxxxxxxxxx>
---
 drivers/gpu/nova-core/fb.rs       | 25 ++++++++++++++++++++-----
 drivers/gpu/nova-core/gsp.rs      |  1 +
 drivers/gpu/nova-core/gsp/boot.rs |  2 +-
 drivers/gpu/nova-core/gsp/fw.rs   | 10 ++++++++++
 4 files changed, 32 insertions(+), 6 deletions(-)

diff --git a/drivers/gpu/nova-core/fb.rs b/drivers/gpu/nova-core/fb.rs
index 725e428154cf..b5b1f6c13edc 100644
--- a/drivers/gpu/nova-core/fb.rs
+++ b/drivers/gpu/nova-core/fb.rs
@@ -24,7 +24,8 @@
     gpu::Chipset,
     gsp,
     num::FromSafeCast,
-    regs, //
+    regs,
+    vgpu::VgpuManager, //
 };
 
 mod hal;
@@ -171,7 +172,12 @@ pub(crate) struct FbLayout {
 
 impl FbLayout {
     /// Computes the FB layout for `chipset` required to run the `gsp_fw` GSP firmware.
-    pub(crate) fn new(chipset: Chipset, bar: Bar0<'_>, gsp_fw: &GspFirmware) -> Result<Self> {
+    pub(crate) fn new(
+        chipset: Chipset,
+        bar: Bar0<'_>,
+        gsp_fw: &GspFirmware,
+        vgpu: &VgpuManager,
+    ) -> Result<Self> {
         let hal = hal::fb_hal(chipset);
 
         let fb = {
@@ -234,10 +240,19 @@ pub(crate) fn new(chipset: Chipset, bar: Bar0<'_>, gsp_fw: &GspFirmware) -> Resu
             FbRange(elf_addr..elf_addr + elf_size)
         };
 
+        let vf_partition_count = if vgpu.enabled() {
+            vgpu.total_vfs().try_into().map_err(|_| EINVAL)?
+        } else {
+            0
+        };
+
         let wpr2_heap = {
             const WPR2_HEAP_DOWN_ALIGN: Alignment = Alignment::new::<SZ_1M>();
-            let wpr2_heap_size =
-                gsp::LibosParams::from_chipset(chipset).wpr_heap_size(chipset, fb.end)?;
+            let wpr2_heap_size = if vgpu.enabled() {
+                gsp::vgpu_fw_heap_size(u32::from(vgpu.total_vfs()))
+            } else {
+                gsp::LibosParams::from_chipset(chipset).wpr_heap_size(chipset, fb.end)?
+            };
             let wpr2_heap_addr = (elf.start - wpr2_heap_size).align_down(WPR2_HEAP_DOWN_ALIGN);
 
             FbRange(wpr2_heap_addr..(elf.start).align_down(WPR2_HEAP_DOWN_ALIGN))
@@ -265,7 +280,7 @@ pub(crate) fn new(chipset: Chipset, bar: Bar0<'_>, gsp_fw: &GspFirmware) -> Resu
             wpr2_heap,
             wpr2,
             heap,
-            vf_partition_count: 0,
+            vf_partition_count,
             pmu_reserved_size: hal.pmu_reserved_size(),
         })
     }
diff --git a/drivers/gpu/nova-core/gsp.rs b/drivers/gpu/nova-core/gsp.rs
index 6821008d48d9..a159f36bf704 100644
--- a/drivers/gpu/nova-core/gsp.rs
+++ b/drivers/gpu/nova-core/gsp.rs
@@ -26,6 +26,7 @@
 mod sequencer;
 
 pub(crate) use fw::{
+    vgpu_fw_heap_size,
     GspFmcBootParams,
     GspFwWprMeta,
     LibosParams, //
diff --git a/drivers/gpu/nova-core/gsp/boot.rs b/drivers/gpu/nova-core/gsp/boot.rs
index c607081e8242..dea028f4b434 100644
--- a/drivers/gpu/nova-core/gsp/boot.rs
+++ b/drivers/gpu/nova-core/gsp/boot.rs
@@ -56,7 +56,7 @@ pub(crate) fn boot(
 
         let gsp_fw = KBox::pin_init(GspFirmware::new(dev, chipset, FIRMWARE_VERSION), GFP_KERNEL)?;
 
-        let fb_layout = FbLayout::new(chipset, bar, &gsp_fw)?;
+        let fb_layout = FbLayout::new(chipset, bar, &gsp_fw, ctx.vgpu)?;
         dev_dbg!(dev, "{:#x?}\n", fb_layout);
 
         let wpr_meta = Coherent::init(dev, GFP_KERNEL, GspFwWprMeta::new(&gsp_fw, &fb_layout))?;
diff --git a/drivers/gpu/nova-core/gsp/fw.rs b/drivers/gpu/nova-core/gsp/fw.rs
index 4db0cfa4dc4d..2fb72f8dc3c9 100644
--- a/drivers/gpu/nova-core/gsp/fw.rs
+++ b/drivers/gpu/nova-core/gsp/fw.rs
@@ -101,6 +101,16 @@ pub(in crate::gsp) fn advance_cpu_write_ptr(qs: &Coherent<GspMem>, count: u32) {
 pub(crate) const GSP_MSG_QUEUE_ELEMENT_SIZE_MAX: usize =
     num::u32_as_usize(bindings::GSP_MSG_QUEUE_ELEMENT_SIZE_MAX);
 
+const GSP_FW_HEAP_SIZE_VGPU_DEFAULT: u64 = 581 * u64::SZ_1M;
+const GSP_FW_HEAP_SIZE_VGPU_48VMS: u64 = 1370 * u64::SZ_1M;
+
+pub(crate) fn vgpu_fw_heap_size(total_vfs: u32) -> u64 {
+    match total_vfs {
+        2..=32 => GSP_FW_HEAP_SIZE_VGPU_DEFAULT,
+        _ => GSP_FW_HEAP_SIZE_VGPU_48VMS,
+    }
+}
+
 /// Empty type to group methods related to heap parameters for running the GSP firmware.
 enum GspFwHeapParams {}
 
-- 
2.51.0