Improve setup for register coloring