host/tensor_fill.h

Go to the documentation of this file.

/***************************************************************************************************
 * Copyright (c) 2017-2019, NVIDIA CORPORATION.  All rights reserved.
 *
 * Redistribution and use in source and binary forms, with or without modification, are permitted
 * provided that the following conditions are met:
 *     * Redistributions of source code must retain the above copyright notice, this list of
 *       conditions and the following disclaimer.
 *     * Redistributions in binary form must reproduce the above copyright notice, this list of
 *       conditions and the following disclaimer in the documentation and/or other materials
 *       provided with the distribution.
 *     * Neither the name of the NVIDIA CORPORATION nor the names of its contributors may be used
 *       to endorse or promote products derived from this software without specific prior written
 *       permission.
 *
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR
 * IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND
 * FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE
 * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING,
 * BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS;
 * OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT,
 * STRICT LIABILITY, OR TOR (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
 * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 *
 **************************************************************************************************/
/* \file
  \brief Provides several functions for filling tensors with data.
*/

#pragma once

// Standard Library includes
#include <utility>
#include <cstdlib>
#include <cmath>

// Cutlass includes
#include "cutlass/cutlass.h"
#include "cutlass/complex.h"
#include "cutlass/array.h"
#include "cutlass/numeric_types.h"

#include "cutlass/util/distribution.h"
#include "tensor_foreach.h"


namespace cutlass {
namespace reference {
namespace host {


namespace detail {

template <
  typename Element,               
  typename Layout>                
struct TensorFillFunc {

  using TensorView = TensorView<Element, Layout>;

  //
  // Data members
  //

  TensorView view;
  Element value;

  //
  // Methods
  //

  TensorFillFunc(
    TensorView const &view_ = TensorView(), 
    Element value_ = Element(0)
  ): view(view_), value(value_) { }

  void operator()(Coord<Layout::kRank> const & coord) const {
    view.at(coord) = value;
  }
};

} // namespace detail


template <
  typename Element,               
  typename Layout>                
void TensorFill(
  TensorView<Element, Layout> dst,    
  Element val = Element(0)) {               

  detail::TensorFillFunc<Element, Layout> func(dst, val);

  TensorForEach(
    dst.extent(),
    func
  );
}


namespace detail {

template <typename Element>
struct RandomGaussianFunc {

  uint64_t seed;
  double mean;
  double stddev;
  int int_scale;
  double pi;

  //
  // Methods
  //
  RandomGaussianFunc(
    uint64_t seed_ = 0, 
    double mean_ = 0, 
    double stddev_ = 1,
    int int_scale_ = -1
  ):
    seed(seed_), mean(mean_), stddev(stddev_), int_scale(int_scale_), pi(std::acos(-1)) {
      std::srand((unsigned)seed);
  }

  Element operator()() const {

    // Box-Muller transform to generate random numbers with Normal distribution
    double u1 = double(std::rand()) / double(RAND_MAX);
    double u2 = double(std::rand()) / double(RAND_MAX);

    // Compute Gaussian random value
    double rnd = std::sqrt(-2 * std::log(u1)) * std::cos(2 * pi * u2);
    rnd = mean + stddev * rnd;

    // Scale and convert final result
    Element result;

    if (int_scale >= 0) {
      rnd = double(int64_t(rnd * double(1 << int_scale))) / double(1 << int_scale);
      result = static_cast<Element>(rnd);
    }
    else {
      result = static_cast<Element>(rnd);
    }

    return result;
  }
};

template <typename Element>
struct RandomGaussianFunc<complex<Element> > {

  uint64_t seed;
  double mean;
  double stddev;
  int int_scale;
  double pi;

  //
  // Methods
  //
  RandomGaussianFunc(
    uint64_t seed_ = 0, 
    double mean_ = 0, 
    double stddev_ = 1,
    int int_scale_ = -1
  ):
    seed(seed_), mean(mean_), stddev(stddev_), int_scale(int_scale_), pi(std::acos(-1)) {
      std::srand((unsigned)seed);
  }

