Merge "tests: Stablize some benchmarks" into main
diff --git a/docs/custom_vm.md b/docs/custom_vm.md
index d52aa95..840acc3 100644
--- a/docs/custom_vm.md
+++ b/docs/custom_vm.md
@@ -190,6 +190,10 @@
             "writable": true
         }
     ],
+    "gpu": {
+        "backend": "virglrenderer",
+        "context_types": ["virgl2"]
+    },
     "params": "root=/dev/vda3 rootwait noinitrd ro enforcing=0 cros_debug cros_secure",
     "protected": false,
     "cpu_topology": "match_host",
diff --git a/java/framework/src/android/system/virtualmachine/VirtualMachineConfig.java b/java/framework/src/android/system/virtualmachine/VirtualMachineConfig.java
index 1e0f6c6..8b444fc 100644
--- a/java/framework/src/android/system/virtualmachine/VirtualMachineConfig.java
+++ b/java/framework/src/android/system/virtualmachine/VirtualMachineConfig.java
@@ -711,6 +711,10 @@
                 Optional.ofNullable(customImageConfig.getDisplayConfig())
                         .map(dc -> dc.toParcelable())
                         .orElse(null);
+        config.gpuConfig =
+                Optional.ofNullable(customImageConfig.getGpuConfig())
+                        .map(dc -> dc.toParcelable())
+                        .orElse(null);
         config.protectedVm = this.mProtectedVm;
         config.memoryMib = bytesToMebiBytes(mMemoryBytes);
         config.cpuTopology = (byte) this.mCpuTopology;
diff --git a/java/framework/src/android/system/virtualmachine/VirtualMachineCustomImageConfig.java b/java/framework/src/android/system/virtualmachine/VirtualMachineCustomImageConfig.java
index 2fcad20..c0ff11a 100644
--- a/java/framework/src/android/system/virtualmachine/VirtualMachineCustomImageConfig.java
+++ b/java/framework/src/android/system/virtualmachine/VirtualMachineCustomImageConfig.java
@@ -36,6 +36,7 @@
     private static final String KEY_TOUCH = "touch";
     private static final String KEY_KEYBOARD = "keyboard";
     private static final String KEY_MOUSE = "mouse";
+    private static final String KEY_GPU = "gpu";
 
     @Nullable private final String name;
     @Nullable private final String kernelPath;
@@ -47,6 +48,7 @@
     private final boolean touch;
     private final boolean keyboard;
     private final boolean mouse;
+    @Nullable private final GpuConfig gpuConfig;
 
     @Nullable
     public Disk[] getDisks() {
@@ -101,7 +103,8 @@
             DisplayConfig displayConfig,
             boolean touch,
             boolean keyboard,
-            boolean mouse) {
+            boolean mouse,
+            GpuConfig gpuConfig) {
         this.name = name;
         this.kernelPath = kernelPath;
         this.initrdPath = initrdPath;
@@ -112,6 +115,7 @@
         this.touch = touch;
         this.keyboard = keyboard;
         this.mouse = mouse;
+        this.gpuConfig = gpuConfig;
     }
 
     static VirtualMachineCustomImageConfig from(PersistableBundle customImageConfigBundle) {
@@ -142,6 +146,7 @@
         builder.useTouch(customImageConfigBundle.getBoolean(KEY_TOUCH));
         builder.useKeyboard(customImageConfigBundle.getBoolean(KEY_KEYBOARD));
         builder.useMouse(customImageConfigBundle.getBoolean(KEY_MOUSE));
+        builder.setGpuConfig(GpuConfig.from(customImageConfigBundle.getPersistableBundle(KEY_GPU)));
         return builder.build();
     }
 
@@ -173,6 +178,9 @@
         pb.putBoolean(KEY_TOUCH, touch);
         pb.putBoolean(KEY_KEYBOARD, keyboard);
         pb.putBoolean(KEY_MOUSE, mouse);
+        pb.putPersistableBundle(
+                KEY_GPU,
+                Optional.ofNullable(gpuConfig).map(gc -> gc.toPersistableBundle()).orElse(null));
         return pb;
     }
 
@@ -181,6 +189,11 @@
         return displayConfig;
     }
 
+    @Nullable
+    public GpuConfig getGpuConfig() {
+        return gpuConfig;
+    }
+
     /** @hide */
     public static final class Disk {
         private final boolean writable;
@@ -224,6 +237,7 @@
         private boolean touch;
         private boolean keyboard;
         private boolean mouse;
+        private GpuConfig gpuConfig;
 
         /** @hide */
         public Builder() {}
@@ -271,6 +285,12 @@
         }
 
         /** @hide */
+        public Builder setGpuConfig(GpuConfig gpuConfig) {
+            this.gpuConfig = gpuConfig;
+            return this;
+        }
+
+        /** @hide */
         public Builder useTouch(boolean touch) {
             this.touch = touch;
             return this;
@@ -300,7 +320,8 @@
                     displayConfig,
                     touch,
                     keyboard,
-                    mouse);
+                    mouse,
+                    gpuConfig);
         }
     }
 
