Stage files for compilation and execution, returning commands for the CLI.
| 78 | |
| 79 | |
| 80 | class ProblemPackager: |
| 81 | """Stage files for compilation and execution, returning commands for the CLI.""" |
| 82 | |
| 83 | def __init__( |
| 84 | self, |
| 85 | definition: Definition, |
| 86 | workloads: list[Workload], |
| 87 | solution: Solution, |
| 88 | config: BenchmarkConfig, |
| 89 | output_dir: Path, |
| 90 | keep_output_dir: bool = False, |
| 91 | ): |
| 92 | self.output_dir = output_dir |
| 93 | self.output_dir.mkdir(parents=True, exist_ok=True) |
| 94 | self.keep_output_dir = keep_output_dir |
| 95 | |
| 96 | self.definition = definition |
| 97 | self.workloads = workloads |
| 98 | self.solution = solution |
| 99 | self.config = config |
| 100 | |
| 101 | # Write problem files to staging directory up front. |
| 102 | (self.output_dir / "definition.json").write_text(definition.model_dump_json()) |
| 103 | (self.output_dir / "workload.jsonl").write_text( |
| 104 | "\n".join(w.model_dump_json() for w in workloads) |
| 105 | ) |
| 106 | (self.output_dir / "solution.json").write_text(solution.model_dump_json()) |
| 107 | (self.output_dir / "config.json").write_text( |
| 108 | json.dumps(dataclasses.asdict(config)) |
| 109 | ) |
| 110 | self._write_sources() |
| 111 | |
| 112 | def __del__(self): |
| 113 | if not self.keep_output_dir: |
| 114 | shutil.rmtree(self.output_dir, ignore_errors=True) |
| 115 | |
| 116 | @property |
| 117 | def _is_cpp(self) -> bool: |
| 118 | return any(lang in _CPP_LANGUAGES for lang in self.solution.spec.languages) |
| 119 | |
| 120 | def _inject_gencode_flags(self, sol_dict: dict) -> dict: |
| 121 | """Auto-inject -gencode flags when no explicit arch flag is set. |
| 122 | |
| 123 | Blackwell targets get sm_100a (required for tcgen05/TMEM instructions). |
| 124 | LOCAL target detects the compile machine's GPU. |
| 125 | """ |
| 126 | spec = sol_dict["spec"] |
| 127 | compile_options = dict(spec.get("compile_options") or {}) |
| 128 | cuda_cflags = list(compile_options.get("cuda_cflags", [])) |
| 129 | |
| 130 | if any("-gencode" in f or "-arch" in f for f in cuda_cflags): |
| 131 | return sol_dict |
| 132 | |
| 133 | gencode_sms: list[str] = [] |
| 134 | target_hw = {h.upper() for h in spec.get("target_hardware", [])} |
| 135 | |
| 136 | if any(h == hw.value for h in target_hw for hw in _BLACKWELL_HARDWARE): |
| 137 | gencode_sms.append("sm_100a") |
no outgoing calls