From b6cf27507c079573dc212b4f8f54d6313186dc61 Mon Sep 17 00:00:00 2001
From: Hsy-Intel <siyuan.hui@intel.com>
Date: Wed, 19 Mar 2025 19:03:29 +0000
Subject: [PATCH] Use TDVMCALL serial output in the early stages of kernel
 initialization

---
 ostd/src/arch/x86/mod.rs       |  9 ++++-----
 ostd/src/lib.rs                |  9 ++++++++-
 ostd/src/mm/frame/allocator.rs | 10 ++++++++--
 ostd/src/mm/frame/meta.rs      | 15 ++++++++++-----
 ostd/src/mm/kspace/mod.rs      |  9 ++++++++-
 5 files changed, 38 insertions(+), 14 deletions(-)

diff --git a/ostd/src/arch/x86/mod.rs b/ostd/src/arch/x86/mod.rs
index 442e0a02e..12943ed87 100644
--- a/ostd/src/arch/x86/mod.rs
+++ b/ostd/src/arch/x86/mod.rs
@@ -27,10 +27,7 @@ cfg_if! {
     if #[cfg(feature = "cvm_guest")] {
         pub(crate) mod tdx_guest;
 
-        use {
-            crate::early_println,
-            ::tdx_guest::{init_tdx, tdcall::InitError},
-        };
+        use ::tdx_guest::{init_tdx, tdcall::InitError};
     }
 }
 
@@ -44,9 +41,11 @@ use log::{info, warn};
 
 #[cfg(feature = "cvm_guest")]
 pub(crate) fn init_cvm_guest() {
+    use ::tdx_guest::serial_println;
+
     match init_tdx() {
         Ok(td_info) => {
-            early_println!(
+            serial_println!(
                 "[kernel] Intel TDX initialized\n[kernel] td gpaw: {}, td attributes: {:?}",
                 td_info.gpaw,
                 td_info.attributes
diff --git a/ostd/src/lib.rs b/ostd/src/lib.rs
index 95d99f027..fe6be00ea 100644
--- a/ostd/src/lib.rs
+++ b/ostd/src/lib.rs
@@ -77,7 +77,10 @@ unsafe fn init() {
         mm::frame::allocator::init_early_allocator();
     }
 
-    arch::serial::init();
+    if_tdx_enabled!({
+    } else {
+        arch::serial::init();
+    });
 
     logger::init();
 
@@ -106,6 +109,10 @@ unsafe fn init() {
 
     unsafe { arch::late_init_on_bsp() };
 
+    if_tdx_enabled!({
+        arch::serial::init();
+    });
+
     smp::init();
 
     // SAFETY: This function is called only once on the BSP.
diff --git a/ostd/src/mm/frame/allocator.rs b/ostd/src/mm/frame/allocator.rs
index 006e24b49..8d1b4648e 100644
--- a/ostd/src/mm/frame/allocator.rs
+++ b/ostd/src/mm/frame/allocator.rs
@@ -10,7 +10,7 @@ use super::{meta::AnyFrameMeta, segment::Segment, Frame};
 use crate::{
     boot::memory_region::MemoryRegionType,
     error::Error,
-    impl_frame_meta_for,
+    if_tdx_enabled, impl_frame_meta_for,
     mm::{paddr_to_vaddr, Paddr, PAGE_SIZE},
     prelude::*,
     util::range_difference,
@@ -212,7 +212,13 @@ pub(crate) unsafe fn init() {
             // Truncate the early allocated frames if there is an overlap.
             for r1 in range_difference(&(region.base()..region.end()), &range_1) {
                 for r2 in range_difference(&r1, &range_2) {
-                    log::info!("Adding free frames to the allocator: {:x?}", r2);
+                    if_tdx_enabled!({
+                        use tdx_guest::serial_println;
+
+                        serial_println!("Adding free frames to the allocator: {:x?}", r2);
+                    } else {
+                        log::info!("Adding free frames to the allocator: {:x?}", r2);
+                    });
                     get_global_frame_allocator().add_free_memory(r2.start, r2.len());
                 }
             }
diff --git a/ostd/src/mm/frame/meta.rs b/ostd/src/mm/frame/meta.rs
index 48639b755..e08da8434 100644
--- a/ostd/src/mm/frame/meta.rs
+++ b/ostd/src/mm/frame/meta.rs
@@ -53,7 +53,7 @@ use log::info;
 
 use crate::{
     arch::mm::PagingConsts,
-    const_assert,
+    const_assert, if_tdx_enabled,
     mm::{
         frame::allocator::{self, EarlyAllocatedFrameMeta},
         kspace::LINEAR_MAPPING_BASE_VADDR,
@@ -447,10 +447,15 @@ pub(crate) unsafe fn init() -> Segment<MetaPageMeta> {
         regions.iter().map(|r| r.base() + r.len()).max().unwrap()
     };
 
-    info!(
-        "Initializing frame metadata for physical memory up to {:x}",
-        max_paddr
-    );
+    if_tdx_enabled!({
+        use tdx_guest::serial_println;
+
+        serial_println!("Initializing frame metadata for physical memory up to {:x}",
+            max_paddr)
+    } else {
+        info!("Initializing frame metadata for physical memory up to {:x}",
+            max_paddr);
+    });
 
     add_temp_linear_mapping(max_paddr);
 
diff --git a/ostd/src/mm/kspace/mod.rs b/ostd/src/mm/kspace/mod.rs
index 40875b495..be6c58c36 100644
--- a/ostd/src/mm/kspace/mod.rs
+++ b/ostd/src/mm/kspace/mod.rs
@@ -59,6 +59,7 @@ use super::{
 use crate::{
     arch::mm::{PageTableEntry, PagingConsts},
     boot::memory_region::MemoryRegionType,
+    if_tdx_enabled,
 };
 
 /// The shortest supported address width is 39 bits. And the literal
@@ -132,7 +133,13 @@ pub static KERNEL_PAGE_TABLE: Once<PageTable<KernelMode, PageTableEntry, PagingC
 /// This function should be called before:
 ///  - any initializer that modifies the kernel page table.
 pub fn init_kernel_page_table(meta_pages: Segment<MetaPageMeta>) {
-    info!("Initializing the kernel page table");
+    if_tdx_enabled!({
+        use tdx_guest::serial_println;
+
+        serial_println!("Initializing the kernel page table");
+    } else {
+        info!("Initializing the kernel page table");
+    });
 
     let regions = &crate::boot::EARLY_INFO.get().unwrap().memory_regions;
     let phys_mem_cap = regions.iter().map(|r| r.base() + r.len()).max().unwrap();