@@ -437,4 +458,223 @@
             }
         }
     }
+
+    /** @hide */
+    public static final class GpuConfig {
+        private static final String KEY_BACKEND = "backend";
+        private static final String KEY_CONTEXT_TYPES = "context_types";
+        private static final String KEY_PCI_ADDRESS = "pci_address";
+        private static final String KEY_RENDERER_FEATURES = "renderer_features";
+        private static final String KEY_RENDERER_USE_EGL = "renderer_use_egl";
+        private static final String KEY_RENDERER_USE_GLES = "renderer_use_gles";
+        private static final String KEY_RENDERER_USE_GLX = "renderer_use_glx";
+        private static final String KEY_RENDERER_USE_SURFACELESS = "renderer_use_surfaceless";
+        private static final String KEY_RENDERER_USE_VULKAN = "renderer_use_vulkan";
+
+        private final String backend;
+        private final String[] contextTypes;
+        private final String pciAddress;
+        private final String rendererFeatures;
+        private final boolean rendererUseEgl;
+        private final boolean rendererUseGles;
+        private final boolean rendererUseGlx;
+        private final boolean rendererUseSurfaceless;
+        private final boolean rendererUseVulkan;
+
+        private GpuConfig(
+                String backend,
+                String[] contextTypes,
+                String pciAddress,
+                String rendererFeatures,
+                boolean rendererUseEgl,
+                boolean rendererUseGles,
+                boolean rendererUseGlx,
+                boolean rendererUseSurfaceless,
+                boolean rendererUseVulkan) {
+            this.backend = backend;
+            this.contextTypes = contextTypes;
+            this.pciAddress = pciAddress;
+            this.rendererFeatures = rendererFeatures;
+            this.rendererUseEgl = rendererUseEgl;
+            this.rendererUseGles = rendererUseGles;
+            this.rendererUseGlx = rendererUseGlx;
+            this.rendererUseSurfaceless = rendererUseSurfaceless;
+            this.rendererUseVulkan = rendererUseVulkan;
+        }
+
+        /** @hide */
+        public String getBackend() {
+            return backend;
+        }
+
+        /** @hide */
+        public String[] getContextTypes() {
+            return contextTypes;
+        }
+
+        /** @hide */
+        public String getPciAddress() {
+            return pciAddress;
+        }
+
+        /** @hide */
+        public String getRendererFeatures() {
+            return rendererFeatures;
+        }
+
+        /** @hide */
+        public boolean getRendererUseEgl() {
+            return rendererUseEgl;
+        }
+
+        /** @hide */
+        public boolean getRendererUseGles() {
+            return rendererUseGles;
+        }
+
+        /** @hide */
+        public boolean getRendererUseGlx() {
+            return rendererUseGlx;
+        }
+
+        /** @hide */
+        public boolean getRendererUseSurfaceless() {
+            return rendererUseSurfaceless;
+        }
+
+        /** @hide */
+        public boolean getRendererUseVulkan() {
+            return rendererUseVulkan;
+        }
+
+        android.system.virtualizationservice.GpuConfig toParcelable() {
+            android.system.virtualizationservice.GpuConfig parcelable =
+                    new android.system.virtualizationservice.GpuConfig();
+            parcelable.backend = this.backend;
+            parcelable.contextTypes = this.contextTypes;
+            parcelable.pciAddress = this.pciAddress;
+            parcelable.rendererFeatures = this.rendererFeatures;
+            parcelable.rendererUseEgl = this.rendererUseEgl;
+            parcelable.rendererUseGles = this.rendererUseGles;
+            parcelable.rendererUseGlx = this.rendererUseGlx;
+            parcelable.rendererUseSurfaceless = this.rendererUseSurfaceless;
+            parcelable.rendererUseVulkan = this.rendererUseVulkan;
+            return parcelable;
+        }
+
+        private static GpuConfig from(PersistableBundle pb) {
+            if (pb == null) {
+                return null;
+            }
+            Builder builder = new Builder();
+            builder.setBackend(pb.getString(KEY_BACKEND));
+            builder.setContextTypes(pb.getStringArray(KEY_CONTEXT_TYPES));
+            builder.setPciAddress(pb.getString(KEY_PCI_ADDRESS));
+            builder.setRendererFeatures(pb.getString(KEY_RENDERER_FEATURES));
+            builder.setRendererUseEgl(pb.getBoolean(KEY_RENDERER_USE_EGL));
+            builder.setRendererUseGles(pb.getBoolean(KEY_RENDERER_USE_GLES));
+            builder.setRendererUseGlx(pb.getBoolean(KEY_RENDERER_USE_GLX));
+            builder.setRendererUseSurfaceless(pb.getBoolean(KEY_RENDERER_USE_SURFACELESS));
+            builder.setRendererUseVulkan(pb.getBoolean(KEY_RENDERER_USE_VULKAN));
+            return builder.build();
+        }
+
+        private PersistableBundle toPersistableBundle() {
+            PersistableBundle pb = new PersistableBundle();
+            pb.putString(KEY_BACKEND, this.backend);
+            pb.putStringArray(KEY_CONTEXT_TYPES, this.contextTypes);
+            pb.putString(KEY_PCI_ADDRESS, this.pciAddress);
+            pb.putString(KEY_RENDERER_FEATURES, this.rendererFeatures);
+            pb.putBoolean(KEY_RENDERER_USE_EGL, this.rendererUseEgl);
+            pb.putBoolean(KEY_RENDERER_USE_GLES, this.rendererUseGles);
+            pb.putBoolean(KEY_RENDERER_USE_GLX, this.rendererUseGlx);
+            pb.putBoolean(KEY_RENDERER_USE_SURFACELESS, this.rendererUseSurfaceless);
+            pb.putBoolean(KEY_RENDERER_USE_VULKAN, this.rendererUseVulkan);
+            return pb;
+        }
+
+        /** @hide */
+        public static class Builder {
+            private String backend;
+            private String[] contextTypes;
+            private String pciAddress;
+            private String rendererFeatures;
+            private boolean rendererUseEgl = true;
+            private boolean rendererUseGles = true;
+            private boolean rendererUseGlx = false;
+            private boolean rendererUseSurfaceless = true;
+            private boolean rendererUseVulkan = false;
+
+            /** @hide */
+            public Builder() {}
+
+            /** @hide */
+            public Builder setBackend(String backend) {
+                this.backend = backend;
+                return this;
+            }
+
+            /** @hide */
+            public Builder setContextTypes(String[] contextTypes) {
+                this.contextTypes = contextTypes;
+                return this;
+            }
+
+            /** @hide */
+            public Builder setPciAddress(String pciAddress) {
+                this.pciAddress = pciAddress;
+                return this;
+            }
+
+            /** @hide */
+            public Builder setRendererFeatures(String rendererFeatures) {
+                this.rendererFeatures = rendererFeatures;
+                return this;
+            }
+
+            /** @hide */
+            public Builder setRendererUseEgl(Boolean rendererUseEgl) {
+                this.rendererUseEgl = rendererUseEgl;
+                return this;
+            }
+
+            /** @hide */
+            public Builder setRendererUseGles(Boolean rendererUseGles) {
+                this.rendererUseGles = rendererUseGles;
+                return this;
+            }
+
+            /** @hide */
+            public Builder setRendererUseGlx(Boolean rendererUseGlx) {
+                this.rendererUseGlx = rendererUseGlx;
+                return this;
+            }
+
+            /** @hide */
+            public Builder setRendererUseSurfaceless(Boolean rendererUseSurfaceless) {
+                this.rendererUseSurfaceless = rendererUseSurfaceless;
+                return this;
+            }
+
+            /** @hide */
+            public Builder setRendererUseVulkan(Boolean rendererUseVulkan) {
+                this.rendererUseVulkan = rendererUseVulkan;
+                return this;
+            }
+
+            /** @hide */
+            public GpuConfig build() {
+                return new GpuConfig(
+                        backend,
+                        contextTypes,
+                        pciAddress,
+                        rendererFeatures,
+                        rendererUseEgl,
+                        rendererUseGles,
+                        rendererUseGlx,
+                        rendererUseSurfaceless,
+                        rendererUseVulkan);
+            }
+        }
+    }
 }
