<commit_before>#include <cassert>
#include <algorithm>
#include <iostream>
#include <mutex>
#include <thread>
#include <agency/executor/executor_array.hpp>
#include <agency/cuda/executor/parallel_executor.hpp>

__managed__ int result;

struct make_7
{
  __host__ __device__
  int operator()() const
  {
    return 7;
  }
};

struct make_42
{
  __host__ __device__
  int operator()() const
  {
    return 42;
  }
};

struct make_1
{
  __host__ __device__
  int operator()() const
  {
    return 1;
  }
};

struct functor
{
  int *result;

  template<class Index>
  __device__
  void operator()(const Index&, int& past, int& outer_shared, int& inner_shared, int& inner_inner_shared)
  {
    atomicAdd(result, past + outer_shared + inner_shared + inner_inner_shared);
  }
};

int main()
{
  using namespace agency;

  using outer_executor_type = cuda::this_thread::parallel_executor;
  using inner_executor_type = cuda::grid_executor;

  {
    // test executor_array then_execute()
    using executor_type = executor_array<inner_executor_type, outer_executor_type>;
    using traits = agency::executor_traits<executor_type>;
    using shape_type = typename traits::shape_type;
    using index_type = typename traits::index_type;

    executor_type exec(2);

    for(size_t i = 0; i < 2; ++i)
    {
      exec[i].device(i);
    }

    auto shape = exec.make_shape(2,{2,2});

    auto past = traits::make_ready_future<int>(exec, 13);

    auto f = exec.then_execute([=] __host__ __device__ (const index_type& idx, int& past, int& outer_shared, int& inner_shared, int& inner_inner_shared)
    {
      printf("hello from agent %d %d %d\n", (int)agency::detail::get<0>(idx), (int)agency::detail::get<1>(idx), (int)agency::detail::get<2>(idx));
      return past + outer_shared + inner_shared + inner_inner_shared;
    },
    [](shape_type shape)
    {
      return traits::container<int>(shape);
    },
    shape,
    past,
    make_7(),
    make_42(),
    make_1());

    auto results = f.get();

    assert(results.size() == agency::detail::shape_size(shape));
    assert(std::all_of(results.begin(), results.end(), [](int x){ return x == 13 + 7 + 42 + 1; }));
  }

  {
    // test executor_array then_execute() returning void
    using executor_type = executor_array<inner_executor_type, outer_executor_type>;
    using traits = agency::executor_traits<executor_type>;
    using shape_type = typename traits::shape_type;

    executor_type exec(2);

    auto past = traits::make_ready_future<int>(exec,13);

    auto shape = exec.make_shape(3,5);

    result = 0;

    auto f = traits::then_execute(exec,
    functor{&result},
    shape,
    past,
    make_7(),
    make_42(),
    make_1());

    f.wait();

    assert(result == agency::detail::shape_size(shape) * (13 + 7 + 42 + 1));
  }

  {
    // test executor_array async_execute()
    using executor_type = executor_array<inner_executor_type, outer_executor_type>;
    using traits = agency::executor_traits<executor_type>;
    using shape_type = typename traits::shape_type;
    using index_type = typename traits::index_type;

    executor_type exec(2);

    for(size_t i = 0; i < 2; ++i)
    {
      exec[i].device(i);
    }

    auto shape = exec.make_shape(2,{2,2});

    auto f = traits::async_execute(exec, [] __host__ __device__ (const index_type& idx, int& outer_shared, int& inner_shared, int& inner_inner_shared)
    {
      return 13 + outer_shared + inner_shared + inner_inner_shared;
    },
    [](shape_type shape)
    {
      return traits::container<int>(shape);
    },
    shape,
    make_7(),
    make_42(),
    make_1());

    auto results = f.get();

    assert(results.size() == agency::detail::shape_size(shape));
    assert(std::all_of(results.begin(), results.end(), [](int x){ return x == 13 + 7 + 42 + 1; }));
  }

  {
    // test flattened executor_array
    using executor_array_type = executor_array<inner_executor_type, outer_executor_type>;
    using executor_type = flattened_executor<executor_array_type>;

    using traits = agency::executor_traits<executor_type>;
    using shape_type = typename traits::shape_type;
    using index_type = typename traits::index_type;

    executor_array_type exec_array(2);

    for(size_t i = 0; i < 2; ++i)
    {
      exec_array[i].device(i);
    }

    executor_type exec{exec_array};

    shape_type shape{2, 2};

    auto ready = traits::make_ready_future<void>(exec);<commit_msg>Make this test robustly handle the actual number of GPUs in the machine<commit_after>#include <cassert>
#include <algorithm>
#include <iostream>
#include <mutex>
#include <thread>
#include <agency/executor/executor_array.hpp>
#include <agency/cuda/executor/parallel_executor.hpp>

__managed__ int result;

struct make_7
{
  __host__ __device__
  int operator()() const
  {
    return 7;
  }
};

struct make_42
{
  __host__ __device__
  int operator()() const
  {
    return 42;
  }
};

struct make_1
{
  __host__ __device__
  int operator()() const
  {
    return 1;
  }
};

struct functor
{
  int *result;

  template<class Index>
  __device__
  void operator()(const Index&, int& past, int& outer_shared, int& inner_shared, int& inner_inner_shared)
  {
    atomicAdd(result, past + outer_shared + inner_shared + inner_inner_shared);
  }
};

int main()
{
  using namespace agency;

  using outer_executor_type = cuda::this_thread::parallel_executor;
  using inner_executor_type = cuda::grid_executor;

  int num_devices = 0;
  cudaError_t error = cudaGetDeviceCount(&num_devices);
  if(error)
  {
    std::string what("CUDA error after cudaGetDeviceCount(): ");
    what += std::string(cudaGetErrorString(error));
    throw std::runtime_error(what);
  }

  {
    // test executor_array then_execute()
    using executor_type = executor_array<inner_executor_type, outer_executor_type>;
    using traits = agency::executor_traits<executor_type>;
    using shape_type = typename traits::shape_type;
    using index_type = typename traits::index_type;

    executor_type exec(num_devices);

    for(size_t i = 0; i < exec.size(); ++i)
    {
      exec[i].device(i);
    }

    auto shape = exec.make_shape(exec.size(),{2,2});

    auto past = traits::make_ready_future<int>(exec, 13);

    auto f = exec.then_execute([=] __host__ __device__ (const index_type& idx, int& past, int& outer_shared, int& inner_shared, int& inner_inner_shared)
    {
      printf("hello from agent %d %d %d\n", (int)agency::detail::get<0>(idx), (int)agency::detail::get<1>(idx), (int)agency::detail::get<2>(idx));
      return past + outer_shared + inner_shared + inner_inner_shared;
    },
    [](shape_type shape)
    {
      return traits::container<int>(shape);
    },
    shape,
    past,
    make_7(),
    make_42(),
    make_1());

    auto results = f.get();

    assert(results.size() == agency::detail::shape_size(shape));
    assert(std::all_of(results.begin(), results.end(), [](int x){ return x == 13 + 7 + 42 + 1; }));
  }

  {
    // test executor_array then_execute() returning void
    using executor_type = executor_array<inner_executor_type, outer_executor_type>;
    using traits = agency::executor_traits<executor_type>;
    using shape_type = typename traits::shape_type;

    executor_type exec(num_devices);

    auto past = traits::make_ready_future<int>(exec,13);

    auto shape = exec.make_shape(3,5);

    result = 0;

    auto f = traits::then_execute(exec,
    functor{&result},
    shape,
    past,
    make_7(),
    make_42(),
    make_1());

    f.wait();

    assert(result == agency::detail::shape_size(shape) * (13 + 7 + 42 + 1));
  }

  {
    // test executor_array async_execute()
    using executor_type = executor_array<inner_executor_type, outer_executor_type>;
    using traits = agency::executor_traits<executor_type>;
    using shape_type = typename traits::shape_type;
    using index_type = typename traits::index_type;

    executor_type exec(num_devices);

    for(size_t i = 0; i < exec.size(); ++i)
    {
      exec[i].device(i);
    }

    auto shape = exec.make_shape(exec.size(),{2,2});

    auto f = traits::async_execute(exec, [] __host__ __device__ (const index_type& idx, int& outer_shared, int& inner_shared, int& inner_inner_shared)
    {
      return 13 + outer_shared + inner_shared + inner_inner_shared;
    },
    [](shape_type shape)
    {
      return traits::container<int>(shape);
    },
    shape,
    make_7(),
    make_42(),
    make_1());

    auto results = f.get();

    assert(results.size() == agency::detail::shape_size(shape));
    assert(std::all_of(results.begin(), results.end(), [](int x){ return x == 13 + 7 + 42 + 1; }));
  }

  {
    // test flattened executor_array
    using executor_array_type = executor_array<inner_executor_type, outer_executor_type>;
    using executor_type = flattened_executor<executor_array_type>;

    using traits = agency::executor_traits<executor_type>;
    using shape_type = typename traits::shape_type;
    using index_type = typename traits::index_type;

    executor_array_type exec_array(num_devices);

    for(size_t i = 0; i < exec_array.size(); ++i)
    {
      exec_array[i].device(i);
    }

    executor_type exec{exec_array};

    shape_type shape{exec_array.size(), 2};

    auto ready = traits::make_ready_future<void>(exec);
{'fraction_non_alphanumeric': 0.13456750027065065, 'fraction_numerical': 0.013749052722745481, 'mean_word_length': 3.831589958158996, 'pattern_counts': {'":': 0, '<': 59, '<?xml version=': 0, '>': 53, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2759730', 'n_tokens_mistral': 3281, 'n_tokens_neox': 3101, 'n_words': 663}