GSP-RM allocates independent RM sub-heaps for each VF partition inside
the WPR2 region. The default baremetal heap sizing is far too small for
vGPU instance, causing GSP-RM to hit out-of-memory failures during VF
initialization.

The host driver must reserve the correct heap size before GSP boots,
because the WPR2 region is locked down by the hardware after boot and
cannot be resized at runtime. The firmware determines the per-VF carve
from the gspFwHeapVfPartitionCount field in the WPR2 metadata header.

Select a pre-calibrated static heap size based on total_vfs (174 MB for
1 VM, 581 MB for 2-32 VFs, 1370 MB for 48 VFs) and set
vf_partition_count accordingly. Extend FbLayout::new() and
GspBootContext to propagate total_vfs through the boot path.

Signed-off-by: Zhi Wang <zhiw@nvidia.com>
---
 drivers/gpu/nova-core/fb.rs       | 17 +++++++++++++----
 drivers/gpu/nova-core/gsp.rs      |  2 +-
 drivers/gpu/nova-core/gsp/boot.rs | 14 +++++++++++---
 drivers/gpu/nova-core/gsp/fw.rs   | 12 ++++++++++++
 4 files changed, 37 insertions(+), 8 deletions(-)

diff --git a/drivers/gpu/nova-core/fb.rs b/drivers/gpu/nova-core/fb.rs
index 725e428154cf..fb4e6aa9fda4 100644
--- a/drivers/gpu/nova-core/fb.rs
+++ b/drivers/gpu/nova-core/fb.rs
@@ -171,7 +171,13 @@ 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_requested: bool,
+        total_vfs: u16,
+    ) -> Result<Self> {
         let hal = hal::fb_hal(chipset);
 
         let fb = {
@@ -236,8 +242,11 @@ pub(crate) fn new(chipset: Chipset, bar: Bar0<'_>, gsp_fw: &GspFirmware) -> Resu
 
         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_requested {
+                gsp::vgpu_fw_heap_size(u32::from(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 +274,7 @@ pub(crate) fn new(chipset: Chipset, bar: Bar0<'_>, gsp_fw: &GspFirmware) -> Resu
             wpr2_heap,
             wpr2,
             heap,
-            vf_partition_count: 0,
+            vf_partition_count: if vgpu_requested { total_vfs as u8 } else { 0 },
             pmu_reserved_size: hal.pmu_reserved_size(),
         })
     }
diff --git a/drivers/gpu/nova-core/gsp.rs b/drivers/gpu/nova-core/gsp.rs
index 94cd4a784b79..921b92c9eb92 100644
--- a/drivers/gpu/nova-core/gsp.rs
+++ b/drivers/gpu/nova-core/gsp.rs
@@ -27,6 +27,7 @@
 mod sequencer;
 
 pub(crate) use fw::{
+    vgpu_fw_heap_size,
     GspFmcBootParams,
     GspFwWprMeta,
     LibosParams, //
@@ -59,7 +60,6 @@ pub(crate) struct GspBootContext<'a> {
     pub(crate) gsp_falcon: &'a Falcon<GspFalcon>,
     pub(crate) sec2_falcon: &'a Falcon<Sec2Falcon>,
     pub(crate) vgpu_requested: Cell<bool>,
-    #[expect(dead_code)]
     pub(crate) total_vfs: u16,
 }
 
diff --git a/drivers/gpu/nova-core/gsp/boot.rs b/drivers/gpu/nova-core/gsp/boot.rs
index 2981d02d15ad..7c1f3f962fbe 100644
--- a/drivers/gpu/nova-core/gsp/boot.rs
+++ b/drivers/gpu/nova-core/gsp/boot.rs
@@ -111,7 +111,13 @@ pub(crate) fn boot(
             GFP_KERNEL,
         )?;
 
-        let fb_layout = FbLayout::new(ctx.chipset, ctx.bar, &gsp_fw)?;
+        let fb_layout = FbLayout::new(
+            ctx.chipset,
+            ctx.bar,
+            &gsp_fw,
+            ctx.vgpu_requested.get(),
+            ctx.total_vfs,
+        )?;
         dev_dbg!(dev, "{:#x?}\n", fb_layout);
 
         let wpr_meta = Coherent::init(dev, GFP_KERNEL, GspFwWprMeta::new(&gsp_fw, &fb_layout))?;
@@ -138,8 +144,10 @@ pub(crate) fn boot(
 
         self.cmdq
             .send_command_no_wait(ctx.bar, commands::SetSystemInfo::new(ctx.pdev, ctx.chipset))?;
-        self.cmdq
-            .send_command_no_wait(ctx.bar, commands::SetRegistry::new(ctx.vgpu_requested.get())?)?;
+        self.cmdq.send_command_no_wait(
+            ctx.bar,
+            commands::SetRegistry::new(ctx.vgpu_requested.get())?,
+        )?;
 
         hal.post_boot(&self, ctx, &gsp_fw)?;
 
diff --git a/drivers/gpu/nova-core/gsp/fw.rs b/drivers/gpu/nova-core/gsp/fw.rs
index 14424a2c2d83..2f3cbc5d5114 100644
--- a/drivers/gpu/nova-core/gsp/fw.rs
+++ b/drivers/gpu/nova-core/gsp/fw.rs
@@ -101,6 +101,18 @@ 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_1VM: u64 = 174 * u64::SZ_1M;
+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 {
+        1 => GSP_FW_HEAP_SIZE_VGPU_1VM,
+        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