diff --git a/libs/hypervisor_props/src/lib.rs b/libs/hypervisor_props/src/lib.rs
index 14614fd..6665bc5 100644
--- a/libs/hypervisor_props/src/lib.rs
+++ b/libs/hypervisor_props/src/lib.rs
@@ -37,3 +37,8 @@
 pub fn version() -> Result<Option<String>> {
     Ok(hypervisorproperties::hypervisor_version()?)
 }
+
+/// Returns if the hypervisor is pKVM
+pub fn is_pkvm() -> Result<bool> {
+    Ok(version()?.unwrap_or_default().starts_with("kvm") && is_protected_vm_supported()?)
+}
diff --git a/microdroid_manager/src/vm_secret.rs b/microdroid_manager/src/vm_secret.rs
index ec40b45..c16a45e 100644
--- a/microdroid_manager/src/vm_secret.rs
+++ b/microdroid_manager/src/vm_secret.rs
@@ -20,7 +20,7 @@
 use secretkeeper_comm::data_types::request::Request;
 use binder::{Strong};
 use coset::{CoseKey, CborSerializable, CborOrdering};
-use dice_policy_builder::{CertIndex, ConstraintSpec, ConstraintType, policy_for_dice_chain, MissingAction, WILDCARD_FULL_ARRAY};
+use dice_policy_builder::{TargetEntry, ConstraintSpec, ConstraintType, policy_for_dice_chain, MissingAction, WILDCARD_FULL_ARRAY};
 use diced_open_dice::{DiceArtifacts, OwnedDiceArtifacts};
 use keystore2_crypto::ZVec;
 use openssl::hkdf::hkdf;
