Refactor merge_adjacent_partitions to take a Context parameter and to dynamically allocate smem.