[PATCH] vcpu size/align for TDX

[Kai Huang <kai.huang@xxxxxxxxx>]

Signed-off-by: Kai Huang <kai.huang@xxxxxxxxx>
---
 arch/x86/include/asm/kvm_host.h |  2 ++
 arch/x86/kvm/svm/svm.c          |  6 ++++--
 arch/x86/kvm/vmx/main.c         | 20 +++-----------------
 arch/x86/kvm/vmx/tdx.c          |  7 +++++++
 4 files changed, 16 insertions(+), 19 deletions(-)

diff --git a/arch/x86/include/asm/kvm_host.h b/arch/x86/include/asm/kvm_host.h
index ed9b65785a24..7b96b6b30a5c 100644
--- a/arch/x86/include/asm/kvm_host.h
+++ b/arch/x86/include/asm/kvm_host.h
@@ -1684,6 +1684,8 @@ struct kvm_x86_ops {
        void (*vcpu_after_set_cpuid)(struct kvm_vcpu *vcpu);

        unsigned int vm_size;
+       unsigned int vcpu_size;
+       unsigned int vcpu_align;
        int (*vm_init)(struct kvm *kvm);
        void (*vm_destroy)(struct kvm *kvm);
        void (*vm_pre_destroy)(struct kvm *kvm);
diff --git a/arch/x86/kvm/svm/svm.c b/arch/x86/kvm/svm/svm.c
index dc8e9af49f11..6a43d6402219 100644
--- a/arch/x86/kvm/svm/svm.c
+++ b/arch/x86/kvm/svm/svm.c
@@ -5094,6 +5094,9 @@ static struct kvm_x86_ops svm_x86_ops __initdata = {
        .vcpu_reset = svm_vcpu_reset,

        .vm_size = sizeof(struct kvm_svm),
+       .vcpu_size = sizeof(struct vcpu_svm),
+       .vcpu_align = __alignof__(struct vcpu_svm),
+
        .vm_init = svm_vm_init,
        .vm_destroy = svm_vm_destroy,

@@ -5543,8 +5546,7 @@ static int __init svm_init(void)
         * Common KVM initialization _must_ come last, after this, /dev/kvm is
         * exposed to userspace!
         */
-       r = kvm_init(sizeof(struct vcpu_svm), __alignof__(struct vcpu_svm),
-                    THIS_MODULE);
+       r = kvm_init(svm_x86_ops.vcpu_size, svm_x86_ops.vcpu_align,
THIS_MODULE);
        if (r)
                goto err_kvm_init;

diff --git a/arch/x86/kvm/vmx/main.c b/arch/x86/kvm/vmx/main.c
index c064d79b7550..86fdcaca7061 100644
--- a/arch/x86/kvm/vmx/main.c
+++ b/arch/x86/kvm/vmx/main.c
@@ -879,6 +879,8 @@ struct kvm_x86_ops vt_x86_ops __initdata = {
        .has_emulated_msr = vt_has_emulated_msr,

        .vm_size = sizeof(struct kvm_vmx),
+       .vcpu_size = sizeof(struct vcpu_vmx),
+       .vcpu_align = __alignof__(struct vcpu_vmx),

        .vm_init = vt_vm_init,
        .vm_pre_destroy = vt_vm_pre_destroy,
@@ -1035,7 +1037,6 @@ module_exit(vt_exit);

 static int __init vt_init(void)
 {
-       unsigned vcpu_size, vcpu_align;
        int r;

        r = vmx_init();
@@ -1047,26 +1048,11 @@ static int __init vt_init(void)
        if (r)
                goto err_tdx_bringup;

-       /*
-        * TDX and VMX have different vCPU structures.  Calculate the
-        * maximum size/align so that kvm_init() can use the larger
-        * values to create the kmem_vcpu_cache.
-        */
-       vcpu_size = sizeof(struct vcpu_vmx);
-       vcpu_align = __alignof__(struct vcpu_vmx);
-       if (enable_tdx) {
-               vcpu_size = max_t(unsigned, vcpu_size,
-                               sizeof(struct vcpu_tdx));
-               vcpu_align = max_t(unsigned, vcpu_align,
-                               __alignof__(struct vcpu_tdx));
-               kvm_caps.supported_vm_types |= BIT(KVM_X86_TDX_VM);
-       }
-
        /*
         * Common KVM initialization _must_ come last, after this, /dev/kvm is
         * exposed to userspace!
         */
-       r = kvm_init(vcpu_size, vcpu_align, THIS_MODULE);
+       r = kvm_init(vt_x86_ops.vcpu_size, vt_x86_ops.vcpu_align, THIS_MODULE);
        if (r)
                goto err_kvm_init;

diff --git a/arch/x86/kvm/vmx/tdx.c b/arch/x86/kvm/vmx/tdx.c
index b4985a64501c..dc44cabf89e6 100644
--- a/arch/x86/kvm/vmx/tdx.c
+++ b/arch/x86/kvm/vmx/tdx.c
@@ -3522,6 +3522,9 @@ int __init tdx_bringup(void)
                enable_tdx = 0;
        }

+       if (enable_tdx)
+               kvm_caps.supported_vm_types |= BIT(KVM_X86_TDX_VM);
+
        return r;

 success_disable_tdx:
@@ -3537,6 +3540,10 @@ void __init tdx_hardware_setup(void)
         * disabled but KVM will continue loading (see tdx_bringup()).
         */
        vt_x86_ops.vm_size = max_t(unsigned int, vt_x86_ops.vm_size,
sizeof(struct kvm_tdx));
+       vt_x86_ops.vcpu_size = max_t(unsigned int, vt_x86_ops.vcpu_size,
+                       sizeof(struct vcpu_tdx));
+       vt_x86_ops.vcpu_align = max_t(unsigned int, vt_x86_ops.vcpu_align,
+                       __alignof__(struct vcpu_tdx));
 
        vt_x86_ops.link_external_spt = tdx_sept_link_private_spt;
        vt_x86_ops.set_external_spte = tdx_sept_set_private_spte;
-- 
2.49.0