  complex<Element> operator()() const {

    Element reals[2];

    for (int i = 0; i < 2; ++i) {
      // Box-Muller transform to generate random numbers with Normal distribution
      double u1 = double(std::rand()) / double(RAND_MAX);
      double u2 = double(std::rand()) / double(RAND_MAX);

      // Compute Gaussian random value
      double rnd = std::sqrt(-2 * std::log(u1)) * std::cos(2 * pi * u2);
      rnd = mean + stddev * rnd;

      if (int_scale >= 0) {
        rnd = double(int(rnd * double(1 << int_scale)));
        reals[i] = from_real<Element>(rnd / double(1 << int_scale));
      }
      else {
        reals[i] = from_real<Element>(rnd);
      }
    }

    return complex<Element>(reals[0], reals[1]);
  }
};

template <
  typename Element,               
  typename Layout>                
struct TensorFillGaussianFunc {

  using TensorView = TensorView<Element, Layout>;

  //
  // Data members
  //

  TensorView view;
  RandomGaussianFunc<Element> func;

  //
  // Methods
  //

  TensorFillGaussianFunc(
    TensorView view_ = TensorView(),
    RandomGaussianFunc<Element> func_ = RandomGaussianFunc<Element>()
  ):
    view(view_), func(func_) {

  }

  void operator()(Coord<Layout::kRank> const &coord) const {
    view.at(coord) = func();
  }
};

} // namespace detail


template <
  typename Element,               
  typename Layout>                
void TensorFillRandomGaussian(
  TensorView<Element, Layout> dst,        
  uint64_t seed,                          
  double mean = 0,                        
  double stddev = 1,                      
  int bits = -1) {                        
  
  detail::RandomGaussianFunc<Element> random_func(seed, mean, stddev, bits);

  detail::TensorFillGaussianFunc<Element, Layout> func(
    dst,
    random_func
  );

  TensorForEach(
    dst.extent(),
    func
  );
}


template <
  typename Element                        
>
void BlockFillRandomGaussian(
  Element *ptr,                           
  size_t capacity,                        
  uint64_t seed,                          
  double mean = 0,                        
  double stddev = 1,                      
  int bits = -1) {                        
  

  detail::RandomGaussianFunc<Element> random_func(seed, mean, stddev, bits);

  for (size_t i = 0; i < capacity; ++i) {
    ptr[i] = random_func();
  }
}


namespace detail {

template <typename Element>
struct RandomUniformFunc {

  using Real = typename RealType<Element>::Type;
  
  uint64_t seed;
  double range;
  double min;
  int int_scale;

  //
  // Methods
  //

  RandomUniformFunc(
    uint64_t seed_ = 0, 
    double max = 1,
    double min_ = 0,
    int int_scale_ = -1
  ):
    seed(seed_), range(max - min_), min(min_), int_scale(int_scale_) {
      std::srand((unsigned)seed);
    }


  Element operator()() const {

    double rnd = double(std::rand()) / double(RAND_MAX);

    rnd = min + range * rnd;

    // Random values are cast to integer after scaling by a power of two to facilitate error
    // testing
    Element result;
    
    if (int_scale >= 0) {
      rnd = double(int64_t(rnd * double(1 << int_scale))) / double(1 << int_scale);
      result = static_cast<Element>(Real(rnd));
    }
    else {
      result = static_cast<Element>(Real(rnd));
    }

    return result;
  }
};

template <typename Element>
struct RandomUniformFunc<complex<Element> > {

  using Real = typename RealType<Element>::Type;
  
  uint64_t seed;
  double range;
  double min;
  int int_scale;

  //
  // Methods
  //

  RandomUniformFunc(
    uint64_t seed_ = 0, 
    double max = 1,
    double min_ = 0,
    int int_scale_ = -1
  ):
    seed(seed_), range(max - min_), min(min_), int_scale(int_scale_) {
      std::srand((unsigned)seed);
    }


  complex<Element> operator()() const {

    Element reals[2];

    for (int i = 0; i < 2; ++i) {
      double rnd = double(std::rand()) / double(RAND_MAX);

      rnd = min + range * rnd;

      // Random values are cast to integer after scaling by a power of two to facilitate error
      // testing
      
      if (int_scale >= 0) {
        rnd = double(int(rnd * double(1 << int_scale)));
        reals[i] = from_real<Element>(Real(rnd / double(1 << int_scale)));
      }
      else {
        reals[i] = from_real<Element>(Real(rnd));
      }
    }

    return complex<Element>(reals[0], reals[1]);
  }
};

template <
  typename Element,               
  typename Layout>                
struct TensorFillRandomUniformFunc {