@@ -45,9 +45,10 @@
 const SUBCOMPONENT_DESCRIPTORS: i64 = -71002;
 const SUBCOMPONENT_SECURITY_VERSION: i64 = 2;
 const SUBCOMPONENT_AUTHORITY_HASH: i64 = 4;
-// Index of DiceChainEntry corresponding to Payload (relative to the end considering DICE Chain
-// as an array)
-const PAYLOAD_INDEX_FROM_END: usize = 0;
+// See dice_for_avf_guest.cddl for the `component_name` used by different boot stages in guest VM.
+const MICRODROID_PAYLOAD_COMPONENT_NAME: &str = "Microdroid payload";
+const GUEST_OS_COMPONENT_NAME: &str = "vm_entry";
+const INSTANCE_HASH_KEY: i64 = -71003;
 
 // Generated using hexdump -vn32 -e'14/1 "0x%02X, " 1 "\n"' /dev/urandom
 const SALT_ENCRYPTED_STORE: &[u8] = &[
@@ -173,25 +174,27 @@
 //     microdroid_manager/src/vm_config.cddl):
 //       - GreaterOrEqual on SECURITY_VERSION (Required)
 //       - ExactMatch on AUTHORITY_HASH (Required).
+//  5. ExactMatch on Instance Hash (Required) - This uniquely identifies one VM instance from
+//     another even if they are running the exact same images.
 fn sealing_policy(dice: &[u8]) -> Result<Vec<u8>, String> {
-    let constraint_spec = [
+    let constraint_spec = vec![
         ConstraintSpec::new(
             ConstraintType::ExactMatch,
             vec![AUTHORITY_HASH],
             MissingAction::Fail,
-            CertIndex::All,
+            TargetEntry::All,
         ),
         ConstraintSpec::new(
             ConstraintType::ExactMatch,
             vec![MODE],
             MissingAction::Fail,
-            CertIndex::All,
+            TargetEntry::All,
         ),
         ConstraintSpec::new(
             ConstraintType::GreaterOrEqual,
             vec![CONFIG_DESC, SECURITY_VERSION],
             MissingAction::Ignore,
-            CertIndex::All,
+            TargetEntry::All,
         ),
         ConstraintSpec::new(
             ConstraintType::GreaterOrEqual,
@@ -202,7 +205,7 @@
                 SUBCOMPONENT_SECURITY_VERSION,
             ],
             MissingAction::Fail,
-            CertIndex::FromEnd(PAYLOAD_INDEX_FROM_END),
+            TargetEntry::ByName(MICRODROID_PAYLOAD_COMPONENT_NAME.to_string()),
         ),
         ConstraintSpec::new(
             ConstraintType::ExactMatch,
@@ -213,11 +216,17 @@
                 SUBCOMPONENT_AUTHORITY_HASH,
             ],
             MissingAction::Fail,
-            CertIndex::FromEnd(PAYLOAD_INDEX_FROM_END),
+            TargetEntry::ByName(MICRODROID_PAYLOAD_COMPONENT_NAME.to_string()),
+        ),
+        ConstraintSpec::new(
+            ConstraintType::ExactMatch,
+            vec![CONFIG_DESC, INSTANCE_HASH_KEY],
+            MissingAction::Fail,
+            TargetEntry::ByName(GUEST_OS_COMPONENT_NAME.to_string()),
         ),
     ];
 
-    policy_for_dice_chain(dice, &constraint_spec)?
+    policy_for_dice_chain(dice, constraint_spec)?
         .to_vec()
         .map_err(|e| format!("DicePolicy construction failed {e:?}"))
 }
diff --git a/tests/testapk/src/java/com/android/microdroid/test/MicrodroidTests.java b/tests/testapk/src/java/com/android/microdroid/test/MicrodroidTests.java
index 4d8dac6..4141903 100644
--- a/tests/testapk/src/java/com/android/microdroid/test/MicrodroidTests.java
+++ b/tests/testapk/src/java/com/android/microdroid/test/MicrodroidTests.java
@@ -28,8 +28,6 @@
 import static com.google.common.truth.Truth.assertThat;
 import static com.google.common.truth.Truth.assertWithMessage;
 import static com.google.common.truth.TruthJUnit.assume;
-import com.android.virt.vm_attestation.testservice.IAttestationService.AttestationStatus;
-import com.android.virt.vm_attestation.testservice.IAttestationService.SigningResult;
 
 import static org.junit.Assert.assertThrows;
 import static org.junit.Assert.assertTrue;
@@ -70,8 +68,15 @@
 import com.android.microdroid.testservice.IAppCallback;
 import com.android.microdroid.testservice.ITestService;
 import com.android.microdroid.testservice.IVmCallback;
