(output_file: Path, dump_dir: Path)
| 203 | |
| 204 | |
| 205 | def write_primary_output(output_file: Path, dump_dir: Path) -> None: |
| 206 | write_sass_outputs(dump_dir) |
| 207 | ptx_files = [ |
| 208 | path |
| 209 | for path in sorted(dump_dir.rglob("*.ptx"), key=lambda item: (item.stat().st_mtime_ns, str(item))) |
| 210 | if path.resolve() != output_file.resolve() |
| 211 | ] |
| 212 | mlir_files = sorted(dump_dir.rglob("*.mlir"), key=lambda item: (item.stat().st_mtime_ns, str(item))) |
| 213 | |
| 214 | if ptx_files: |
| 215 | with output_file.open("w") as output: |
| 216 | for path in ptx_files: |
| 217 | output.write(f"// {path.name}\n") |
| 218 | output.write(path.read_text(errors="replace")) |
| 219 | output.write("\n\n") |
| 220 | return |
| 221 | |
| 222 | message = [ |
| 223 | "// No CuTe DSL PTX was generated.", |
| 224 | "// Make sure the source imports cutlass.cute and calls cute.compile(...).", |
| 225 | ] |
| 226 | if mlir_files: |
| 227 | message.append("// Generated MLIR is available in the device-code view.") |
| 228 | output_file.write_text("\n".join(message) + "\n") |
| 229 | |
| 230 | |
| 231 | def main() -> None: |
no test coverage detected