  using TensorView = TensorView<Element, Layout>;

  //
  // Data members
  //

  TensorView view;
  RandomUniformFunc<Element> func;

  //
  // Methods
  //

  TensorFillRandomUniformFunc(
    TensorView view_ = TensorView(),
    RandomUniformFunc<Element> func_ = RandomUniformFunc<Element>()
  ):
    view(view_), func(func_) {

  }

  void operator()(Coord<Layout::kRank> const &coord) const {

    view.at(coord) = func();
  }
};

} // namespace detail


template <
  typename Element,               
  typename Layout>                
void TensorFillRandomUniform(
  TensorView<Element, Layout> dst,        
  uint64_t seed,                          
  double max = 1,                         
  double min = 0,                         
  int bits = -1) {                        
  detail::RandomUniformFunc<Element> random_func(seed, max, min, bits);

  detail::TensorFillRandomUniformFunc<Element, Layout> func(
    dst,
    random_func
  );

  TensorForEach(
    dst.extent(),
    func
  );
}


template <
  typename Element                        
>
void BlockFillRandomUniform(
  Element *ptr,
  size_t capacity,
  uint64_t seed,                          
  double max = 1,                         
  double min = 0,                         
  int bits = -1) {                        
  detail::RandomUniformFunc<Element> random_func(seed, max, min, bits);

  for (size_t i = 0; i < capacity; ++i) {
    ptr[i] = random_func();
  }
}


namespace detail {

template <
  typename Element,               
  typename Layout>                
struct TensorFillDiagonalFunc {

  using TensorView = TensorView<Element, Layout>;

  //
  // Data members
  //

  TensorView view;
  Element diag;
  Element other;

  //
  // Methods
  //

  TensorFillDiagonalFunc(
    TensorView const &view_ = TensorView(),
    Element diag_ = Element(1),
    Element other_ = Element(0)
  ):
    view(view_), diag(diag_), other(other_) { }

  void operator()(Coord<Layout::kRank> const & coord) const {
    bool is_diag = true;
    
    CUTLASS_PRAGMA_UNROLL
    for (int i = 1; i < Layout::kRank; ++i) {
      if (coord[i] != coord[i - 1]) {
        is_diag = false;
        break;
      }
    }

    view.at(coord) = (is_diag ? diag : other);
  }
};

} // namespace detail


template <
  typename Element,               
  typename Layout>                
void TensorFillDiagonal(
  TensorView<Element, Layout> dst,        
  Element diag = Element(1),              
  Element other = Element(0)) {           

  detail::TensorFillDiagonalFunc<Element, Layout> func(
    dst,
    diag,
    other
  );

  TensorForEach(
    dst.extent(),
    func
  );
}


template <
  typename Element,               
  typename Layout>                
void TensorFillIdentity(
  TensorView<Element, Layout> dst) {               

  TensorFillDiagonal(dst, Element(1), Element(0));
}


template <
  typename Element,               
  typename Layout>                
void TensorUpdateDiagonal(
  TensorView<Element, Layout> dst,                 
  Element val = Element(1)) {

  typename Layout::Index extent = dst.extent().min();

  for (typename Layout::Index i = 0; i < extent; ++i) {
    Coord<Layout::kRank> coord(i);
    dst.at(coord) = val;
  }
}


namespace detail {

template <
  typename Element,               
  typename Layout>                
struct TensorUpdateOffDiagonalFunc {

  using TensorView = TensorView<Element, Layout>;

  //
  // Data members
  //

  TensorView view;
  Element other;

  //
  // Methods
  //

  TensorUpdateOffDiagonalFunc(
    TensorView const &view_ = TensorView(),
    Element other_ = Element(0)
  ):
    view(view_), other(other_) { }