+import com.android.virt.vm_attestation.testservice.IAttestationService.AttestationStatus;
+import com.android.virt.vm_attestation.testservice.IAttestationService.SigningResult;
 import com.android.virt.vm_attestation.util.X509Utils;
 
+import co.nstant.in.cbor.CborDecoder;
+import co.nstant.in.cbor.model.Array;
+import co.nstant.in.cbor.model.DataItem;
+import co.nstant.in.cbor.model.MajorType;
+
 import com.google.common.base.Strings;
 import com.google.common.truth.BooleanSubject;
 
@@ -113,17 +118,13 @@
 import java.util.concurrent.atomic.AtomicReference;
 import java.util.stream.Stream;
 
-import co.nstant.in.cbor.CborDecoder;
-import co.nstant.in.cbor.model.Array;
-import co.nstant.in.cbor.model.DataItem;
-import co.nstant.in.cbor.model.MajorType;
-
 @RunWith(Parameterized.class)
 public class MicrodroidTests extends MicrodroidDeviceTestBase {
     private static final String TAG = "MicrodroidTests";
     private static final String TEST_APP_PACKAGE_NAME = "com.android.microdroid.test";
     private static final String VM_ATTESTATION_PAYLOAD_PATH = "libvm_attestation_test_payload.so";
     private static final String VM_ATTESTATION_MESSAGE = "Hello RKP from AVF!";
+    private static final int ENCRYPTED_STORAGE_BYTES = 4_000_000;
 
     @Rule public Timeout globalTimeout = Timeout.seconds(300);
 
@@ -1602,7 +1603,7 @@
                 newVmConfigBuilderWithPayloadBinary("MicrodroidTestNativeLib.so")
                         .setDebugLevel(DEBUG_LEVEL_FULL);
         if (encryptedStoreEnabled) {
-            builder.setEncryptedStorageBytes(4_000_000);
+            builder.setEncryptedStorageBytes(ENCRYPTED_STORAGE_BYTES);
         }
         VirtualMachineConfig config = builder.build();
         String vmNameOrig = "test_vm_orig";
@@ -1656,7 +1657,7 @@
         VirtualMachineConfig config =
                 newVmConfigBuilderWithPayloadBinary("MicrodroidTestNativeLib.so")
                         .setMemoryBytes(minMemoryRequired())
-                        .setEncryptedStorageBytes(4_000_000)
+                        .setEncryptedStorageBytes(ENCRYPTED_STORAGE_BYTES)
                         .setDebugLevel(DEBUG_LEVEL_FULL)
                         .build();
         VirtualMachine vm = forceCreateNewVirtualMachine("test_vm", config);
@@ -1683,7 +1684,7 @@
         VirtualMachineConfig config =
                 newVmConfigBuilderWithPayloadBinary("MicrodroidTestNativeLib.so")
                         .setMemoryBytes(minMemoryRequired())
-                        .setEncryptedStorageBytes(4_000_000)
+                        .setEncryptedStorageBytes(ENCRYPTED_STORAGE_BYTES)
                         .setDebugLevel(DEBUG_LEVEL_FULL)
                         .build();
 
@@ -1790,7 +1791,7 @@
         VirtualMachineConfig config =
                 newVmConfigBuilderWithPayloadBinary("MicrodroidTestNativeLib.so")
                         .setMemoryBytes(minMemoryRequired())
-                        .setEncryptedStorageBytes(4_000_000)
+                        .setEncryptedStorageBytes(ENCRYPTED_STORAGE_BYTES)
                         .setDebugLevel(DEBUG_LEVEL_FULL)
                         .build();
         VirtualMachine vm = forceCreateNewVirtualMachine("test_vm_a", config);
@@ -2282,7 +2283,7 @@
         VirtualMachineConfig vmConfig =
                 newVmConfigBuilderWithPayloadBinary("MicrodroidTestNativeLib.so")
                         .setDebugLevel(DEBUG_LEVEL_FULL)
-                        .setEncryptedStorageBytes(4_000_000)
+                        .setEncryptedStorageBytes(ENCRYPTED_STORAGE_BYTES)
                         .build();
         VirtualMachine vm = forceCreateNewVirtualMachine("test_vm_encstore_no_exec", vmConfig);
 
diff --git a/virtualizationmanager/src/aidl.rs b/virtualizationmanager/src/aidl.rs
index 8e1b6bb..ac70509 100644
--- a/virtualizationmanager/src/aidl.rs
+++ b/virtualizationmanager/src/aidl.rs
@@ -17,7 +17,7 @@
 use crate::{get_calling_pid, get_calling_uid, get_this_pid};
 use crate::atom::{write_vm_booted_stats, write_vm_creation_stats};
 use crate::composite::make_composite_image;
-use crate::crosvm::{CrosvmConfig, DiskFile, DisplayConfig, InputDeviceOption, PayloadState, VmContext, VmInstance, VmState};
+use crate::crosvm::{CrosvmConfig, DiskFile, DisplayConfig, GpuConfig, InputDeviceOption, PayloadState, VmContext, VmInstance, VmState};
 use crate::debug_config::DebugConfig;
 use crate::dt_overlay::{create_device_tree_overlay, VM_DT_OVERLAY_MAX_SIZE, VM_DT_OVERLAY_PATH};
 use crate::payload::{add_microdroid_payload_images, add_microdroid_system_images, add_microdroid_vendor_image};
@@ -535,6 +535,16 @@
         } else {
             None
         };
