Enable Separate GMMU Page Table Support

driver/driver.go

@@ -143,6 +143,7 @@ type DeviceProperties struct {
 func (d *Driver) RegisterGPU(
 	commandProcessorPort sim.Port,
 	properties DeviceProperties,
+	gmmuPageTable vm.PageTable,
 ) {
 	d.GPUs = append(d.GPUs, commandProcessorPort)
 
@@ -157,8 +158,13 @@ func (d *Driver) RegisterGPU(
 	}
 	gpuDevice.SetTotalMemSize(properties.DRAMSize)
 	d.memAllocator.RegisterDevice(gpuDevice)
-
 	d.devices = append(d.devices, gpuDevice)
+
+	for _, page := range d.memAllocator.GetVAddrToPageMapping() {
+		if page.DeviceID == uint64(gpuDevice.ID) {
+			gmmuPageTable.Insert(page)
+		}
+	}
 }
 
 // Tick ticks

driver/internal/memoryallocator.go

@@ -22,6 +22,8 @@ type MemoryAllocator interface {
 		vAddr uint64,
 		unified bool,
 	) vm.Page
+
+	GetVAddrToPageMapping() map[uint64]vm.Page
 }
 
 // NewMemoryAllocator creates a new memory allocator.
@@ -285,3 +287,13 @@ func (a *memoryAllocatorImpl) Free(ptr uint64) {
 
 	a.removePage(ptr)
 }
+
+func (a *memoryAllocatorImpl) GetVAddrToPageMapping() map[uint64]vm.Page {
+	a.Lock()
+	defer a.Unlock()
+	copy := make(map[uint64]vm.Page, len(a.vAddrToPageMapping))
+	for vAddr, page := range a.vAddrToPageMapping {
+		copy[vAddr] = page
+	}
+	return copy
+}