Fix the bug with reg alloc