+        let gpu_config = if cfg!(paravirtualized_devices) {
+            config
+                .gpuConfig
+                .as_ref()
+                .map(GpuConfig::new)
+                .transpose()
+                .or_binder_exception(ExceptionCode::ILLEGAL_ARGUMENT)?
+        } else {
+            None
+        };
 
         let input_device_options = if cfg!(paravirtualized_devices) {
             config
@@ -600,6 +610,7 @@
             virtio_snd_backend,
             console_input_device: config.consoleInputDevice.clone(),
             boost_uclamp: config.boostUclamp,
+            gpu_config,
         };
         let instance = Arc::new(
             VmInstance::new(
diff --git a/virtualizationmanager/src/crosvm.rs b/virtualizationmanager/src/crosvm.rs
index 79c7b7f..3722d4d 100644
--- a/virtualizationmanager/src/crosvm.rs
+++ b/virtualizationmanager/src/crosvm.rs
@@ -47,6 +47,7 @@
     MemoryTrimLevel::MemoryTrimLevel,
     VirtualMachineAppConfig::DebugLevel::DebugLevel,
     DisplayConfig::DisplayConfig as DisplayConfigParcelable,
+    GpuConfig::GpuConfig as GpuConfigParcelable,
 };
 use android_system_virtualizationservice_internal::aidl::android::system::virtualizationservice_internal::IGlobalVmContext::IGlobalVmContext;
 use android_system_virtualizationservice_internal::aidl::android::system::virtualizationservice_internal::IBoundDevice::IBoundDevice;
@@ -133,6 +134,7 @@
     pub virtio_snd_backend: Option<String>,
     pub console_input_device: Option<String>,
     pub boost_uclamp: bool,
+    pub gpu_config: Option<GpuConfig>,
 }
 
 #[derive(Debug)]
@@ -155,6 +157,37 @@
     }
 }
 
+#[derive(Debug)]
+pub struct GpuConfig {
+    pub backend: Option<String>,
+    pub context_types: Option<Vec<String>>,
+    pub pci_address: Option<String>,
+    pub renderer_features: Option<String>,
+    pub renderer_use_egl: Option<bool>,
+    pub renderer_use_gles: Option<bool>,
+    pub renderer_use_glx: Option<bool>,
+    pub renderer_use_surfaceless: Option<bool>,
+    pub renderer_use_vulkan: Option<bool>,
+}
+
+impl GpuConfig {
+    pub fn new(raw_config: &GpuConfigParcelable) -> Result<GpuConfig> {
+        Ok(GpuConfig {
+            backend: raw_config.backend.clone(),
+            context_types: raw_config.contextTypes.clone().map(|context_types| {
+                context_types.iter().filter_map(|context_type| context_type.clone()).collect()
+            }),
+            pci_address: raw_config.pciAddress.clone(),
+            renderer_features: raw_config.rendererFeatures.clone(),
+            renderer_use_egl: Some(raw_config.rendererUseEgl),
+            renderer_use_gles: Some(raw_config.rendererUseGles),
+            renderer_use_glx: Some(raw_config.rendererUseGlx),
+            renderer_use_surfaceless: Some(raw_config.rendererUseSurfaceless),
+            renderer_use_vulkan: Some(raw_config.rendererUseVulkan),
+        })
+    }
+}
+
 fn try_into_non_zero_u32(value: i32) -> Result<NonZeroU32> {
     let u32_value = value.try_into()?;
     NonZeroU32::new(u32_value).ok_or(anyhow!("value should be greater than 0"))
@@ -855,6 +888,8 @@
         command.arg("--no-balloon");
     }
 
