Support a Separate GMMU Page Table

.vscode/launch.json

@@ -166,7 +166,7 @@
             "program": "${workspaceFolder}/samples/fir",
             "args": [
                 "-timing",
-                "-length=8192",
+                "-length=81920",
                 "-report-all",
             ],
         },

driver/builder.go

@@ -17,6 +17,7 @@ type Builder struct {
 	useMagicMemoryCopy  bool
 	middlewareD2HCycles int
 	middlewareH2DCycles int
+	memorySize          uint64
 }
 
 // MakeBuilder creates a driver builder with some default configuration
@@ -74,6 +75,12 @@ func (b Builder) WithH2DCycles(h2dCycles int) Builder {
 	return b
 }
 
+// WithMemorySize sets the memory size of the CPU.
+func (b Builder) WithMemorySize(memorySize uint64) Builder {
+	b.memorySize = memorySize
+	return b
+}
+
 // Build creates a driver.
 func (b Builder) Build(name string) *Driver {
 	driver := new(Driver)
@@ -125,7 +132,7 @@ func (b *Builder) createCPU(d *Driver) {
 		Type:     internal.DeviceTypeCPU,
 		MemState: internal.NewDeviceMemoryState(d.Log2PageSize),
 	}
-	cpu.SetTotalMemSize(4 * mem.GB)
+	cpu.SetTotalMemSize(b.memorySize)
 
 	d.memAllocator.RegisterDevice(cpu)
 	d.devices = append(d.devices, cpu)

driver/driver.go

@@ -1,6 +1,7 @@
 package driver
 
 import (
+	"fmt"
 	"log"
 	"reflect"
 	"runtime/debug"
@@ -142,6 +143,7 @@ type DeviceProperties struct {
 func (d *Driver) RegisterGPU(
 	commandProcessorPort sim.Port,
 	properties DeviceProperties,
+	gmmuPageTable vm.PageTable,
 ) {
 	d.GPUs = append(d.GPUs, commandProcessorPort)
 
@@ -156,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
@@ -455,6 +462,9 @@ func (d *Driver) distributeWGToGPUs(
 		panic("not all wg allocated")
 	}
 
+	// fmt.Sprintln("total WG: %d WG Per CU %d\n", totalWGCount, wgPerCU)
+	fmt.Printf("total WG: %d WG Per CU %d\n", totalWGCount, wgPerCU)
+
 	return wgDist
 }
 

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
+}

timing/cp/builder.go

@@ -86,9 +86,9 @@ func (b Builder) Build(name string) *CommandProcessor {
 		tracing.CollectTrace(cp, b.visTracer)
 	}
 
-	if b.perfAnalyzer != nil {
-		b.perfAnalyzer.RegisterComponent(cp)
-	}
+	// if b.perfAnalyzer != nil {
+	// 	b.perfAnalyzer.RegisterComponent(cp)
+	// }
 
 	return cp
 }