From 04cfbf21a447b3b13a8917f07ac0479ac5fe9998 Mon Sep 17 00:00:00 2001
From: Antoine Viallon <antoine@lesviallon.fr>
Date: Sun, 16 Apr 2023 22:35:50 +0200
Subject: [PATCH] [Hardware/Nvidia] Fix suspend with Nvidia GPUs

---
 hardware/nvidia/default.nix     |  2 +-
 hardware/nvidia/proprietary.nix | 14 ++++++++++----
 2 files changed, 11 insertions(+), 5 deletions(-)

diff --git a/hardware/nvidia/default.nix b/hardware/nvidia/default.nix
index 83831e1..4a75ef2 100644
--- a/hardware/nvidia/default.nix
+++ b/hardware/nvidia/default.nix
@@ -8,12 +8,12 @@ in {
     ./vgpu.nix
     ./proprietary.nix
     ./opensource.nix
+    ( mkRenamedOptionModule [ "aviallon" "hardware" "nvidia" "saveAllVram" ] [ "aviallon" "hardware" "nvidia" "proprietary" "saveAllVram" ] )
   ];
 
   options.aviallon.hardware.nvidia = {
     enable = mkEnableOption "enable Nvidia hardware config";
     useProprietary = mkEnableOption "nvidia proprietary drivers";
-    saveAllVram = mkEnableOption "back up all VRAM in /tmp before going to sleep. May reduce artifacts after resuming";
   };
 
   config = mkIf cfg.enable {
diff --git a/hardware/nvidia/proprietary.nix b/hardware/nvidia/proprietary.nix
index 8c0726f..f67ab72 100644
--- a/hardware/nvidia/proprietary.nix
+++ b/hardware/nvidia/proprietary.nix
@@ -21,10 +21,14 @@ in {
         example = 28;
         type = types.int;
       };
+      saveAllVram = mkEnableOption "back up all VRAM in /var/tmp before going to sleep. May reduce artifacts after resuming";
     };
   };
   
   config = mkIf (cfg.enable && cfg.useProprietary) {
+
+    assertions = [];
+  
     boot.initrd.kernelModules = [
       "nvidia"
       "nvidia_drm"
@@ -44,8 +48,10 @@ in {
     services.xserver.exportConfiguration = true;
 
     hardware.nvidia = {
-      powerManagement.enable = true;
-      powerManagement.finegrained = mkIf config.hardware.nvidia.prime.offload.enable true;
+      powerManagement = mkIf (config.hardware.nvidia.prime.offload.enable || cfg.proprietary.saveAllVram) {
+        enable = true;
+        finegrained = true;
+      };
       modesetting.enable = true;
       nvidiaSettings = true;
     };
@@ -55,8 +61,8 @@ in {
         "nvidia-drm.modeset" = 1;
         "nvidia.NVreg_UsePageAttributeTable" = 1;
       }
-      // optionalAttrs cfg.saveAllVram {
-        "nvidia.NVreg_PreserveVideoMemoryAllocations" = 1;
+      // {
+        "nvidia.NVreg_PreserveVideoMemoryAllocations" = toValue cfg.proprietary.saveAllVram;
         "nvidia.NVreg_TemporaryFilePath" = "/var/tmp/nvidia-gpu.vram.img";
       }
     ;