From 146e85c050905fb28b1f74ac6519157f82ba1b02 Mon Sep 17 00:00:00 2001
From: Yifan Sun <sunyifan112358@gmail.com>
Date: Tue, 3 Dec 2024 00:35:26 +0000
Subject: [PATCH] Export Device's PageTable

---
 driver/driver.go                   |  2 +-
 driver/internal/device.go          |  2 +-
 driver/internal/memoryallocator.go |  8 ++++----
 go.work                            |  3 ---
 go.work.sum                        | 10 ----------
 5 files changed, 6 insertions(+), 19 deletions(-)
 delete mode 100644 go.work
 delete mode 100644 go.work.sum

diff --git a/driver/driver.go b/driver/driver.go
index 576ae9d9..2cdb3118 100644
--- a/driver/driver.go
+++ b/driver/driver.go
@@ -153,7 +153,7 @@ func (d *Driver) RegisterGPU(
 			CUCount:  properties.CUCount,
 			DRAMSize: properties.DRAMSize,
 		},
-		pageTable: vm.PageTable,
+		PageTable: vm.NewPageTable(d.Log2PageSize),
 	}
 	gpuDevice.SetTotalMemSize(properties.DRAMSize)
 	d.memAllocator.RegisterDevice(gpuDevice)
diff --git a/driver/internal/device.go b/driver/internal/device.go
index 6175b2b1..4033d058 100644
--- a/driver/internal/device.go
+++ b/driver/internal/device.go
@@ -40,7 +40,7 @@ type Device struct {
 	nextActualGPUIndex int
 	MemState           DeviceMemoryState
 	Properties         DeviceProperties
-	pageTable          vm.PageTable
+	PageTable          vm.PageTable
 }
 
 // SetTotalMemSize sets total memory size
diff --git a/driver/internal/memoryallocator.go b/driver/internal/memoryallocator.go
index 68071009..7aaea995 100644
--- a/driver/internal/memoryallocator.go
+++ b/driver/internal/memoryallocator.go
@@ -64,7 +64,7 @@ func (a *memoryAllocatorImpl) RegisterDevice(device *Device) {
 	state := device.MemState
 	state.setInitialAddress(a.totalStorageByteSize)
 	a.totalStorageByteSize += state.getStorageSize()
-	device.pageTable = vm.NewPageTable(a.log2PageSize)
+	device.PageTable = vm.NewPageTable(a.log2PageSize)
 	a.devices[device.ID] = device
 }
 
@@ -162,7 +162,7 @@ func (a *memoryAllocatorImpl) allocatePages(
 
 		// fmt.Printf("page.addr is %x piage Device ID is %d \n", page.PAddr, page.DeviceID)
 		// debug.PrintStack()
-		device.pageTable.Insert(page)
+		device.PageTable.Insert(page)
 		a.vAddrToPageMapping[page.VAddr] = page
 	}
 
@@ -189,7 +189,7 @@ func (a *memoryAllocatorImpl) Remap(
 	device := a.devices[deviceID]
 	pages := a.allocateMultiplePagesWithGivenVAddrs(pid, deviceID, vAddrs, false)
 	for _, page := range pages {
-		device.pageTable.Insert(page)
+		device.PageTable.Insert(page)
 	}
 }
 
@@ -246,7 +246,7 @@ func (a *memoryAllocatorImpl) allocatePageWithGivenVAddr(
 		DeviceID: uint64(deviceID),
 		Unified:  isUnified,
 	}
-	device.pageTable.Insert(page)
+	device.PageTable.Insert(page)
 	a.vAddrToPageMapping[page.VAddr] = page
 
 	return page
diff --git a/go.work b/go.work
deleted file mode 100644
index 01cddcb9..00000000
--- a/go.work
+++ /dev/null
@@ -1,3 +0,0 @@
-go 1.22.4
-
-use ../akita
diff --git a/go.work.sum b/go.work.sum
deleted file mode 100644
index 7488f7e4..00000000
--- a/go.work.sum
+++ /dev/null
@@ -1,10 +0,0 @@
-github.com/chzyer/readline v1.5.1/go.mod h1:Eh+b79XXUwfKfcPLepksvw2tcLE/Ct21YObkaSkeBlk=
-github.com/golang/protobuf v1.5.3/go.mod h1:XVQd3VNwM+JqD3oG2Ue2ip4fOMUkwXdXDdiuN0vRsmY=
-github.com/ianlancetaylor/demangle v0.0.0-20220517205856-0058ec4f073c/go.mod h1:aYm2/VgdVmcIU8iMfdMvDMsRAQjcfZSKFby6HOFvi/w=
-github.com/yuin/goldmark v1.4.13/go.mod h1:6yULJ656Px+3vBD8DxQVa3kxgyrAnzto9xy5taEt/CY=
-golang.org/x/crypto v0.21.0/go.mod h1:0BP7YvVV9gBbVKyeTG0Gyn+gZm94bibOW5BjDEYAOMs=
-golang.org/x/mod v0.10.0/go.mod h1:iBbtSCu2XBx23ZKBPSOrRkjjQPZFPuis4dIYUhu/chs=
-golang.org/x/sync v0.2.0/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
-golang.org/x/term v0.18.0/go.mod h1:ILwASektA3OnRv7amZ1xhE/KTR+u50pbXfZ03+6Nx58=
-google.golang.org/protobuf v1.28.0/go.mod h1:HV8QOd/L58Z+nl8r43ehVNZIU/HEI6OcFqwMG9pJV4I=
-gopkg.in/yaml.v2 v2.4.0/go.mod h1:RDklbk79AGWmwhnvt/jBztapEOGDOx6ZbXqjP6csGnQ=