From 00ebf5b6d9be9472f57ebbe590dc08e7302dbb0d Mon Sep 17 00:00:00 2001
From: DX990307 <xudaoxuan@ufl.edu>
Date: Fri, 20 Sep 2024 21:14:43 +0000
Subject: [PATCH 1/2] For Wafer-scale GPU Heterogeneous Design Porject

---
 .vscode/launch.json  | 2 +-
 driver/builder.go    | 9 ++++++++-
 driver/driver.go     | 4 ++++
 timing/cp/builder.go | 6 +++---
 4 files changed, 16 insertions(+), 5 deletions(-)

diff --git a/.vscode/launch.json b/.vscode/launch.json
index 83822569..d6f08f5e 100644
--- a/.vscode/launch.json
+++ b/.vscode/launch.json
@@ -166,7 +166,7 @@
             "program": "${workspaceFolder}/samples/fir",
             "args": [
                 "-timing",
-                "-length=8192",
+                "-length=81920",
                 "-report-all",
             ],
         },
diff --git a/driver/builder.go b/driver/builder.go
index 3cf47ed8..62f62060 100644
--- a/driver/builder.go
+++ b/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)
diff --git a/driver/driver.go b/driver/driver.go
index 28d17459..bb882d9c 100644
--- a/driver/driver.go
+++ b/driver/driver.go
@@ -1,6 +1,7 @@
 package driver
 
 import (
+	"fmt"
 	"log"
 	"reflect"
 	"runtime/debug"
@@ -455,6 +456,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
 }
 
diff --git a/timing/cp/builder.go b/timing/cp/builder.go
index 797a5bfb..99e65e1e 100644
--- a/timing/cp/builder.go
+++ b/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
 }

From 8da82c35d554276bda54f1199ec6c96dc5060baa Mon Sep 17 00:00:00 2001
From: Yuwei Sun <ys4680@nyu.edu>
Date: Fri, 1 Nov 2024 20:49:02 +0000
Subject: [PATCH 2/2] Enable Separate GMMU Page Table Support

---
 driver/driver.go                   |  8 +++++++-
 driver/internal/memoryallocator.go | 12 ++++++++++++
 2 files changed, 19 insertions(+), 1 deletion(-)

diff --git a/driver/driver.go b/driver/driver.go
index bb882d9c..0dcf0a0f 100644
--- a/driver/driver.go
+++ b/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
diff --git a/driver/internal/memoryallocator.go b/driver/internal/memoryallocator.go
index 1a887ea5..83f4c909 100644
--- a/driver/internal/memoryallocator.go
+++ b/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
+}