[Feature] Update the memory interface to support unified memory.