  void operator()(Coord<Layout::kRank> const & coord) const {
    bool is_diag = true;
    
    CUTLASS_PRAGMA_UNROLL
    for (int i = 1; i < Layout::kRank; ++i) {
      if (coord[i] != coord[i - 1]) {
        is_diag = false;
        break;
      }
    }

    if (!is_diag) {
      view.at(coord) = other;
    }
  }
};

} // namespace detail


template <
  typename Element,               
  typename Layout>                
void TensorUpdateOffDiagonal(
  TensorView<Element, Layout> dst,      
  Element other = Element(1)) {

  detail::TensorUpdateOffDiagonalFunc<Element, Layout> func(
    dst,
    other
  );

  TensorForEach(
    dst.extent(),
    func
  );
}



namespace detail {

template <
  typename Element,               
  typename Layout>                
struct TensorFillLinearFunc {

  using TensorView = TensorView<Element, Layout>;

  //
  // Data members
  //

  TensorView view;
  Array<Element, Layout::kRank> v;
  Element s;

  //
  // Methods
  //
  
  TensorFillLinearFunc() { }

  TensorFillLinearFunc(
    TensorView const &view_,
    Array<Element, Layout::kRank> const & v_,
    Element s_ = Element(0)
  ):
    view(view_), v(v_), s(s_) { }

  void operator()(Coord<Layout::kRank> const & coord) const {
    
    Element sum(s);

    CUTLASS_PRAGMA_UNROLL
    for (int i = 0; i < Layout::kRank; ++i) {
      sum += Element(coord[i]) * v[i];
    }

    view.at(coord) = sum;
  }
};

} // namespace detail


template <
  typename Element,               
  typename Layout>                
void TensorFillLinear(
  TensorView<Element, Layout> dst,      
  Array<Element, Layout::kRank> const & v,
  Element s = Element(0)) {

  detail::TensorFillLinearFunc<Element, Layout> func(
    dst,
    v,
    s
  );

  TensorForEach(
    dst.extent(),
    func
  );
}


template <
  typename Element,               
  typename Layout>                
void TensorFillSequential(
  TensorView<Element, Layout> dst,     
  Element s = Element(0)) {

  Array<Element, Layout::kRank> stride;

  stride[0] = Element(1);

  CUTLASS_PRAGMA_UNROLL
  for (int i = 1; i < Layout::kRank; ++i) {
    stride[i] = stride[i - 1] * Element(dst.extent()[i - 1]);
  }

  TensorFillLinear(dst, stride, s);
}


template <
  typename Element
>
void BlockFillSequential(
  Element *ptr,
  int64_t capacity,
  Element v = Element(1),
  Element s = Element(0)) {
  int i = 0;

  while (i < capacity) {
    cutlass::ReferenceFactory<Element, (cutlass::sizeof_bits<Element>::value <
                                        8)>::get(ptr, i) = s;

    s = Element(s + v);
    ++i;
  }
}


template <
  typename Element
>
void BlockFillRandom(
  Element *ptr,
  size_t capacity,
  uint64_t seed,
  Distribution dist) {

  if (dist.kind == Distribution::Gaussian) {
    BlockFillRandomGaussian<Element>(
      ptr, 
      capacity, 
      seed, 
      dist.gaussian.mean, 
      dist.gaussian.stddev, 
      dist.int_scale);
  }
  else if (dist.kind == Distribution::Uniform) {
    BlockFillRandomUniform<Element>(
      ptr, 
      capacity, 
      seed, 
      dist.uniform.max,
      dist.uniform.min, 
      dist.int_scale);
  }
}


template <
  typename Element,               
  typename Layout>                
void TensorCopyDiagonalIn(
  TensorView<Element, Layout> dst,          
  Element const *ptr) {                     

  typename Layout::Index extent = dst.extent().min();
  
  for (typename Layout::Index i = 0; i < extent; ++i) {
    Coord<Layout::kRank> coord(i);
    dst.at(coord) = ptr[i];
  }
}


template <
  typename Element,               
  typename Layout>                
void TensorCopyDiagonalOut(
  Element *ptr,                               
  TensorView<Element, Layout> src) {          

  typename Layout::Index extent = src.extent().min();
  
  for (typename Layout::Index i = 0; i < extent; ++i) {
    Coord<Layout::kRank> coord(i);
    ptr[i] = src.at(coord);
  }
}


} // namespace host
} // namespace reference
} // namespace cutlass