Add the ability to flash NixOS to eMMC on AGX Orin

default.nix

@@ -89,7 +89,7 @@ in rec {
   # GST plugins
 
   # Generate a flash script using the built configuration options set in a NixOS configuration
-  flashScriptFromNixos = config: let
+  flashScriptFromNixos = flashEmmc: config: let
     cfg = config.hardware.nvidia-jetpack;
   in callPackage ./flash-script.nix {
     name = config.networking.hostName;
@@ -108,6 +108,10 @@ in rec {
     };
 
     dtbsDir = config.hardware.deviceTree.package;
+
+    sdImage = if (flashEmmc)
+      then (pkgsAarch64.nixos (import ./emmc-image.nix { inherit jetson-firmware cfg; })).config.system.build.sdImage
+      else null;
   };
 
   flash-scripts = rec {
@@ -121,7 +125,7 @@ in rec {
         hardware.nvidia-jetpack.enable = true;
       }).config.hardware.deviceTree.package;
     };
-  } // (lib.mapAttrs' (n: c: lib.nameValuePair "flash-${n}" (flashScriptFromNixos (pkgs.nixos {
+  } // (lib.mapAttrs' (n: c: lib.nameValuePair "flash-${n}" (flashScriptFromNixos false (pkgs.nixos {
     imports = [ ./modules/default.nix { hardware.nvidia-jetpack = c; } ];
     hardware.nvidia-jetpack.enable = true;
     networking.hostName = n; # Just so it sets the flash binary name.
@@ -130,6 +134,27 @@ in rec {
     "xavier-agx-devkit" = { som = "xavier-agx"; carrierBoard = "devkit"; };
     "xavier-nx-devkit" = { som = "xavier-nx"; carrierBoard = "devkit"; };
     "xavier-nx-devkit-emmc" = { som = "xavier-nx-emmc"; carrierBoard = "devkit"; };
+  })
+  //({
+    "flash-orin-agx-devkit-emmc" = let
+      partitionsToRemove = [
+        "RECNAME" "RECDTB-NAME" "RP1" "RP2" "RECROOTFS" # Recovery
+      ];
+      xpathMatch = lib.concatMapStringsSep " or " (p: "@name = \"${p}\"") partitionsToRemove;
+      filterPartitions = basefile: pkgs.runCommand "flash.xml" { nativeBuildInputs = [ pkgs.buildPackages.xmlstarlet ]; } ''
+      xmlstarlet ed -d '//partition[${xpathMatch}]' ${basefile} >$out
+    '';
+    in
+    (flashScriptFromNixos true (pkgs.nixos {
+      imports = [ ./modules/default.nix { hardware.nvidia-jetpack = {
+        som = "orin-agx";
+        carrierBoard = "devkit";
+        dtbName = "tegra234-p3701-0000-p3737-0000.dtb";
+        flashScriptOverrides.partitionTemplate = (filterPartitions "${bspSrc}/bootloader/t186ref/cfg/flash_t234_qspi_sdmmc.xml");
+      }; } ];
+      hardware.nvidia-jetpack.enable = true;
+      networking.hostName = "orin-agx-devkit";
+    }).config);
   });
 }
 // l4t

emmc-image.nix

@@ -0,0 +1,54 @@
+{ jetson-firmware, cfg }: { config, pkgs, modulesPath, lib, ... }:
+
+with lib;
+
+{
+  imports = [
+    ./modules/default.nix
+    (modulesPath + "/profiles/base.nix")
+    (modulesPath + "/profiles/installation-device.nix")
+    (modulesPath + "/installer/sd-card/sd-image.nix")
+  ];
+
+  hardware.nvidia-jetpack = cfg;
+
+  # Avoids a bunch of extra modules we don't have in the tegra_defconfig, like "ata_piix",
+  disabledModules = [ (modulesPath + "/profiles/all-hardware.nix") ];
+
+  boot.loader.grub.enable = false;
+
+  sdImage = let
+    kernelPath = "${config.boot.kernelPackages.kernel}/" + "${config.system.boot.loader.kernelFile}";
+    initrdPath = "${config.system.build.initialRamdisk}/" + "${config.system.boot.loader.initrdFile}";
+    fdtPath = "${config.hardware.deviceTree.package}/" + "${config.hardware.nvidia-jetpack.dtbName}";
+    extlinux = pkgs.writeText "extlinux.conf" ''
+      TIMEOUT 30
+      DEFAULT primary
+
+      MENU TITLE NixOS boot options
+
+      LABEL primary
+        MENU LABEL primary kernel
+        LINUX /boot/${config.system.boot.loader.kernelFile}
+        FDT /boot/${config.hardware.nvidia-jetpack.dtbName}
+        INITRD /boot/${config.system.boot.loader.initrdFile}
+        APPEND init=${config.system.build.toplevel}/init ${toString config.boot.kernelParams}
+    '';
+  in {
+    populateFirmwareCommands = ''
+        mkdir -p firmware/EFI/BOOT
+        cp ${jetson-firmware}/L4TLauncher.efi firmware/EFI/BOOT/BOOTAA64.efi
+    '';
+    postBuildCommands = ''
+        cp firmware_part.img $out
+        cp root-fs.img $out
+    '';
+    populateRootCommands = ''
+        mkdir -p ./files/boot/extlinux
+        cp ${extlinux} ./files/boot/extlinux/extlinux.conf
+        cp ${kernelPath} "./files/boot/${config.system.boot.loader.kernelFile}"
+        cp ${initrdPath} "./files/boot/${config.system.boot.loader.initrdFile}"
+        cp ${fdtPath} "./files/boot/${config.hardware.nvidia-jetpack.dtbName}"
+    '';
+  };
+}

flash-script.nix

@@ -8,6 +8,9 @@
 
   # Optional package containing uefi_jetson.efi to replace prebuilt version
   jetson-firmware ? null,
+
+  # Optional package containing sdcard image
+  sdImage ? null,
 }:
 
 writeShellScriptBin "flash-${name}" (''
@@ -42,10 +45,16 @@ writeShellScriptBin "flash-${name}" (''
   cp ${jetson-firmware}/dtbs/*.dtbo kernel/dtb/
   ''}
 
+  ${lib.optionalString (sdImage != null) ''
+    # Copy rootfs and boot partitions
+    cp -r ${sdImage}/firmware_part.img bootloader/esp.img
+    cp -r ${sdImage}/root-fs.img bootloader/system.img
+  ''}
+
   chmod -R u+w .
 
 '' + (if (flashArgs != null) then ''
-  ./flash.sh ${lib.optionalString (partitionTemplate != null) "-c flash.xml"} $@ ${flashArgs}
+  ./flash.sh ${lib.optionalString (sdImage != null) "-r"} ${lib.optionalString (partitionTemplate != null) "-c flash.xml"} $@ ${flashArgs}
 '' else ''
   ${runtimeShell}
 ''))

modules/default.nix

@@ -61,6 +61,12 @@ in
         type = types.bool;
         description = "Enable PREEMPT_RT patches";
       };
+
+      dtbName = mkOption {
+        default = "";
+        type = types.str;
+        description = "Devicetree filename";
+      };
     };
   };