Function main

etc/scripts/cutedsl_wrapper.py:231–279

()

def main() -> None:
    parser = argparse.ArgumentParser(description="CuTe DSL Compiler Explorer wrapper")
    parser.add_argument("input_file", type=Path, help="Path to the input Python file")
    parser.add_argument("--output_file", type=Path, required=True, help="Path to the output PTX file")
    parser.add_argument("--arch", default="sm_90a", help="Value for CUTE_DSL_ARCH")
    parser.add_argument("--keep", default=None, help="Value for CUTE_DSL_KEEP")
    parser.add_argument("--python_path", default="", help="Extra PYTHONPATH entries separated by os.pathsep")
    parser.add_argument("--no_runtime_patch", action="store_true", help="Allow compiled kernels to run")
    args, source_args = parser.parse_known_args()

    output_file = args.output_file
    dump_dir = output_file.parent
    dump_dir.mkdir(parents=True, exist_ok=True)
    output_file.unlink(missing_ok=True)

    add_python_paths(args.python_path)
    add_python_paths(os.environ.get("CUTE_DSL_PYTHONPATH", ""))

    os.environ["CUTE_DSL_ARCH"] = args.arch
    os.environ["CUTE_DSL_DUMP_DIR"] = str(dump_dir)
    os.environ["CUTE_DSL_KEEP"] = args.keep or os.environ.get("CUTE_DSL_KEEP", "ir,ptx,cubin")
    os.environ.setdefault("CUTE_DSL_KEEP_IR", "1")
    os.environ.setdefault("CUTE_DSL_KEEP_PTX", "1")
    os.environ.setdefault("CUTE_DSL_KEEP_CUBIN", "1")
    os.environ.setdefault("CUTE_DSL_NO_CACHE", "1")
    os.environ.setdefault("CUTE_DSL_DISABLE_FILE_CACHING", "1")
    os.environ.setdefault("CUDA_TOOLKIT_PATH", "/usr/local")

    try:
        import cutlass.cute as cute
    except ModuleNotFoundError as exc:
        raise SystemExit(
            "CuTe DSL is not importable. Use a Python with cutlass.cute installed, "
            f"or pass --python_path /path/to/cutlass/python/CuTeDSL. Missing module: {exc.name}."
        ) from exc

    patch_compile_defaults(cute, args.arch, dump_dir)

    if not args.no_runtime_patch:
        patch_runtime_calls()

    old_argv = sys.argv
    try:
        sys.argv = [str(args.input_file), *source_args]
        runpy.run_path(str(args.input_file), run_name="__main__")
    finally:
        sys.argv = old_argv

    write_primary_output(output_file, dump_dir)


if __name__ == "__main__":