diff --git a/arch/x86/mm/Makefile b/arch/x86/mm/Makefile
index c80febc44cd2..2998eaac0dbb 100644
--- a/arch/x86/mm/Makefile
+++ b/arch/x86/mm/Makefile
@@ -67,3 +67,5 @@ obj-$(CONFIG_AMD_MEM_ENCRYPT)	+= mem_encrypt_amd.o
 
 obj-$(CONFIG_AMD_MEM_ENCRYPT)	+= mem_encrypt_identity.o
 obj-$(CONFIG_AMD_MEM_ENCRYPT)	+= mem_encrypt_boot.o
+
+obj-$(CONFIG_HEKI)		+= heki.o
diff --git a/arch/x86/mm/heki.c b/arch/x86/mm/heki.c
new file mode 100644
index 000000000000..c495df0d8772
--- /dev/null
+++ b/arch/x86/mm/heki.c
@@ -0,0 +1,56 @@
+// SPDX-License-Identifier: GPL-2.0
+/*
+ * Hypervisor Enforced Kernel Integrity (Heki) - Arch specific.
+ *
+ * Copyright © 2023 Microsoft Corporation
+ */
+
+#include <linux/heki.h>
+#include <linux/kvm_mem_attr.h>
+
+#ifdef pr_fmt
+#undef pr_fmt
+#endif
+
+#define pr_fmt(fmt) "heki-guest: " fmt
+
+static unsigned long kernel_va;
+static unsigned long kernel_end;
+static unsigned long direct_map_va;
+static unsigned long direct_map_end;
+
+__init void heki_arch_early_init(void)
+{
+	/* Kernel virtual address space range, not yet compatible with KASLR. */
+	if (pgtable_l5_enabled()) {
+		kernel_va = 0xff00000000000000UL;
+		kernel_end = 0xffffffffffe00000UL;
+		direct_map_va = 0xff11000000000000UL;
+		direct_map_end = 0xff91000000000000UL;
+	} else {
+		kernel_va = 0xffff800000000000UL;
+		kernel_end = 0xffffffffffe00000UL;
+		direct_map_va = 0xffff888000000000UL;
+		direct_map_end = 0xffffc88000000000UL;
+	}
+
+	/*
+	 * Initialize the counters for all existing kernel mappings except
+	 * for direct map.
+	 */
+	heki_map(kernel_va, direct_map_va);
+	heki_map(direct_map_end, kernel_end);
+}
+
+unsigned long heki_flags_to_permissions(unsigned long flags)
+{
+	unsigned long permissions;
+
+	permissions = MEM_ATTR_READ | MEM_ATTR_EXEC;
+	if (flags & _PAGE_RW)
+		permissions |= MEM_ATTR_WRITE;
+	if (flags & _PAGE_NX)
+		permissions &= ~MEM_ATTR_EXEC;
+
+	return permissions;
+}
diff --git a/include/linux/heki.h b/include/linux/heki.h
index a7ae0b387dfe..86c787d121e0 100644
--- a/include/linux/heki.h
+++ b/include/linux/heki.h
@@ -19,6 +19,16 @@
 
 #ifdef CONFIG_HEKI
 
+/*
+ * This structure keeps track of the collective permissions for a guest page
+ * across all of its mappings.
+ */
+struct heki_counters {
+	int read;
+	int write;
+	int execute;
+};
+
 /*
  * This structure contains a guest physical range and its permissions (RWX).
  */
@@ -56,9 +66,17 @@ struct heki_hypervisor {
 /*
  * If the active hypervisor supports Heki, it will plug its heki_hypervisor
  * pointer into this heki structure.
+ *
+ * During guest kernel boot, permissions counters for each guest page are
+ * initialized based on the page's current permissions.
  */
 struct heki {
 	struct heki_hypervisor *hypervisor;
+	struct mem_table *counters;
+};
+
+enum heki_cmd {
+	HEKI_MAP,
 };
 
 /*
@@ -72,6 +90,9 @@ struct heki_args {
 	phys_addr_t pa;
 	size_t size;
 	unsigned long flags;
+
+	/* Command passed by caller. */
+	enum heki_cmd cmd;
 };
 
 /* Callback function called by the table walker. */
