[Ironwood] Refine native API host-device benchmark implementation

[linamy85]

• Implement simple baseline using jax.device_put and jax.device_get.
• Ensure 2D data layout (rows, 128) for memory alignment.

To execute:

kl apply -f Ironwood/guides/host_device/tpu7x-host-device-benchmark.yaml

Sample 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.yaml                                                                                                                                               │
│ --- Starting Host-Device Transfer Benchmark (H2D/D2H) ---                                                                                                                                                                                                   │
│ ********************************************************                                                                                                                                                                                                    │
│ WARNING: This benchmark is currently a WORK IN PROGRESS                                                                                                                                                                                                     │
│ ********************************************************                                                                                                                                                                                                    │
│                                                                                                                                                                                                                                                             │
│ Configuration:                                                                                                                                                                                                                                              │
│     Interleaved: false                                                                                                                                                                                                                                      │
│                                                                                                                                                                                                                                                             │
│ --- Running Config: Ironwood/configs/host_device/host_device.yaml ---                                                                                                                                                                                       │
│                                                                                                                                                                                                                                                             │
│ ==============================Starting benchmark 'host_device'==============================                                                                                                                                                                │
│                                                                                                                                                                                                                                                             │
│ Running benchmark: host_device with params: {'data_size_mb': 1, 'num_runs': 20, 'trace_dir': '../microbenchmarks/host_device/trace/benchmark_0'}                                                                                                            │
│ Benchmarking (Simple) Transfer with Data Size: 1 MB on 2 devices for 20 iterations                                                                                                                                                                          │
│ Running benchmark: host_device with params: {'data_size_mb': 16, 'num_runs': 20, 'trace_dir': '../microbenchmarks/host_device/trace/benchmark_1'}                                                                                                           │
│ Benchmarking (Simple) Transfer with Data Size: 16 MB on 2 devices for 20 iterations                                                                                                                                                                         │
│ Running benchmark: host_device with params: {'data_size_mb': 128, 'num_runs': 20, 'trace_dir': '../microbenchmarks/host_device/trace/benchmark_2'}                                                                                                          │
│ Benchmarking (Simple) Transfer with Data Size: 128 MB on 2 devices for 20 iterations                                                                                                                                                                        │
│ Running benchmark: host_device with params: {'data_size_mb': 256, 'num_runs': 20, 'trace_dir': '../microbenchmarks/host_device/trace/benchmark_3'}                                                                                                          │
│ Benchmarking (Simple) Transfer with Data Size: 256 MB on 2 devices for 20 iterations                                                                                                                                                                        │
│ Running benchmark: host_device with params: {'data_size_mb': 512, 'num_runs': 20, 'trace_dir': '../microbenchmarks/host_device/trace/benchmark_4'}                                                                                                          │
│ Benchmarking (Simple) Transfer with Data Size: 512 MB on 2 devices for 20 iterations                                                                                                                                                                        │
│ Running benchmark: host_device with params: {'data_size_mb': 1024, 'num_runs': 20, 'trace_dir': '../microbenchmarks/host_device/trace/benchmark_5'}                                                                                                         │
│ Benchmarking (Simple) Transfer with Data Size: 1024 MB on 2 devices for 20 iterations                                                                                                                                                                       │
│ Running benchmark: host_device with params: {'data_size_mb': 2048, 'num_runs': 20, 'trace_dir': '../microbenchmarks/host_device/trace/benchmark_6'}                                                                                                         │
│ Benchmarking (Simple) Transfer with Data Size: 2048 MB on 2 devices for 20 iterations                                                                                                                                                                       │
│ Running benchmark: host_device with params: {'data_size_mb': 4096, 'num_runs': 20, 'trace_dir': '../microbenchmarks/host_device/trace/benchmark_7'}                                                                                                         │
│ Benchmarking (Simple) Transfer with Data Size: 4096 MB on 2 devices for 20 iterations                                                                                                                                                                       │
│ Running benchmark: host_device with params: {'data_size_mb': 8192, 'num_runs': 20, 'trace_dir': '../microbenchmarks/host_device/trace/benchmark_8'}                                                                                                         │
│ Benchmarking (Simple) Transfer with Data Size: 8192 MB on 2 devices for 20 iterations                                                                                                                                                                       │
│ Running benchmark: host_device with params: {'data_size_mb': 16384, 'num_runs': 20, 'trace_dir': '../microbenchmarks/host_device/trace/benchmark_9'}                                                                                                        │
│ Benchmarking (Simple) Transfer with Data Size: 16384 MB on 2 devices for 20 iterations                                                                                                                                                                      │
│ Running benchmark: host_device with params: {'data_size_mb': 32768, 'num_runs': 20, 'trace_dir': '../microbenchmarks/host_device/trace/benchmark_10'}                                                                                                       │
│ Benchmarking (Simple) Transfer with Data Size: 32768 MB on 2 devices for 20 iterations                                                                                                                                                                      │
│ Metrics written to CSV at ../microbenchmarks/host_device/t_host_device_TDZKVGG9P2.tsv.                                                                                                                                                                      │
│ --- Finished Config: Ironwood/configs/host_device/host_device.yaml ---

[chishuen]

Discussed offline. Overall LGTM. Left some minor comments.