Rework plugin/PJRT wheel build to use Bazel Build

[lucbruni-amd]

Motivation

Addresses #175
Fixes #178 (the first half of which)

Technical Details

• Use bazel build for plugin/PJRT wheels and copy built artifacts from bazel-bin to --output_path (no bazel run script args).

• Add copy_individual_files and copy_dir_recursively in build/tools/utils.py for copying .whl files and editable directory trees.

• Add WHEEL_BAZEL_DIST_DIR and WHEEL_GLOB_PREFIX in build.py and a copy step that uses them after the build.

• Add WHEEL_BUILD_TARGET_EDITABLE_DICT and make the copy step use copy_dir_recursively for editable and copy_individual_files for normal wheels.

• Split plugin sources into ROCM_PLUGIN_STATIC_SRCS and ROCM_PLUGIN_SOURCES; add build_gpu_kernels_wheel_via_sources (no data) and point the plugin jax_wheel at it.

• In wheel_sources for the plugin, use data_srcs and py_srcs for @jax ROCm targets so the aspect collects .so files; keep static_srcs for non-@jax files.

• Add jax_rocm7_plugin_wheel_editable and jax_rocm7_pjrt_wheel_editable jax_wheel targets.

• In build_gpu_kernels_wheel.py, add _find_patchelf() and use the resolved path when calling patchelf; copy .so from --srcs when present with fallback to runfiles and a clear error

• Set default wheel type to release via --repo_env=ML_WHEEL_TYPE=release (still overridable with --bazel_options=--repo_env=ML_WHEEL_TYPE=snapshot, for example). This addresses the second clause under Motivation above.

Test Plan

Run the build/ci_build script; run build/build.py directly with prefetching separated and combined with the build step; build editable wheels as previous

Test Result

Observe expected wheelhouse/dist output wheels for the given test. Also observe that bazel build is supported allowing for prefetch + build separation.

Submission Checklist

• Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.