diff --git a/dmc-lib/src/ir/mod.rs b/dmc-lib/src/ir/mod.rs
index bcf1f86..6f8cf84 100644
--- a/dmc-lib/src/ir/mod.rs
+++ b/dmc-lib/src/ir/mod.rs
@@ -35,9 +35,10 @@ mod variable_locations;
 
 pub fn compile_dvm_function(
     ir_function: &IrFunction,
+    register_count: usize,
     constants_table: &mut ConstantsTable,
 ) -> Function {
-    let variable_locations = assign_registers(ir_function);
+    let variable_locations = assign_registers(ir_function, register_count);
     let instructions = assemble_ir_function(ir_function, &variable_locations, constants_table);
     Function::new(
         ir_function.fqn().into(),
@@ -47,7 +48,7 @@ pub fn compile_dvm_function(
     )
 }
 
-fn assign_registers(ir_function: &IrFunction) -> VariableLocations {
+fn assign_registers(ir_function: &IrFunction, register_count: usize) -> VariableLocations {
     if ir_function.blocks().is_empty() {
         return VariableLocations::new(HashMap::new(), HashMap::new());
     }
@@ -55,5 +56,5 @@ fn assign_registers(ir_function: &IrFunction) -> VariableLocations {
         unimplemented!("having more than one block in a function is not yet implemented.")
     }
     let block = &ir_function.blocks()[0];
-    block_assign_registers(block, 16) // TODO: make this configurable
+    block_assign_registers(block, register_count)
 }
diff --git a/dmc-lib/src/lib.rs b/dmc-lib/src/lib.rs
index c4714f1..57abfa6 100644
--- a/dmc-lib/src/lib.rs
+++ b/dmc-lib/src/lib.rs
@@ -39,6 +39,7 @@ pub struct CompileCompilationUnitResult {
 
 pub fn compile_compilation_unit(
     input: &str,
+    register_count: usize,
     constants_table: &mut ConstantsTable,
 ) -> Result<CompileCompilationUnitResult, Diagnostics> {
     let (compilation_unit, diagnostics) = parse_compilation_unit(input, None);
@@ -63,7 +64,7 @@ pub fn compile_compilation_unit(
     let mut dvm_functions = HashMap::new();
 
     for ir_function in &lower_to_ir_result.functions {
-        let dvm_function = compile_dvm_function(ir_function, constants_table);
+        let dvm_function = compile_dvm_function(ir_function, register_count, constants_table);
         dvm_functions.insert(dvm_function.name_owned(), dvm_function);
     }
 
diff --git a/e2e-tests/src/lib.rs b/e2e-tests/src/lib.rs
index c43843b..cb79b44 100644
--- a/e2e-tests/src/lib.rs
+++ b/e2e-tests/src/lib.rs
@@ -25,7 +25,8 @@ mod e2e_tests {
     fn prepare_context(input: &str) -> Result<DvmContext, Diagnostics> {
         let mut constants_table = ConstantsTable::new();
 
-        let compile_compilation_unit_result = compile_compilation_unit(input, &mut constants_table);
+        let compile_compilation_unit_result =
+            compile_compilation_unit(input, REGISTER_COUNT, &mut constants_table);
 
         let mut dvm_context = DvmContext::new();