diff --git a/pvmfw/Android.bp b/pvmfw/Android.bp
index 287e14a..e5fdfe8 100644
--- a/pvmfw/Android.bp
+++ b/pvmfw/Android.bp
@@ -12,7 +12,6 @@
     flags: ["-Dunsafe_op_in_unsafe_fn"],
     features: [
         "legacy",
-        "cpu_feat_hafdbs",
     ],
     rustlibs: [
         "libaarch64_paging",
diff --git a/pvmfw/src/helpers.rs b/pvmfw/src/helpers.rs
index 1d39e1f..dcfd99c 100644
--- a/pvmfw/src/helpers.rs
+++ b/pvmfw/src/helpers.rs
@@ -152,13 +152,3 @@
         core::ffi::CStr::from_bytes_with_nul(concat!($str, "\0").as_bytes()).unwrap()
     }};
 }
-
-/// Returns `true` if hardware dirty state management is available.
-pub fn dbm_available() -> bool {
-    if !cfg!(feature = "cpu_feat_hafdbs") {
-        return false;
-    }
-    // Hardware dirty bit management available flag (ID_AA64MMFR1_EL1.HAFDBS[1])
-    const DBM_AVAILABLE: usize = 1 << 1;
-    read_sysreg!("id_aa64mmfr1_el1") & DBM_AVAILABLE != 0
-}
diff --git a/pvmfw/src/memory.rs b/pvmfw/src/memory.rs
index c316dd2..3205a4d 100644
--- a/pvmfw/src/memory.rs
+++ b/pvmfw/src/memory.rs
@@ -16,7 +16,7 @@
 
 #![deny(unsafe_op_in_unsafe_fn)]
 
-use crate::helpers::{self, dbm_available, page_4kb_of, RangeExt, PVMFW_PAGE_SIZE, SIZE_4MB};
+use crate::helpers::{self, page_4kb_of, RangeExt, PVMFW_PAGE_SIZE, SIZE_4MB};
 use crate::mmu;
 use aarch64_paging::paging::{Attributes, Descriptor, MemoryRegion as VaRange};
 use alloc::alloc::alloc_zeroed;
@@ -40,7 +40,7 @@
 use once_cell::race::OnceBox;
 use spin::mutex::SpinMutex;
 use tinyvec::ArrayVec;
-use vmbase::{dsb, isb, read_sysreg, tlbi, write_sysreg};
+use vmbase::{dsb, isb, memory::set_dbm_enabled, tlbi};
 
 /// Base of the system's contiguous "main" memory.
 pub const BASE_ADDR: usize = 0x8000_0000;
@@ -219,15 +219,13 @@
     const CAPACITY: usize = 5;
     const MMIO_CAPACITY: usize = 5;
     const PVMFW_RANGE: MemoryRange = (BASE_ADDR - SIZE_4MB)..BASE_ADDR;
-    // TCR_EL1.{HA,HD} bits controlling hardware management of access and dirty state
-    const TCR_EL1_HA_HD_BITS: usize = 3 << 39;
 
     /// Create a new instance from an active page table, covering the maximum RAM size.
     pub fn new(mut page_table: mmu::PageTable) -> Self {
         // Activate dirty state management first, otherwise we may get permission faults immediately
         // after activating the new page table. This has no effect before the new page table is
         // activated because none of the entries in the initial idmap have the DBM flag.
-        Self::set_dbm_enabled(true);
+        set_dbm_enabled(true);
 
         debug!("Activating dynamic page table...");
         // SAFETY - page_table duplicates the static mappings for everything that the Rust code is
@@ -434,25 +432,11 @@
             .modify_range(&(addr..addr + 1), &mark_dirty_block)
             .map_err(|_| MemoryTrackerError::SetPteDirtyFailed)
     }
-
-    fn set_dbm_enabled(enabled: bool) {
-        if dbm_available() {
-            let mut tcr = read_sysreg!("tcr_el1");
-            if enabled {
-                tcr |= Self::TCR_EL1_HA_HD_BITS
-            } else {
-                tcr &= !Self::TCR_EL1_HA_HD_BITS
-            };
-            // Safe because it writes to a system register and does not affect Rust.
-            unsafe { write_sysreg!("tcr_el1", tcr) }
-            isb!();
-        }
-    }
 }
 
 impl Drop for MemoryTracker {
     fn drop(&mut self) {
-        Self::set_dbm_enabled(false);
+        set_dbm_enabled(false);
         self.flush_dirty_pages().unwrap();
         self.unshare_all_memory();
     }
diff --git a/vmbase/Android.bp b/vmbase/Android.bp
index 5ed436c..237a377 100644
--- a/vmbase/Android.bp
+++ b/vmbase/Android.bp
@@ -70,6 +70,9 @@
     whole_static_libs: [
         "librust_baremetal",
     ],
+    features: [
+        "cpu_feat_hafdbs",
+    ],
     apex_available: ["com.android.virt"],
 }
 
diff --git a/vmbase/src/lib.rs b/vmbase/src/lib.rs
index 0eca1e4..d086f1c 100644
--- a/vmbase/src/lib.rs
+++ b/vmbase/src/lib.rs
@@ -23,6 +23,7 @@
 pub mod layout;
 mod linker;
 pub mod logger;
+pub mod memory;
 pub mod power;
 pub mod uart;
 
diff --git a/vmbase/src/memory/dbm.rs b/vmbase/src/memory/dbm.rs
new file mode 100644
index 0000000..3a52c2d
--- /dev/null
+++ b/vmbase/src/memory/dbm.rs
@@ -0,0 +1,47 @@
+// Copyright 2023, 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.
+
+//! Hardware management of the access flag and dirty state.
+
+use crate::{isb, read_sysreg, write_sysreg};
+
+/// Sets whether the hardware management of access and dirty state is enabled with
+/// the given boolean.
+pub fn set_dbm_enabled(enabled: bool) {
+    if !dbm_available() {
+        return;
+    }
+    // TCR_EL1.{HA,HD} bits controlling hardware management of access and dirty state
+    const TCR_EL1_HA_HD_BITS: usize = 3 << 39;
+
+    let mut tcr = read_sysreg!("tcr_el1");
+    if enabled {
+        tcr |= TCR_EL1_HA_HD_BITS
+    } else {
+        tcr &= !TCR_EL1_HA_HD_BITS
+    };
+    // Safe because it writes to a system register and does not affect Rust.
+    unsafe { write_sysreg!("tcr_el1", tcr) }
+    isb!();
+}
+
+/// Returns `true` if hardware dirty state management is available.
+fn dbm_available() -> bool {
+    if !cfg!(feature = "cpu_feat_hafdbs") {
+        return false;
+    }
+    // Hardware dirty bit management available flag (ID_AA64MMFR1_EL1.HAFDBS[1])
+    const DBM_AVAILABLE: usize = 1 << 1;
+    read_sysreg!("id_aa64mmfr1_el1") & DBM_AVAILABLE != 0
+}
diff --git a/vmbase/src/memory/mod.rs b/vmbase/src/memory/mod.rs
new file mode 100644
index 0000000..c4990f2
--- /dev/null
+++ b/vmbase/src/memory/mod.rs
@@ -0,0 +1,19 @@
+// Copyright 2023, 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.
+
+//! Memory management.
+
+mod dbm;
+
+pub use dbm::set_dbm_enabled;