TheRock/rocm-libraries/projects/rocthrust/examples/cuda/range_view.cu

// Copyright (c) 2020-2025 Advanced Micro Devices, Inc. All rights reserved.
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in all
// copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
// SOFTWARE.

#include <thrust/device_vector.h>
#include <thrust/execution_policy.h>
#include <thrust/for_each.h>
#include <thrust/iterator/counting_iterator.h>

#include <iostream>

#include "../include/host_device.h"

// This example demonstrates the use of a view: a non-owning wrapper for an
// iterator range which presents a container-like interface to the user.
//
// For example, a view of a device_vector's data can be helpful when we wish to
// access that data from a device function. Even though device_vectors are not
// accessible from device functions, the range_view class allows us to access
// and manipulate its data as if we were manipulating a real container.

template <class Iterator>
class range_view
{
public:
  using iterator        = Iterator;
  using value_type      = typename thrust::iterator_traits<iterator>::value_type;
  using pointer         = typename thrust::iterator_traits<iterator>::pointer;
  using difference_type = typename thrust::iterator_traits<iterator>::difference_type;
  using reference       = typename thrust::iterator_traits<iterator>::reference;

private:
  const iterator first;
  const iterator last;

public:
  __host__ __device__ range_view(Iterator first, Iterator last)
      : first(first)
      , last(last)
  {}
  __host__ __device__ ~range_view() {}

  __host__ __device__ difference_type size() const
  {
    return thrust::distance(first, last);
  }

  __host__ __device__ reference operator[](difference_type n)
  {
    return *(first + n);
  }
  __host__ __device__ const reference operator[](difference_type n) const
  {
    return *(first + n);
  }

  __host__ __device__ iterator begin()
  {
    return first;
  }
  __host__ __device__ const iterator cbegin() const
  {
    return first;
  }
  __host__ __device__ iterator end()
  {
    return last;
  }
  __host__ __device__ const iterator cend() const
  {
    return last;
  }

  __host__ __device__ thrust::reverse_iterator<iterator> rbegin()
  {
    return thrust::reverse_iterator<iterator>(end());
  }
  __host__ __device__ const thrust::reverse_iterator<const iterator> crbegin() const
  {
    return thrust::reverse_iterator<const iterator>(cend());
  }
  __host__ __device__ thrust::reverse_iterator<iterator> rend()
  {
    return thrust::reverse_iterator<iterator>(begin());
  }
  __host__ __device__ const thrust::reverse_iterator<const iterator> crend() const
  {
    return thrust::reverse_iterator<const iterator>(cbegin());
  }
  __host__ __device__ reference front()
  {
    return *begin();
  }
  __host__ __device__ const reference front() const
  {
    return *cbegin();
  }

  __host__ __device__ reference back()
  {
    return *end();
  }
  __host__ __device__ const reference back() const
  {
    return *cend();
  }

  __host__ __device__ bool empty() const
  {
    return size() == 0;
  }
};

// This helper function creates a range_view from iterator and the number of
// elements
template <class Iterator, class Size>
range_view<Iterator> __host__ __device__ make_range_view(Iterator first, Size n)
{
  return range_view<Iterator>(first, first + n);
}

// This helper function creates a range_view from a pair of iterators
template <class Iterator>
range_view<Iterator> __host__ __device__ make_range_view(Iterator first, Iterator last)
{
  return range_view<Iterator>(first, last);
}

// This helper function creates a range_view from a Vector
template <class Vector>
range_view<typename Vector::iterator> __host__ make_range_view(Vector& v)
{
  return range_view<typename Vector::iterator>(v.begin(), v.end());
}

// This saxpy functor stores view of X, Y, Z array, and accesses them in
// vector-like way
template <class View1, class View2, class View3>
struct saxpy_functor
{
  const float a;
  View1 x;
  View2 y;
  View3 z;

  __host__ __device__ saxpy_functor(float _a, View1 _x, View2 _y, View3 _z)
      : a(_a)
      , x(_x)
      , y(_y)
      , z(_z)
  {}

  __host__ __device__ void operator()(int i)
  {
    z[i] = a * x[i] + y[i];
  }
};

// saxpy function, which can either be called form host or device
// The views are passed by value
template <class View1, class View2, class View3>
__host__ __device__ void saxpy(float A, View1 X, View2 Y, View3 Z)
{
  // Z = A * X + Y
  const int size = X.size();
  thrust::for_each(thrust::device,
                   thrust::make_counting_iterator(0),
                   thrust::make_counting_iterator(size),
                   saxpy_functor<View1, View2, View3>(A, X, Y, Z));
}

struct f1
{
  __host__ __device__ float operator()(float x) const
  {
    return x * 3;
  }
};

int main()
{
  using std::cout;
  using std::endl;

  // initialize host arrays
  float x[4] = {1.0, 1.0, 1.0, 1.0};
  float y[4] = {1.0, 2.0, 3.0, 4.0};
  float z[4] = {0.0};

  thrust::device_vector<float> X(x, x + 4);
  thrust::device_vector<float> Y(y, y + 4);
  thrust::device_vector<float> Z(z, z + 4);

  saxpy(
    2.0,

    // make a range view of a pair of transform_iterators
    make_range_view(thrust::make_transform_iterator(X.cbegin(), f1()), thrust::make_transform_iterator(X.cend(), f1())),

    // range view of normal_iterators
    make_range_view(Y.begin(), thrust::distance(Y.begin(), Y.end())),

    // range view of naked pointers
    make_range_view(Z.data().get(), 4));

  // print values from original device_vector<float> Z
  // to ensure that range view was mapped to this vector
  for (std::size_t i = 0, n = Z.size(); i < n; ++i)
  {
    cout << "z[" << i << "]= " << Z[i] << endl;
  }

  return 0;
}