Call DropAllControlDeps before removing an instruction from the computation during RematerializeInstructions, or RemoveInstruction will fail.
Call DropAllControlDeps before removing an instruction from the computation during RematerializeInstructions, or RemoveInstruction will fail.

Before this change, hlo-opt was failing with this stack:
INTERNAL: RET_CHECK failure (third_party/tensorflow/compiler/xla/hlo/ir/hlo_computation.cc:426) ignore_safety_check || IsSafelyRemovable(instruction) cannot remove instruction: %all-reduce-start.285 = (s32[8,64]{1,0}, s32[], s32[], s32[], s32[], /*index=5*/s32[], s32[], s32[], s32[], s32[], /*index=10*/s32[], s32[], s32[], s32[], s32[], /*index=15*/s32[], s32[], s32[], s32[], s32[]) all-reduce-start(s32[8,64]{1,0} %input_scatter_fusion.16, s32[] %convert.81230.0.remat, s32[] %copy.13646, s32[] %copy.13647, s32[] %copy.13648, /*index=5*/s32[] %copy.13649, s32[] %copy.13650, s32[] %copy.13651, s32[] %copy.13652, s32[] %copy.13653, /*index=10*/s32[] %copy.13654, s32[] %copy.13655, s32[] %copy.13656, s32[] %copy.13657, s32[] %copy.13658, /*index=15*/s32[] %copy.13659, s32[] %copy.13660, s32[] %copy.13661, s32[] %copy.13662, s32[] %copy.13663), channel_id=711, replica_groups={{0,4,8,12,16,20,24,28,32,36,40,44,48,52,56,60,64,68,72,76,80,84,88,92,96,100,104,108,112,116,120,124},{1,5,9,13,17,21,25,29,33,37,41,45,49,53,57,61,65,69,73,77,81,85,89,93,97,101,105,109,113,117,121,125},{2,6,10,14,18,22,26,30,34,38,42,46,50,54,58,62,66,70,74,78,82,86,90,94,98,102,106,110,114,118,122,126},{3,7,11,15,19,23,27,31,35,39,43,47,51,55,59,63,67,71,75,79,83,87,91,95,99,103,107,111,115,119,123,127}}, use_global_device_ids=true, to_apply=%region_597.10276.clone.1, control-predecessors={%copy.14013, %copy.14014, %copy.14015, %copy.14016, %copy.14017, %copy.14018, %copy.14019, %copy.14020, %copy.14021, %copy.14022, %copy.14023, %copy.14024, %copy.14025, %copy.14026, %copy.14027, %copy.14028, %copy.14029, %copy.14030, %copy.14031, %copy.14032, %copy.14033, %copy.14034, %copy.14035, %copy.14036, %copy.14037, %copy.14038, %copy.14039, %copy.14040, %copy.14041, %copy.14042}, backend_config={"operation_queue_id":"0","wait_on_operation_queues":[],"collective_backend_config":{"is_sync":false,"no_parallel_custom_call":false}}
=== Source Location Trace: === 
third_party/tensorflow/compiler/xla/status_macros.cc:80
third_party/tensorflow/compiler/xla/service/hlo_rematerialization.cc:2165
third_party/tensorflow/compiler/xla/service/hlo_rematerialization.cc:2564
third_party/tensorflow/compiler/xla/service/hlo_rematerialization.cc:2708
third_party/tensorflow/compiler/xla/service/hlo_rematerialization.cc:2767
third_party/tensorflow/compiler/xla/service/hlo_rematerialization.cc:2905
third_party/tensorflow/compiler/xla/service/hlo_pass_pipeline.h:140
third_party/tensorflow/compiler/xla/service/hlo_pass_pipeline.cc:185
third_party/tensorflow/compiler/xla/service/gpu/gpu_compiler.cc:2183
third_party/tensorflow/compiler/xla/service/gpu/gpu_compiler.cc:1923
third_party/tensorflow/compiler/xla/service/gpu/gpu_compiler.cc:2012
third_party/tensorflow/compiler/xla/tools/hlo_opt/opt_lib.cc:155
third_party/tensorflow/compiler/xla/tools/hlo_opt/opt_lib.cc:108
third_party/tensorflow/compiler/xla/tools/hlo_opt/gpu_opt.cc:80
third_party/tensorflow/compiler/xla/tools/hlo_opt/opt_main.cc:166
third_party/tensorflow/compiler/xla/tools/hlo_opt/opt_main.cc:177

After this change, hlo-opt fails with an expected failure due to no JAX:
UNIMPLEMENTED: No registered implementation for custom call to cu_threefry2x32 for platform CUDA
=== Source Location Trace: === 
third_party/tensorflow/compiler/xla/service/gpu/ir_emitter_unnested.cc:1391
third_party/tensorflow/compiler/xla/service/gpu/ir_emitter_unnested.cc:3023
third_party/tensorflow/compiler/xla/service/gpu/ir_emitter_unnested.cc:2574
third_party/tensorflow/compiler/xla/service/gpu/ir_emitter_unnested.cc:2015
third_party/tensorflow/compiler/xla/service/gpu/ir_emitter_unnested.cc:3023
third_party/tensorflow/compiler/xla/service/gpu/compile_module_to_llvm_ir.cc:205
third_party/tensorflow/compiler/xla/service/gpu/gpu_compiler.cc:1949
third_party/tensorflow/compiler/xla/service/gpu/gpu_compiler.cc:2028
third_party/tensorflow/compiler/xla/tools/hlo_opt/opt_lib.cc:155
third_party/tensorflow/compiler/xla/tools/hlo_opt/opt_lib.cc:108
third_party/tensorflow/compiler/xla/tools/hlo_opt/gpu_opt.cc:80
third_party/tensorflow/compiler/xla/tools/hlo_opt/opt_main.cc:166
third_party/tensorflow/compiler/xla/tools/hlo_opt/opt_main.cc:177
