fn generate_torch()

in build2cmake/src/main.rs [115:187]

• 63 lines of code
• 5 McCabe index (conditional complexity)

fn generate_torch(
    backend: Option<Backend>,
    build_toml: PathBuf,
    target_dir: Option<PathBuf>,
    force: bool,
    ops_id: Option<String>,
) -> Result<()> {
    let target_dir = check_or_infer_target_dir(&build_toml, target_dir)?;

    let build_compat = parse_and_validate(build_toml)?;

    if matches!(build_compat, BuildCompat::V1(_)) {
        eprintln!(
            "build.toml is in the deprecated V1 format, use `build2cmake update-build` to update."
        )
    }

    let build: Build = build_compat
        .try_into()
        .context("Cannot update build configuration")?;

    let mut env = Environment::new();
    env.set_trim_blocks(true);
    minijinja_embed::load_templates!(&mut env);

    let backend = match (backend, build.general.universal) {
        (None, true) => {
            let file_set = write_torch_ext_universal(&env, &build, target_dir.clone(), ops_id)?;
            file_set.write(&target_dir, force)?;
            return Ok(());
        }
        (Some(backend), true) => bail!("Universal kernel, cannot generate for backend {}", backend),
        (Some(backend), false) => {
            if !build.has_kernel_with_backend(&backend) {
                bail!("No kernels found for backend {}", backend);
            }

            backend
        }
        (None, false) => {
            let mut kernel_backends = build.backends();
            let backend = if let Some(backend) = kernel_backends.pop_first() {
                backend
            } else {
                bail!("No kernels found in build.toml");
            };

            if !kernel_backends.is_empty() {
                let kernel_backends: Vec<_> = build
                    .backends()
                    .into_iter()
                    .map(|backend| backend.to_string())
                    .collect();
                bail!(
                    "Multiple supported backends found in build.toml: {}. Please specify one with --backend.",
                    kernel_backends.join(", ")
                );
            }

            backend
        }
    };

    let file_set = match backend {
        Backend::Cuda | Backend::Rocm => {
            write_torch_ext_cuda(&env, backend, &build, target_dir.clone(), ops_id)?
        }
        Backend::Metal => write_torch_ext_metal(&env, &build, target_dir.clone(), ops_id)?,
    };
    file_set.write(&target_dir, force)?;

    Ok(())
}