@@ -84,6 +105,14 @@ extern bool __read_mostly enable_mbec;
 
 void heki_early_init(void);
 void heki_late_init(void);
+void heki_counters_init(void);
+void heki_walk(unsigned long va, unsigned long va_end, heki_func_t func,
+	       struct heki_args *args);
+void heki_map(unsigned long va, unsigned long end);
+
+/* Arch-specific functions. */
+void heki_arch_early_init(void);
+unsigned long heki_flags_to_permissions(unsigned long flags);
 
 #else /* !CONFIG_HEKI */
 
@@ -93,6 +122,9 @@ static inline void heki_early_init(void)
 static inline void heki_late_init(void)
 {
 }
+static inline void heki_map(unsigned long va, unsigned long end)
+{
+}
 
 #endif /* CONFIG_HEKI */
 
diff --git a/virt/heki/Kconfig b/virt/heki/Kconfig
index 75a784653e31..6d956eb9d04b 100644
--- a/virt/heki/Kconfig
+++ b/virt/heki/Kconfig
@@ -6,6 +6,8 @@ config HEKI
 	bool "Hypervisor Enforced Kernel Integrity (Heki)"
 	depends on ARCH_SUPPORTS_HEKI && HYPERVISOR_SUPPORTS_HEKI
 	select KVM_GENERIC_MEMORY_ATTRIBUTES
+	depends on !X86_16BIT
+	select SPARSEMEM
 	help
 	  This feature enhances guest virtual machine security by taking
 	  advantage of security features provided by the hypervisor for guests.
diff --git a/virt/heki/Makefile b/virt/heki/Makefile
index a5daa4ff7a4f..564f92faa9d8 100644
--- a/virt/heki/Makefile
+++ b/virt/heki/Makefile
@@ -2,3 +2,4 @@
 
 obj-y += main.o
 obj-y += walk.o
