Implement host-device baseline using device_put/get

[linamy85]

Baseline benchmarking for host-device transfer. This benchmark work is WIP so the implementation/result may change in the future.

Example output:

+ export TPU_VISIBLE_CHIPS=0
+ TPU_VISIBLE_CHIPS=0
+ bash ./Ironwood/scripts/run_host_device_benchmark.sh --config Ironwood/configs/host_device/host_device_single_chip.yaml
--- Starting Host-Device Transfer Benchmark (H2D/D2H) ---
*****************************************************
*  WARNING: THIS BENCHMARK IS A WORK IN PROGRESS    *
*  Results may be unstable or subject to change.    *
*****************************************************
Interleaved: false
--- Running Config: Ironwood/configs/host_device/host_device_single_chip.yaml ---

==============================Starting benchmark 'host_device'==============================

Running benchmark: host_device with params: {'num_devices': 2, 'smart_chunking': False, 'data_size_mb': 1, 'num_runs': 20, 'trace_dir': '../microbenchmarks/host_device/single_chip/trace/benchmark_0'}
tpu_devices [TpuDevice(id=0, process_index=0, coords=(0,0,0), core_on_chip=0), TpuDevice(id=1, process_index=0, coords=(0,0,0), core_on_chip=1)]
Simple Benchmark: 1 MB on 2 devices, 20 runs
Running benchmark: host_device with params: {'num_devices': 2, 'smart_chunking': False, 'data_size_mb': 16, 'num_runs': 20, 'trace_dir': '../microbenchmarks/host_device/single_chip/trace/benchmark_1'}
tpu_devices [TpuDevice(id=0, process_index=0, coords=(0,0,0), core_on_chip=0), TpuDevice(id=1, process_index=0, coords=(0,0,0), core_on_chip=1)]
Simple Benchmark: 16 MB on 2 devices, 20 runs
Running benchmark: host_device with params: {'num_devices': 2, 'smart_chunking': False, 'data_size_mb': 128, 'num_runs': 20, 'trace_dir': '../microbenchmarks/host_device/single_chip/trace/benchmark_2'}
tpu_devices [TpuDevice(id=0, process_index=0, coords=(0,0,0), core_on_chip=0), TpuDevice(id=1, process_index=0, coords=(0,0,0), core_on_chip=1)]
Simple Benchmark: 128 MB on 2 devices, 20 runs
Running benchmark: host_device with params: {'num_devices': 2, 'smart_chunking': False, 'data_size_mb': 256, 'num_runs': 20, 'trace_dir': '../microbenchmarks/host_device/single_chip/trace/benchmark_3'}
tpu_devices [TpuDevice(id=0, process_index=0, coords=(0,0,0), core_on_chip=0), TpuDevice(id=1, process_index=0, coords=(0,0,0), core_on_chip=1)]
Simple Benchmark: 256 MB on 2 devices, 20 runs
Running benchmark: host_device with params: {'num_devices': 2, 'smart_chunking': False, 'data_size_mb': 512, 'num_runs': 20, 'trace_dir': '../microbenchmarks/host_device/single_chip/trace/benchmark_4'}
tpu_devices [TpuDevice(id=0, process_index=0, coords=(0,0,0), core_on_chip=0), TpuDevice(id=1, process_index=0, coords=(0,0,0), core_on_chip=1)]
Simple Benchmark: 512 MB on 2 devices, 20 runs
Running benchmark: host_device with params: {'num_devices': 2, 'smart_chunking': False, 'data_size_mb': 1024, 'num_runs': 20, 'trace_dir': '../microbenchmarks/host_device/single_chip/trace/benchmark_5'}
tpu_devices [TpuDevice(id=0, process_index=0, coords=(0,0,0), core_on_chip=0), TpuDevice(id=1, process_index=0, coords=(0,0,0), core_on_chip=1)]
Simple Benchmark: 1024 MB on 2 devices, 20 runs
Running benchmark: host_device with params: {'num_devices': 2, 'smart_chunking': False, 'data_size_mb': 2048, 'num_runs': 20, 'trace_dir': '../microbenchmarks/host_device/single_chip/trace/benchmark_6'}
tpu_devices [TpuDevice(id=0, process_index=0, coords=(0,0,0), core_on_chip=0), TpuDevice(id=1, process_index=0, coords=(0,0,0), core_on_chip=1)]
Simple Benchmark: 2048 MB on 2 devices, 20 runs
Running benchmark: host_device with params: {'num_devices': 2, 'smart_chunking': False, 'data_size_mb': 4096, 'num_runs': 20, 'trace_dir': '../microbenchmarks/host_device/single_chip/trace/benchmark_7'}
tpu_devices [TpuDevice(id=0, process_index=0, coords=(0,0,0), core_on_chip=0), TpuDevice(id=1, process_index=0, coords=(0,0,0), core_on_chip=1)]
Simple Benchmark: 4096 MB on 2 devices, 20 runs
Running benchmark: host_device with params: {'num_devices': 2, 'smart_chunking': False, 'data_size_mb': 8192, 'num_runs': 20, 'trace_dir': '../microbenchmarks/host_device/single_chip/trace/benchmark_8'}
tpu_devices [TpuDevice(id=0, process_index=0, coords=(0,0,0), core_on_chip=0), TpuDevice(id=1, process_index=0, coords=(0,0,0), core_on_chip=1)]
Simple Benchmark: 8192 MB on 2 devices, 20 runs
Running benchmark: host_device with params: {'num_devices': 2, 'smart_chunking': False, 'data_size_mb': 16384, 'num_runs': 20, 'trace_dir': '../microbenchmarks/host_device/single_chip/trace/benchmark_9'}
tpu_devices [TpuDevice(id=0, process_index=0, coords=(0,0,0), core_on_chip=0), TpuDevice(id=1, process_index=0, coords=(0,0,0), core_on_chip=1)]
Simple Benchmark: 16384 MB on 2 devices, 20 runs
Running benchmark: host_device with params: {'num_devices': 2, 'smart_chunking': False, 'data_size_mb': 32768, 'num_runs': 20, 'trace_dir': '../microbenchmarks/host_device/single_chip/trace/benchmark_10'}
tpu_devices [TpuDevice(id=0, process_index=0, coords=(0,0,0), core_on_chip=0), TpuDevice(id=1, process_index=0, coords=(0,0,0), core_on_chip=1)]
Simple Benchmark: 32768 MB on 2 devices, 20 runs
Metrics written to CSV at ../microbenchmarks/host_device/single_chip/t_host_device_MV1MSGY5MO.tsv.
--- Finished Config: Ironwood/configs/host_device/host_device_single_chip.yaml ---

--- All Benchmarks Finished ---

[chishuen]

Could you add a short description of what this PR enables? Thanks

[hylin2002]

LGTM