diff --git a/README.md b/README.md index 6f902f7..8e57e2c 100644 --- a/README.md +++ b/README.md @@ -1,6 +1,9 @@ # ParametricOperators.jl -[![][license-img]][license-status] +[![][license-img]][license-status] +[![Documenter](https://github.com/slimgroup/ParametricOperators.jl/actions/workflows/Documenter.yml/badge.svg)](https://github.com/slimgroup/ParametricOperators.jl/actions/workflows/Documenter.yml) +[![TagBot](https://github.com/slimgroup/ParametricOperators.jl/actions/workflows/TagBot.yml/badge.svg)](https://github.com/slimgroup/ParametricOperators.jl/actions/workflows/TagBot.yml) + ParametricOperators.jl is a Julia Language-based scientific library designed to facilitate the creation and manipulation of tensor operations involving large-scale data using Kronecker products. It provides an efficient and mathematically consistent way to express tensor programs and distribution in the context of machine learning. @@ -16,13 +19,17 @@ ParametricOperators.jl is a Julia Language-based scientific library designed to ```julia julia> using Pkg - julia> Pkg.activate("path/to/your/project") + julia> Pkg.activate("path/to/your/environment") julia> Pkg.add("ParametricOperators") ``` This will add `ParametricOperators.jl` as dependency to your project -## Examples +## Documentation + +[Detailed Documentation hosted here](https://slimgroup.github.io/ParametricOperators.jl) + +## Some Examples ### 1. FFT of 3D Tensor @@ -55,7 +62,7 @@ Make sure to add necessary dependencies. You might also need to load a proper MP ```julia julia> using Pkg -julia> Pkg.activate("path/to/your/project") +julia> Pkg.activate("path/to/your/environment") julia> Pkg.add("MPI") julia> Pkg.add("CUDA") ``` @@ -111,7 +118,7 @@ Make sure to add necessary dependencies to compute the gradient ```julia julia> using Pkg -julia> Pkg.activate("path/to/your/project") +julia> Pkg.activate("path/to/your/environment") julia> Pkg.add("Zygote") ``` @@ -151,7 +158,7 @@ Make sure to add necessary dependencies. You might also need to load a proper MP ```julia julia> using Pkg -julia> Pkg.activate("path/to/your/project") +julia> Pkg.activate("path/to/your/environment") julia> Pkg.add("MPI") julia> Pkg.add("CUDA") julia> Pkg.add("Zygote") diff --git a/docs/make.jl b/docs/make.jl index 822d157..4f60fbd 100644 --- a/docs/make.jl +++ b/docs/make.jl @@ -5,6 +5,17 @@ makedocs( sitename = "ParametricOperators", format = Documenter.HTML(), # modules = [ParametricOperators] + pages=[ + "Introduction" => "index.md", + "Quick Start" => "quickstart.md", + "Examples" => [ + "3D FFT" => "examples/3D_FFT.md", + "Distributed 3D FFT" => "examples/3D_DFFT.md", + "3D Conv" => "examples/3D_Conv.md", + "Distributed 3D Conv" => "examples/3D_DConv.md", + ], + "API" => "api.md" + ] ) # Automatically deploy documentation to gh-pages. diff --git a/docs/src/api.md b/docs/src/api.md new file mode 100644 index 0000000..7353a2a --- /dev/null +++ b/docs/src/api.md @@ -0,0 +1,3 @@ +# ParametricOperators.jl + +Documentation for ParametricOperators.jl diff --git a/docs/src/examples.md b/docs/src/examples.md new file mode 100644 index 0000000..7353a2a --- /dev/null +++ b/docs/src/examples.md @@ -0,0 +1,3 @@ +# ParametricOperators.jl + +Documentation for ParametricOperators.jl diff --git a/docs/src/examples/3D_Conv.md b/docs/src/examples/3D_Conv.md new file mode 100644 index 0000000..afe502a --- /dev/null +++ b/docs/src/examples/3D_Conv.md @@ -0,0 +1,39 @@ +### Parametrized Convolution on 3D Tensor + +Make sure to add necessary dependencies to compute the gradient + +```julia +julia> using Pkg +julia> Pkg.activate("path/to/your/environment") +julia> Pkg.add("Zygote") +``` + +```julia +using Pkg +Pkg.activate("./path/to/your/environment") + +using ParametricOperators +using Zygote + +T = Float32 + +gt, gx, gy = 100, 100, 100 + +# Define a transform along each dimension +St = ParMatrix(T, gt, gt) +Sx = ParMatrix(T, gx, gx) +Sy = ParMatrix(T, gy, gy) + +# Create a Kronecker operator than chains together the transforms +S = Sy ⊗ Sx ⊗ St + +# Parametrize our transform +θ = init(S) |> gpu + +# Apply the transform on a random input +x = rand(T, gt, gx, gy) |> gpu +y = S(θ) * vec(x) + +# Compute the gradient wrt some objective of our parameters +θ′ = gradient(θ -> sum(S(θ) * vec(x)), θ) +``` \ No newline at end of file diff --git a/docs/src/examples/3D_DConv.md b/docs/src/examples/3D_DConv.md new file mode 100644 index 0000000..6bd55cc --- /dev/null +++ b/docs/src/examples/3D_DConv.md @@ -0,0 +1,62 @@ +### Distributed Parametrized Convolution of a 3D Tensor: + +Make sure to add necessary dependencies. You might also need to load a proper MPI implementation based on your hardware. + +```julia +julia> using Pkg +julia> Pkg.activate("path/to/your/environment") +julia> Pkg.add("MPI") +julia> Pkg.add("CUDA") +julia> Pkg.add("Zygote") +``` + +Copy the following code into a `.jl` file +```julia +using Pkg +Pkg.activate("./path/to/your/environment") + +using ParametricOperators +using CUDA +using MPI + +MPI.Init() + +comm = MPI.COMM_WORLD +rank = MPI.Comm_rank(comm) +size = MPI.Comm_size(comm) + +# Julia requires you to manually assign the gpus, modify to your case. +CUDA.device!(rank % 4) +partition = [1, 1, size] + +T = Float32 + +# Define your Global Size and Data Partition +gt, gx, gy = 100, 100, 100 +nt, nx, ny = [gt, gx, gy] .÷ partition + +# Define a transform along each dimension +St = ParMatrix(T, gt, gt) +Sx = ParMatrix(T, gx, gx) +Sy = ParMatrix(T, gy, gy) + +# Create and distribute the Kronecker operator than chains together the transforms +S = Sy ⊗ Sx ⊗ St +S = distribute(S, partition) + +# Parametrize our transform +θ = init(S) |> gpu + +# Apply the transform on a random input +x = rand(T, nt, nx, ny) |> gpu +y = S(θ) * vec(x) + +# Compute the gradient wrt some objective of our parameters +θ′ = gradient(θ -> sum(S(θ) * vec(x)), θ) + +MPI.Finalize() +``` + +You can run the above by doing: + +`srun -n N_TASKS julia code_above.jl` diff --git a/docs/src/examples/3D_DFFT.md b/docs/src/examples/3D_DFFT.md new file mode 100644 index 0000000..69bbabf --- /dev/null +++ b/docs/src/examples/3D_DFFT.md @@ -0,0 +1,55 @@ +### Distributed FFT of a 3D Tensor: + +Make sure to add necessary dependencies. You might also need to load a proper MPI implementation based on your hardware. + +```julia +julia> using Pkg +julia> Pkg.activate("path/to/your/environment") +julia> Pkg.add("MPI") +julia> Pkg.add("CUDA") +``` + +Copy the following code into a `.jl` file +```julia +using Pkg +Pkg.activate("./path/to/your/environment") + +using ParametricOperators +using CUDA +using MPI + +MPI.Init() + +comm = MPI.COMM_WORLD +rank = MPI.Comm_rank(comm) +size = MPI.Comm_size(comm) + +# Julia requires you to manually assign the gpus, modify to your case. +CUDA.device!(rank % 4) +partition = [1, 1, size] + +T = Float32 + +# Define your Global Size and Data Partition +gt, gx, gy = 100, 100, 100 +nt, nx, ny = [gt, gx, gy] .÷ partition + +# Define a transform along each dimension +Ft = ParDFT(T, gt) +Fx = ParDFT(Complex{T}, gx) +Fy = ParDFT(Complex{T}, gy) + +# Create and distribute the Kronecker operator than chains together the transforms +F = Fy ⊗ Fx ⊗ Ft +F = distribute(F, partition) + +# Apply the transform on a random input +x = rand(T, nt, nx, ny) |> gpu +y = F * vec(x) + +MPI.Finalize() +``` + +You can run the above by doing: + +`srun -n N_TASKS julia code_above.jl` diff --git a/docs/src/examples/3D_FFT.md b/docs/src/examples/3D_FFT.md new file mode 100644 index 0000000..0941fc4 --- /dev/null +++ b/docs/src/examples/3D_FFT.md @@ -0,0 +1,24 @@ +### FFT of 3D Tensor + +```julia +using Pkg +Pkg.activate("./path/to/your/environment") + +using ParametricOperators + +T = Float32 + +gt, gx, gy = 100, 100, 100 + +# Define a transform along each dimension +Ft = ParDFT(T, gt) +Fx = ParDFT(Complex{T}, gx) +Fy = ParDFT(Complex{T}, gy) + +# Create a Kronecker operator than chains together the transforms +F = Fy ⊗ Fx ⊗ Ft + +# Apply the transform on a random input +x = rand(T, gt, gx, gy) |> gpu +y = F * vec(x) +``` diff --git a/docs/src/index.md b/docs/src/index.md index 7353a2a..5016e66 100644 --- a/docs/src/index.md +++ b/docs/src/index.md @@ -1,3 +1,3 @@ -# ParametricOperators.jl +### ParametricOperators.jl Documentation for ParametricOperators.jl diff --git a/docs/src/quickstart.md b/docs/src/quickstart.md new file mode 100644 index 0000000..7f21943 --- /dev/null +++ b/docs/src/quickstart.md @@ -0,0 +1,72 @@ +## Installation + +Add `ParametricOperators.jl` as a dependency to your environment. + +To add, either do: + +```julia +julia> ] +(v1.9) add ParametricOperators +``` + +OR + +```julia +julia> using Pkg +julia> Pkg.activate("path/to/your/environment") +julia> Pkg.add("ParametricOperators") +``` + +## Simple Operator + +Make sure to include the package in your environment + +```julia +using ParametricOperators +``` + +Lets start by defining a Matrix Operator of size `10x10`: + +```julia +A = ParMatrix(10, 10) +``` + +Now, we parametrize our operator with some weights `θ`: + +```julia +θ = init(A) +``` + +We can now apply our operator on some random input: + +```julia +x = rand(10) +A(θ) * x +``` + +## Gradient Computation + +!!! note "Limited AD support" + Current support only provided for Zygote.jl + +Make sure to include an AD package in your environment + +```julia +using Zygote +``` + +Using the example above, one can find the gradient of the weights or your input w.r.t to some objective using a standard AD package: + +```julia +# Gradient w.r.t weights +θ′ = gradient(θ -> sum(A(θ) * x), θ) + +# Gradient w.r.t input +x′ = gradient(x -> sum(A(θ) * x), x) +``` + +## Chaining Operators + +## Kronecker Operator + +## Distributing Kronecker Operator