gen_intrinsics.py

# Generates rust code for every intrinsic from libdevice.json made in gen_libdevice_json.py

import os
import json
import inspect

dirname = os.path.dirname(__file__)
filename = os.path.join(dirname, 'data/libdevice.json')
out_filename = os.path.join(dirname, 'data/std_intrinsics.rs')

txt = open(filename, "r", encoding="utf8").read()
raw = json.loads(txt)

header = '''//! Raw libdevice math intrinsics.
//! 
//! Note that this file was autogenerated using crude text analysis
//! from the libdevice PDF file, therefore many of the descriptions have broken
//! text, especially for math symbols. The link to the libdevice website is provided
//! for every intrinsic so you can view the non-broken description. 
//! 
//! Most of the intrinsics here have "proper" functions, corresponding f32/f64 functions
//! are already codegenned to libdevice intrinsics by the codegen automatically. This module
//! is mostly for exotic intrinsics that have not been added as proper functions yet.
//!
//! The underlying intrinsic functions have a prefix of `__nv_`, however this prefix
//! is stripped from the functions in this module just for convenience.

// Generated file, do not edit by hand, see scripts/gen_intrinsics.py

'''
out = ""

for intrinsic in raw:
    raw_name = intrinsic["name"]
    out += f"#[link_name = \"{raw_name}\"]\n"

    name = raw_name.removeprefix("__nv_")
    description = intrinsic["description"]
    returns = intrinsic["returns"]

    # There isnt actually any availability which is `No` instead of `Yes`, so including it is useless for now
    hyperlink = f"[Nvidia docs](https://docs.nvidia.com/cuda/libdevice-users-guide/{raw_name}.html#{raw_name})"
    full_desc = inspect.cleandoc(
        f"{description}\n\n{hyperlink}\n\n# Returns\n\n{returns}")
    out += f"#[doc = \"{full_desc}\"]\n"
    out += f"pub fn {name}("
    sig = intrinsic["sig"]
    return_ty = sig["returns"]

    params = []
    for param in sig["params"]:
        param_name = param["name"]
        param_ty = param["type"]
        if param_name == "in":
            param_name = "in_"

        params.append(f"{param_name}: {param_ty}")

    out += ", ".join(params)
    out += f") -> {return_ty};\n\n"

out = f"{header}extern \"C\" {{\n{out}}}\n"

open(out_filename, "w", encoding="utf8").write(out)