Skip to content

Feature request: Code generation for pytorch #63

@jarlebring

Description

@jarlebring

(Based on a comment from when I presented this result at EPFL.) The generated code is very GPU-friendly. We could generate code also for for PyTorch.

In fact the Bader, Blanes, Casas (in our code graph_bbc) method is currently default in pytorch matrix_exp, and we can do a lot better than the BBC-method so there is certainly room for improvement. It is implemented in the c++ code: mexp_impl and mexp:

https://github.com/pytorch/pytorch/blob/f4dd88489a77ff9b300bf6f9b34c233ca82f76d7/aten/src/ATen/native/LinearAlgebra.cpp#L2111
https://github.com/pytorch/pytorch/blob/f4dd88489a77ff9b300bf6f9b34c233ca82f76d7/aten/src/ATen/native/LinearAlgebra.cpp#L2185
and lower level code compute_T1 ... :
https://github.com/pytorch/pytorch/blob/f4dd88489a77ff9b300bf6f9b34c233ca82f76d7/aten/src/ATen/native/LinearAlgebra.cpp#L1855

I would suggest we add a language with a variant symbol

struct LangCPP
  variant::Symbol
end

If variant=:plain it will generate self-contained C++ code. If variant=:pytorch it generates pytorch-compatible code.

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions