diff --git a/gateware/main.py b/gateware/main.py
index ea2b277..2c65742 100644
--- a/gateware/main.py
+++ b/gateware/main.py
@@ -122,7 +122,7 @@ class Core(Elaboratable):
         fw = load_firmware_for_mem()
 
         # Hook up memory space
-        self.rom = ROM(fw)
+        self.rom = ROM(size_bytes=8192, data=fw)
         m.submodules.rom = self.rom
         # Problem: not sure to handle how we do byte vs word addressing properly
         # So doing this shift is a bit of a hacky way to impl anything
diff --git a/gateware/memory.py b/gateware/memory.py
index 915b7fc..5bde3ea 100644
--- a/gateware/memory.py
+++ b/gateware/memory.py
@@ -2,25 +2,28 @@ from amaranth import *
 from amaranth_soc.wishbone import *
 from amaranth_soc.memory import *
 
+from math import log2, ceil
+
 # TODO impl select
 
 
 # We sub-class wishbone.Interface here because it needs to be a bus object to be added as a window to Wishbone stuff
 class ROM(Elaboratable, Interface):
-    def __init__(self, data=None):
+    def __init__(self, *, size_bytes=4096, data=None):
         #self.size = len(data)
-        self.data = Memory(width=32, depth=(4096 >> 2), init=data)
+        addr_width = ceil(log2(size_bytes >> 2))
+        self.data = Memory(width=32, depth=size_bytes >> 2, init=data)
         self.r = self.data.read_port()
 
         # Need to init Interface
-        Interface.__init__(self, addr_width=10, data_width=32, granularity=8)
+        Interface.__init__(self, addr_width=addr_width, data_width=32, granularity=8)
 
         # This is effectively a "window", and it has a certain set of resources
         # 12 = log2(4096)
-        memory_map = MemoryMap(addr_width=12, data_width=8)
+        memory_map = MemoryMap(addr_width=addr_width + 2, data_width=8)
         # TODO need to unify how I deal with size
         # In this case, one resource, which is out memory
-        memory_map.add_resource(self.data, name="rom_data", size=(4096 >> 2))
+        memory_map.add_resource(self.data, name="rom_data", size=size_bytes)
 
         self.memory_map = memory_map
 
@@ -52,21 +55,21 @@ class ROM(Elaboratable, Interface):
 # TODO support read segmentation or whatever it's called, where you read/write certian bytes from memory
 # Otherwise we can't store individual bytes, and this will wreck shit in weird ways.
 class RAM(Elaboratable, Interface):
-    def __init__(self):
-        #self.size = len(data)
-        self.data = Memory(width=32, depth=(4096 >> 2))
+    def __init__(self, *, size_bytes=4096):
+        addr_width = ceil(log2(size_bytes >> 2))
+        self.data = Memory(width=32, depth=size_bytes >> 2)
         self.r = self.data.read_port()
         self.w = self.data.write_port()
 
         # Need to init Interface
-        Interface.__init__(self, addr_width=10, data_width=32, granularity=8)
+        Interface.__init__(self, addr_width=addr_width, data_width=32, granularity=8)
 
         # This is effectively a "window", and it has a certain set of resources
         # 12 = log2(4096)
-        memory_map = MemoryMap(addr_width=12, data_width=8)
+        memory_map = MemoryMap(addr_width=addr_width + 2, data_width=8)
         # TODO need to unify how I deal with size
         # In this case, one resource, which is out memory
-        memory_map.add_resource(self.data, name="ram_data", size=(4096 >> 2))
+        memory_map.add_resource(self.data, name="ram_data", size=size_bytes)
 
         self.memory_map = memory_map