()
| 515 | |
| 516 | |
| 517 | def build_shmem(): |
| 518 | # Determine which shmem backend(s) to build based on TRITON_DIST_SHMEM_BACKEND |
| 519 | shmem_backend = os.getenv("TRITON_DIST_SHMEM_BACKEND", "").lower() |
| 520 | |
| 521 | if not _is_hip_platform(): |
| 522 | # Not HIP platform, skip shmem build |
| 523 | return |
| 524 | |
| 525 | # If no backend specified, build both |
| 526 | if not shmem_backend: |
| 527 | print("TRITON_DIST_SHMEM_BACKEND not set, building both mori_shmem and rocshmem") |
| 528 | build_mori_shmem() |
| 529 | build_rocshmem() |
| 530 | elif shmem_backend == "mori_shmem": |
| 531 | print("Building mori_shmem backend") |
| 532 | build_mori_shmem() |
| 533 | elif shmem_backend == "rocshmem": |
| 534 | print("Building rocshmem backend") |
| 535 | build_rocshmem() |
| 536 | else: |
| 537 | raise RuntimeError(f"Unknown TRITON_DIST_SHMEM_BACKEND: {shmem_backend}. Must be 'mori_shmem' or 'rocshmem'") |
| 538 | |
| 539 | # Also build if explicitly requested via env var (for backward compatibility) |
| 540 | if check_env_flag("TRITON_DISTRIBUTED_BUILD_PYROCSHMEM", "0") and shmem_backend != "rocshmem": |
| 541 | build_rocshmem() # (9, 4) |
| 542 | |
| 543 | |
| 544 | class SHMEMBuildOnly(Command): |
no test coverage detected