/* SPDX-License-Identifier: GPL-2.0 */
 #include <asm/asm-offsets.h>
+#include <asm/frame.h>
 #include <asm/tdx.h>
 
 /*
  *            TDX module.
  */
 .macro TDX_MODULE_CALL host:req
+       FRAME_BEGIN
        /*
         * R12 will be used as temporary storage for struct tdx_module_output
         * pointer. Since R12-R15 registers are not used by TDCALL/SEAMCALL
        mov %rsi, %rcx
        /* Leave input param 2 in RDX */
 
-       .if \host
+.if \host
        seamcall
        /*
         * SEAMCALL instruction is essentially a VMExit from VMX root
         * This value will never be used as actual SEAMCALL error code as
         * it is from the Reserved status code class.
         */
-       jnc .Lno_vmfailinvalid
-       mov $TDX_SEAMCALL_VMFAILINVALID, %rax
-.Lno_vmfailinvalid:
-
-       .else
+       jc .Lseamcall_vmfailinvalid
+.else
        tdcall
-       .endif
+.endif
 
        /*
         * Fetch output pointer from stack to R12 (It is used
         * Other registers may contain details of the failure.
         */
        test %r12, %r12
-       jz .Lno_output_struct
+       jz .Lout
 
        /* Copy result registers to output struct: */
        movq %rcx, TDX_MODULE_rcx(%r12)
        movq %r10, TDX_MODULE_r10(%r12)
        movq %r11, TDX_MODULE_r11(%r12)
 
-.Lno_output_struct:
+.Lout:
        /* Restore the state of R12 register */
        pop %r12
+
+       FRAME_END
+       RET
+
+.if \host
+.Lseamcall_vmfailinvalid:
+       mov $TDX_SEAMCALL_VMFAILINVALID, %rax
+       /* pop the unused output pointer back to %r9 */
+       pop %r9
+       jmp .Lout
+.endif /* \host */
+
 .endm