+    let mut memory_mib = config.memory_mib;
+
     if config.protected {
         match system_properties::read(SYSPROP_CUSTOM_PVMFW_PATH)? {
             Some(pvmfw_path) if !pvmfw_path.is_empty() => {
@@ -870,6 +905,12 @@
         let swiotlb_size_mib = 2 * virtio_pci_device_count as u32;
         command.arg("--swiotlb").arg(swiotlb_size_mib.to_string());
 
+        // b/346770542 for consistent "usable" memory across protected and non-protected VMs under
+        // pKVM.
+        if hypervisor_props::is_pkvm()? {
+            memory_mib = memory_mib.map(|m| m.saturating_add(swiotlb_size_mib));
+        }
+
         // Workaround to keep crash_dump from trying to read protected guest memory.
         // Context in b/238324526.
         command.arg("--unmap-guest-memory-on-fork");
@@ -891,7 +932,7 @@
         command.arg("--params").arg("console=hvc0");
     }
 
-    if let Some(memory_mib) = config.memory_mib {
+    if let Some(memory_mib) = memory_mib {
         command.arg("--mem").arg(memory_mib.to_string());
     }
 
@@ -1001,12 +1042,48 @@
     }
 
     if cfg!(paravirtualized_devices) {
+        if let Some(gpu_config) = &config.gpu_config {
+            let mut gpu_args = Vec::new();
+            if let Some(backend) = &gpu_config.backend {
+                gpu_args.push(format!("backend={}", backend));
+            }
+            if let Some(context_types) = &gpu_config.context_types {
+                gpu_args.push(format!("context-types={}", context_types.join(":")));
+            }
+            if let Some(pci_address) = &gpu_config.pci_address {
+                gpu_args.push(format!("pci-address={}", pci_address));
+            }
+            if let Some(renderer_features) = &gpu_config.renderer_features {
+                gpu_args.push(format!("renderer-features={}", renderer_features));
+            }
+            if gpu_config.renderer_use_egl.unwrap_or(false) {
+                gpu_args.push("egl=true".to_string());
+            }
+            if gpu_config.renderer_use_gles.unwrap_or(false) {
+                gpu_args.push("gles=true".to_string());
+            }
+            if gpu_config.renderer_use_glx.unwrap_or(false) {
+                gpu_args.push("glx=true".to_string());
+            }
+            if gpu_config.renderer_use_surfaceless.unwrap_or(false) {
+                gpu_args.push("surfaceless=true".to_string());
+            }
+            if gpu_config.renderer_use_vulkan.unwrap_or(false) {
+                gpu_args.push("vulkan=true".to_string());
+            }
+            command.arg(format!("--gpu={}", gpu_args.join(",")));
+        }
         if let Some(display_config) = &config.display_config {
-            command.arg("--gpu")
-            // TODO(b/331708504): support backend config as well
-            .arg("backend=virglrenderer,context-types=virgl2,egl=true,surfaceless=true,glx=false,gles=true")
-            .arg(format!("--gpu-display=mode=windowed[{},{}],dpi=[{},{}],refresh-rate={}", display_config.width, display_config.height, display_config.horizontal_dpi, display_config.vertical_dpi, display_config.refresh_rate))
-            .arg(format!("--android-display-service={}", config.name));
+            command
+                .arg(format!(
+                    "--gpu-display=mode=windowed[{},{}],dpi=[{},{}],refresh-rate={}",
+                    display_config.width,
+                    display_config.height,
+                    display_config.horizontal_dpi,
+                    display_config.vertical_dpi,
+                    display_config.refresh_rate
+                ))
+                .arg(format!("--android-display-service={}", config.name));
         }
     }
 
diff --git a/virtualizationservice/aidl/android/system/virtualizationservice/GpuConfig.aidl b/virtualizationservice/aidl/android/system/virtualizationservice/GpuConfig.aidl
new file mode 100644
index 0000000..1cd4dc6
--- /dev/null
+++ b/virtualizationservice/aidl/android/system/virtualizationservice/GpuConfig.aidl
@@ -0,0 +1,28 @@
+/*
+ * Copyright 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package android.system.virtualizationservice;
+
+parcelable GpuConfig {
+    @nullable String backend;
+    @nullable String[] contextTypes;
+    @nullable String pciAddress;
+    @nullable String rendererFeatures;
+    boolean rendererUseEgl = false;
+    boolean rendererUseGles = false;
+    boolean rendererUseGlx = false;
+    boolean rendererUseSurfaceless = false;
+    boolean rendererUseVulkan = false;
+}
diff --git a/virtualizationservice/aidl/android/system/virtualizationservice/VirtualMachineRawConfig.aidl b/virtualizationservice/aidl/android/system/virtualizationservice/VirtualMachineRawConfig.aidl
index 0d175dd..69664b4 100644
--- a/virtualizationservice/aidl/android/system/virtualizationservice/VirtualMachineRawConfig.aidl
+++ b/virtualizationservice/aidl/android/system/virtualizationservice/VirtualMachineRawConfig.aidl
@@ -18,6 +18,7 @@
 import android.system.virtualizationservice.CpuTopology;
 import android.system.virtualizationservice.DiskImage;
 import android.system.virtualizationservice.DisplayConfig;
+import android.system.virtualizationservice.GpuConfig;
 import android.system.virtualizationservice.InputDevice;
 
 /** Raw configuration for running a VM. */
@@ -94,4 +95,6 @@
 
     /** Enable boost UClamp for less variance during testing/benchmarking */
     boolean boostUclamp;
+
+    @nullable GpuConfig gpuConfig;
 }
diff --git a/virtualizationservice/src/aidl.rs b/virtualizationservice/src/aidl.rs
index 241baa5..ae8d1da 100644
--- a/virtualizationservice/src/aidl.rs
+++ b/virtualizationservice/src/aidl.rs
@@ -454,7 +454,7 @@
             .context("Failed to allocate instance_id")
             .or_service_specific_exception(-1)?;
         let uid = get_calling_uid();
-        info!("Allocated a VM's instance_id: {:?}, for uid: {:?}", hex::encode(id), uid);
+        info!("Allocated a VM's instance_id: {:?}..., for uid: {:?}", &hex::encode(id)[..8], uid);
         let state = &mut *self.state.lock().unwrap();
         if let Some(sk_state) = &mut state.sk_state {
             let user_id = multiuser_get_user_id(uid);
diff --git a/vmbase/Android.bp b/vmbase/Android.bp
index f01e8aa..ee12e85 100644
--- a/vmbase/Android.bp
+++ b/vmbase/Android.bp
@@ -57,6 +57,8 @@
         hwaddress: false,
     },
     native_coverage: false,
+    // TODO(b/346974429): Workaround pvmfw failure when enabling full LTO
+    lto_O0: true,
 }
 
 // Used by cc_binary when producing the ELF of a vmbase-based binary.