+obj-y += counters.o
diff --git a/virt/heki/counters.c b/virt/heki/counters.c
new file mode 100644
index 000000000000..7067449cabca
--- /dev/null
+++ b/virt/heki/counters.c
@@ -0,0 +1,147 @@
+// SPDX-License-Identifier: GPL-2.0-only
+/*
+ * Hypervisor Enforced Kernel Integrity (Heki) - Permissions counters.
+ *
+ * Copyright © 2023 Microsoft Corporation
+ */
+
+#include <linux/heki.h>
+#include <linux/kvm_mem_attr.h>
+#include <linux/mem_table.h>
+
+#include "common.h"
+
+DEFINE_MUTEX(heki_lock);
+
+static void heki_update_counters(struct heki_counters *counters,
+				 unsigned long perm, unsigned long set,
+				 unsigned long clear)
+{
+	if (WARN_ON_ONCE(!counters))
+		return;
+
+	if ((clear & MEM_ATTR_READ) && (perm & MEM_ATTR_READ))
+		counters->read--;
+	if ((clear & MEM_ATTR_WRITE) && (perm & MEM_ATTR_WRITE))
+		counters->write--;
+	if ((clear & MEM_ATTR_EXEC) && (perm & MEM_ATTR_EXEC))
+		counters->execute--;
+
+	if ((set & MEM_ATTR_READ) && !(perm & MEM_ATTR_READ))
+		counters->read++;
+	if ((set & MEM_ATTR_WRITE) && !(perm & MEM_ATTR_WRITE))
+		counters->write++;
+	if ((set & MEM_ATTR_EXEC) && !(perm & MEM_ATTR_EXEC))
+		counters->execute++;
+}
+
+static struct heki_counters *heki_create_counters(struct mem_table *table,
+						  phys_addr_t pa)
+{
+	struct heki_counters *counters;
+	void **entry;
+
+	entry = mem_table_create(table, pa);
+	if (WARN_ON(!entry))
+		return NULL;
+
+	counters = kzalloc(sizeof(*counters), GFP_KERNEL);
+	if (WARN_ON(!counters))
+		return NULL;
+
+	*entry = counters;
+	return counters;
+}
+
+void heki_callback(struct heki_args *args)
+{
+	/* The VA is only for debug. It is not really used in this function. */
+	unsigned long va;
+	phys_addr_t pa, pa_end;
+	unsigned long permissions;
+	void **entry;
+	struct heki_counters *counters;
+	unsigned int ignore;
+
+	if (!pfn_valid(args->pa >> PAGE_SHIFT))
+		return;
+
+	permissions = heki_flags_to_permissions(args->flags);
+
+	/*
+	 * Handle counters for a leaf entry in the kernel page table.
+	 */
+	pa_end = args->pa + args->size;
+	for (pa = args->pa, va = args->va; pa < pa_end;
+	     pa += PAGE_SIZE, va += PAGE_SIZE) {
+		entry = mem_table_find(heki.counters, pa, &ignore);
+		if (entry)
+			counters = *entry;
+		else
+			counters = NULL;
+
+		switch (args->cmd) {
+		case HEKI_MAP:
+			if (!counters)
+				counters =
+					heki_create_counters(heki.counters, pa);
+			heki_update_counters(counters, 0, permissions, 0);
+			break;
+
+		default:
+			WARN_ON_ONCE(1);
+			break;
+		}
+	}
+}
+
+static void heki_func(unsigned long va, unsigned long end,
+		      struct heki_args *args)
+{
+	if (!heki.counters || va >= end)
+		return;
+
+	va = ALIGN_DOWN(va, PAGE_SIZE);
+	end = ALIGN(end, PAGE_SIZE);
+
+	mutex_lock(&heki_lock);
+
+	heki_walk(va, end, heki_callback, args);
+
+	mutex_unlock(&heki_lock);
+}
+
+/*
+ * Find the mappings in the given range and initialize permission counters for
+ * them.
+ */
+void heki_map(unsigned long va, unsigned long end)
+{
+	struct heki_args args = {
+		.cmd = HEKI_MAP,
+	};
+
+	heki_func(va, end, &args);
+}
+
+/*
+ * Permissions counters are associated with each guest page using the
+ * Memory Table feature. Initialize the permissions counters here.
+ * Note that we don't support large page entries for counters because
+ * it is difficult to merge/split counters for large pages.
+ */
+
+static void heki_counters_free(void *counters)
+{
+	kfree(counters);
+}
+
+static struct mem_table_ops heki_counters_ops = {
+	.free = heki_counters_free,
+};
+
+__init void heki_counters_init(void)
+{
+	heki.counters = mem_table_alloc(&heki_counters_ops);
+	WARN_ON(!heki.counters);
+}
diff --git a/virt/heki/main.c b/virt/heki/main.c
index ff1937e1c946..0ab7de659e6f 100644
--- a/virt/heki/main.c
+++ b/virt/heki/main.c
@@ -21,6 +21,16 @@ __init void heki_early_init(void)
 		pr_warn("Heki is not enabled\n");
 		return;
 	}
+
+	/*
+	 * Static addresses (see heki_arch_early_init) are not compatible with
+	 * KASLR. This will be handled in a next patch series.
+	 */
+	if (IS_ENABLED(CONFIG_RANDOMIZE_BASE)) {
+		pr_warn("Heki is disabled because KASLR is not supported yet\n");
+		return;
+	}
+
 	pr_warn("Heki is enabled\n");
 
 	if (!heki.hypervisor) {
@@ -29,6 +39,9 @@ __init void heki_early_init(void)
 		return;
 	}
 	pr_warn("Heki is supported by the active Hypervisor\n");
+
+	heki_counters_init();
+	heki_arch_early_init();
 }
 
 /*
