var_toys.hpp

Go to the documentation of this file.

#pragma once
 
#include "async_toys.hpp"
#include "var.hpp"
#include <uni20/core/math.hpp>
 
namespace uni20::async
{
 
namespace detail
{
 
// Reverse tasks must observe absence before touching any other dependency. An
// unseeded result gradient prunes the task; a stored exception still propagates.
template <typename T, typename Transform>
AsyncTask accumulate_transformed_gradient(ReadBuffer<T> input_gradient, WriteBuffer<T> output_gradient,
                                          Transform transform)
{
  auto gradient = co_await input_gradient.transfer().or_cancel();
  T const value = gradient.get();
  gradient.release();
  co_await output_gradient += transform(value);
}
 
template <typename T, typename Transform>
AsyncTask accumulate_transformed_gradient(ReadBuffer<T> input_gradient, ReadBuffer<T> forward_value,
                                          WriteBuffer<T> output_gradient, Transform transform)
{
  auto gradient = co_await input_gradient.transfer().or_cancel();
  T const gradient_value = gradient.get();
  gradient.release();
 
  auto value = co_await forward_value.transfer();
  T const input_value = value.get();
  value.release();
  co_await output_gradient += transform(gradient_value, input_value);
}
 
} // namespace detail
 
template <typename T> Var<T> sin(Var<T> x)
{
  Var<T> Result;
 
  schedule(co_sin(x.value.read(), Result.value.write()));
 
  // TRACE("Var Sin", Result.grad.value().queue().latest().get());
  schedule([](ReadBuffer<T> in, ReadBuffer<T> in_grad, WriteBuffer<T> out_grad) static -> AsyncTask {
    using std::cos;
    // GCC 13 workaround for `(co_await ...).get()`: use an explicit owning proxy.
    // GCC 14+ supports the one-liner form, e.g.
    //   auto const in_g = (co_await in_grad.transfer().or_cancel()).get();
    auto in_grad_buffer = co_await in_grad.transfer().or_cancel();
    auto const in_g = in_grad_buffer.get();
    in_grad_buffer.release();
 
    auto in_buffer = co_await in.transfer();
    auto const cos_x = cos(in_buffer.get());
    in_buffer.release();
    co_await out_grad += uni20::conj(cos_x) * in_g;
  }(x.value.read(), Result.grad.input(), x.grad.output()));
 
  return Result;
}
 
template <typename T> Var<T> cos(Var<T> x)
{
  Var<T> Result;
  schedule(co_cos(x.value.read(), Result.value.write()));
  schedule(detail::accumulate_transformed_gradient(Result.grad.input(), x.value.read(), x.grad.output(),
                                                   [](T const& gradient, T const& input) {
                                                     using std::sin;
                                                     return -uni20::conj(sin(input)) * gradient;
                                                   }));
  return Result;
}
 
template <typename T> Var<T> operator-(Var<T> x, T y)
{
  Var<T> Result;
  Result.value = x.value - y;
  x.grad += Result.grad;
  return Result;
}
 
template <typename T> Var<T> operator-(T x, Var<T> y)
{
  Var<T> Result;
  Result.value = x - y.value;
  y.grad -= Result.grad;
  return Result;
}
 
template <typename T> Var<T> operator-(Var<T> x, Var<T> y)
{
  Var<T> Result;
  Result.value = x.value - y.value;
  x.grad += Result.grad;
  y.grad -= Result.grad;
  return Result;
}
 
template <typename T> Var<T> operator+(Var<T> x, T y)
{
  Var<T> Result;
  Result.value = x.value + y;
  x.grad += Result.grad;
  return Result;
}
 
template <typename T> Var<T> operator+(T x, Var<T> y)
{
  Var<T> Result;
  Result.value = x + y.value;
  y.grad += Result.grad;
  return Result;
}
 
template <typename T> Var<T> operator+(Var<T> x, Var<T> y)
{
  Var<T> Result;
  Result.value = x.value + y.value;
  x.grad += Result.grad;
  y.grad += Result.grad;
  return Result;
}
 
template <typename T> Var<T> operator*(T x, Var<T> y)
{
  Var<T> Result;
  Result.value = x * y.value;
  schedule(detail::accumulate_transformed_gradient(
      Result.grad.input(), y.grad.output(),
      [factor = uni20::herm(x)](T const& gradient) { return factor * gradient; }));
  return Result;
}
 
template <typename T> Var<T> operator*(Var<T> x, T y)
{
  Var<T> Result;
  Result.value = x.value * y;
  schedule(detail::accumulate_transformed_gradient(
      Result.grad.input(), x.grad.output(),
      [factor = uni20::herm(y)](T const& gradient) { return gradient * factor; }));
  return Result;
}
 
template <typename T> Var<T> operator*(Var<T> x, Var<T> y)
{
  Var<T> Result;
  Result.value = x.value * y.value;
  schedule(detail::accumulate_transformed_gradient(
      Result.grad.input(), y.value.read(), x.grad.output(),
      [](T const& gradient, T const& value) { return gradient * uni20::herm(value); }));
  schedule(detail::accumulate_transformed_gradient(
      Result.grad.input(), x.value.read(), y.grad.output(),
      [](T const& gradient, T const& value) { return uni20::herm(value) * gradient; }));
  return Result;
}
 
template <typename T> Var<uni20::make_real_t<T>> real(Var<T> z)
{
  using uni20::real;
  using r_type = uni20::make_real_t<T>;
  Var<r_type> Result;
  Result.value = real(z.value);
  schedule([](ReadBuffer<r_type> in_grad, WriteBuffer<T> out_grad) static -> AsyncTask {
    auto in_grad_buffer = co_await in_grad.transfer().or_cancel();
    auto const grad = in_grad_buffer.get();
    in_grad_buffer.release();
    co_await out_grad += grad;
  }(Result.grad.input(), z.grad.output()));
  return Result;
}
 
template <typename T> Var<uni20::make_real_t<T>> imag(Var<T> z)
{
  using uni20::imag;
  using r_type = uni20::make_real_t<T>;
  Var<r_type> Result;
  Result.value = imag(z.value);
  schedule([](ReadBuffer<r_type> in_grad, WriteBuffer<T> out_grad) static -> AsyncTask {
    auto in_grad_buffer = co_await in_grad.transfer().or_cancel();
    auto const grad = in_grad_buffer.get();
    in_grad_buffer.release();
    T value{};
    imag(value) = grad;
    co_await out_grad += std::move(value);
  }(Result.grad.input(), z.grad.output()));
  return Result;
}
 
} // namespace uni20::async