diff --git a/vmlauncher_app/java/com/android/virtualization/vmlauncher/MainActivity.java b/vmlauncher_app/java/com/android/virtualization/vmlauncher/MainActivity.java
index 5355313..33e4755 100644
--- a/vmlauncher_app/java/com/android/virtualization/vmlauncher/MainActivity.java
+++ b/vmlauncher_app/java/com/android/virtualization/vmlauncher/MainActivity.java
@@ -32,6 +32,7 @@
 import android.system.virtualmachine.VirtualMachineConfig;
 import android.system.virtualmachine.VirtualMachineCustomImageConfig;
 import android.system.virtualmachine.VirtualMachineCustomImageConfig.DisplayConfig;
+import android.system.virtualmachine.VirtualMachineCustomImageConfig.GpuConfig;
 import android.system.virtualmachine.VirtualMachineException;
 import android.system.virtualmachine.VirtualMachineManager;
 import android.util.DisplayMetrics;
@@ -63,6 +64,7 @@
 import java.nio.ByteOrder;
 import java.nio.file.Files;
 import java.nio.file.Path;
+import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
@@ -125,6 +127,47 @@
             if (json.has("console_input_device")) {
                 configBuilder.setConsoleInputDevice(json.getString("console_input_device"));
             }
+            if (json.has("gpu")) {
+                JSONObject gpuJson = json.getJSONObject("gpu");
+
+                GpuConfig.Builder gpuConfigBuilder = new GpuConfig.Builder();
+
+                if (gpuJson.has("backend")) {
+                    gpuConfigBuilder.setBackend(gpuJson.getString("backend"));
+                }
+                if (gpuJson.has("context_types")) {
+                    ArrayList<String> contextTypes = new ArrayList<String>();
+                    JSONArray contextTypesJson = gpuJson.getJSONArray("context_types");
+                    for (int i = 0; i < contextTypesJson.length(); i++) {
+                        contextTypes.add(contextTypesJson.getString(i));
+                    }
+                    gpuConfigBuilder.setContextTypes(contextTypes.toArray(new String[0]));
+                }
+                if (gpuJson.has("pci_address")) {
+                    gpuConfigBuilder.setPciAddress(gpuJson.getString("pci_address"));
+                }
+                if (gpuJson.has("renderer_features")) {
+                    gpuConfigBuilder.setRendererFeatures(gpuJson.getString("renderer_features"));
+                }
+                if (gpuJson.has("renderer_use_egl")) {
+                    gpuConfigBuilder.setRendererUseEgl(gpuJson.getBoolean("renderer_use_egl"));
+                }
+                if (gpuJson.has("renderer_use_gles")) {
+                    gpuConfigBuilder.setRendererUseGles(gpuJson.getBoolean("renderer_use_gles"));
+                }
+                if (gpuJson.has("renderer_use_glx")) {
+                    gpuConfigBuilder.setRendererUseGlx(gpuJson.getBoolean("renderer_use_glx"));
+                }
+                if (gpuJson.has("renderer_use_surfaceless")) {
+                    gpuConfigBuilder.setRendererUseSurfaceless(
+                            gpuJson.getBoolean("renderer_use_surfaceless"));
+                }
+                if (gpuJson.has("renderer_use_vulkan")) {
+                    gpuConfigBuilder.setRendererUseVulkan(
+                            gpuJson.getBoolean("renderer_use_vulkan"));
+                }
+                customImageConfigBuilder.setGpuConfig(gpuConfigBuilder.build());
+            }
 
             configBuilder.setMemoryBytes(8L * 1024 * 1024 * 1024 /* 8 GB */);
             WindowMetrics windowMetrics = getWindowManager().getCurrentWindowMetrics();