tensor_ref_collection.h

Go to the documentation of this file.

/***************************************************************************************************
 * Copyright (c) 2017-2018, 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.
 *
 **************************************************************************************************/
#pragma once

#include "cutlass/tensor_ref.h"

namespace cutlass {

//
// TensorRefCollection is a concept for storing a logical collection of TensorRef objects. Classes
// satisfying the TensorRefCollection concept must support the following:
//
//   // Define storage type
//   typedef typename TensorRefCollection::Storage Storage;
//
//   // Define a type for offsets in memory
//   typedef typename TensorRefCollection::LongIndex LongIndex;
//
//   // Define a ConstIterator type satisfying TensorRefIterator
//   typedef typename TensorRefCollection::ConstIterator TensorRefIterator;
//
//   // Implement a begin() method.
//   TensorRefIterator iterator = collection.begin();
//
//
// TensorRefIterator is a concept for accessing an element in a TensorRefCollection. Classes
// satisfying the TensorRefIterator concept must support the following:
//
//   // Define a TensorRef type accessed by the iterator
//   typedef typename TensorRefIterator::TensorRef TensorRef;
//
//   // Access the TensorRef
//   TensorRef ref = *iterator;
//
//   // Pre-increment and post-increment
//   ++iterator;
//   iterator++;
//
//   // Pre-decrement and post-decrement
//   --iterator;
//   iterator--;
//

template <
  typename Storage_,
  int Rank_,
  typename MapFunc_ = IdentityTensorMapFunc<Rank_>,
  int StorageRank_ = MapFunc_::kStorageRank,
  typename Index_ = int,
  typename LongIndex_ = long long
>
struct TensorRefBatchStrided: 
  public TensorRef<Storage_, Rank_, MapFunc_, StorageRank_, Index_, LongIndex_> {

  //
  // Type definitions
  //

  typedef TensorRef<Storage_, Rank_, MapFunc_, StorageRank_, Index_, LongIndex_> Base;

  typedef typename Base::Storage Storage;

  typedef Index_ Index;

  typedef LongIndex_ LongIndex;

  typedef Coord<kRank> TensorCoord;

  typedef Base TensorRef;

  class ConstIterator {
  public:
    typedef Base TensorRef;

  private:

    TensorRefBatchStrided const &ref_;

    LongIndex offset_;

  public:

    CUTLASS_HOST_DEVICE
    ConstIterator(
      TensorRefBatchStrided const &ref, 
      LongIndex offset = 0): ref_(ref), offset_(offset) { }

    CUTLASS_HOST_DEVICE
    TensorRef *operator() const {
      TensorRef ref(ref_);
      ref.add_pointer_offset(offset_);
      return ref;
    }

    CUTLASS_HOST_DEVICE
    ConstIterator &operator++() {
      offset_ += ref_.tensor_stride;
      return *this;
    }

    CUTLASS_HOST_DEVICE
    ConstIterator operator++(int) {
      ConstIterator ret(*this);
      offset_ += ref_.tensor_stride;
      return ret;
    }

    CUTLASS_HOST_DEVICE
    ConstIterator operator+(Index idx) {
      return ConstIterator(ref, offset_ + ref_.tensor_stride * idx);
    }

    CUTLASS_HOST_DEVICE
    ConstIterator &operator+=(Index idx) {
      offset_ += ref_.tensor_stride * idx;
      return *this;
    }

    CUTLASS_HOST_DEVICE
    ConstIterator &operator--() {
      offset_ -= ref_.tensor_stride;
      return *this;
    }

    CUTLASS_HOST_DEVICE
    ConstIterator operator--(int) {
      ConstIterator ret(*this);
      offset_ -= ref_.tensor_stride;
      return ret;
    }

    CUTLASS_HOST_DEVICE
    ConstIterator operator-(Index idx) {
      return ConstIterator(ref_, offset_ - ref_.tensor_stride * idx);
    }

    CUTLASS_HOST_DEVICE
    ConstIterator &operator-=(Index idx) {
      offset_ -= ref_.tensor_stride * idx;
      return *this;
    }

    CUTLASS_HOST_DEVICE
    Stride operator-(ConstIterator const &it) {
      return offset_ - it.offset_;
    }
  };

  //
  // Data members
  //

  LongIndex tensor_stride;

  //
  // Methods
  //

  // Default ctor
  CUTLASS_HOST_DEVICE
  TensorRefBatchStrided(): tensor_stride(0) { }

  // Constructs form a tensor reference and 
  CUTLASS_HOST_DEVICE
  TensorRefBatchStrided(TensorRef const &ref, LongIndex _tensor_stride = 0): 
    TensorRef(ref), 
    tensor_stride(_tensor_stride) { }

  CUTLASS_HOST_DEVICE
  LongIndex get_pointer_offset(Index idx) const {
    return idx * tensor_stride;
  }

  // Returns a reference
  CUTLASS_HOST_DEVICE
  TensorRef at(Index idx) const {
    TensorRef ref(*this);
    ref.add_pointer_offset(get_pointer_offset(idx));
    return ref;
  }

  CUTLASS_HOST_DEVICE
  ConstIterator begin() {
    return ConstIterator(*this);
  }
};


template <
  typename Storage_,
  int Rank_,
  typename MapFunc_ = IdentityTensorMapFunc<Rank_>,
  int StorageRank_ = MapFunc_::kStorageRank,
  typename Index_ = int,
  typename LongIndex_ = long long
>
struct TensorRefArray {
  //
  // Type definitions
  //

  typedef TensorRef<Storage_, Rank_, MapFunc_, StorageRank_, Index_, LongIndex_> TensorRef;

  typedef Storage_ Storage;

  typedef Index_ Index;

  typedef LongIndex_ LongIndex;

  static int const kStorageRank = TensorRef::kStorageRank;

  class ConstIterator {
  public:

    typedef Base TensorRef;

  private:
    TensorRefArray const &ref_;

    int idx_;

  public:

    CUTLASS_HOST_DEVICE
    ConstIterator(TensorArrayRef const &ref, int idx = 0): ref_(ref), idx_(idx) { }

    CUTLASS_HOST_DEVICE
    TensorRef *operator() const {
      return ref_.reference(idx_);
    }

    CUTLASS_HOST_DEVICE
    ConstIterator &operator++() {
      ++idx_;
      return *this;
    }

    CUTLASS_HOST_DEVICE
    ConstIterator operator++(int) {
      ConstIterator ret(*this);
      idx_ ++;
      return ret;
    }

    CUTLASS_HOST_DEVICE
    ConstIterator operator+(Index idx) {
      return ConstIterator(ref_, idx_ + idx);
    }

    CUTLASS_HOST_DEVICE
    ConstIterator &operator+=(Index idx) {
      idx_ += idx;
      return *this;
    }

    CUTLASS_HOST_DEVICE
    ConstIterator &operator--() {
      --idx_;
      return *this;
    }

    CUTLASS_HOST_DEVICE
    ConstIterator operator--(int) {
      ConstIterator ret(*this);
      --idx_;
      return ret;
    }

    CUTLASS_HOST_DEVICE
    ConstIterator &operator-=(Index idx) {
      idx_ -= idx;
      return *this;
    }

    CUTLASS_HOST_DEVICE
    ConstIterator operator-(Index idx) {
      return ConstIterator(ref_, idx_ + idx);
    }
  };

  //
  // Data members
  //

  Storage **pointers;

  Index *strides[kStorageRank - 1];

  //
  // Methods
  //

  // Default ctor
  CUTLASS_HOST_DEVICE
  TensorArrayRef() { }

  // Construct from pointers to arrays to strides
  CUTLASS_HOST_DEVICE
  TensorArrayRef(
    Storage **_pointers,
    Index _strides[kStorageRank - 1]): pointers(_pointers) { 

    // Copy pointers to strides arrays
    for (int i = 0; i < kStorageRank - 1; ++i) {
      strides[i] = _strides[i];
    }
  }

  // Returns a TensorRef at the given index in the collection
  CUTLASS_HOST_DEVICE
  TensorRef at(Index idx) const {
    Coord<kStorageRank - 1, Index> stride;
    CUTLASS_PRAGMA_UNROLL
    for (int i = 0; i < kStorageRank - 1; ++i) {
      stride[i] = stride_[idx][i];
    }
    return TensorRef(pointers[idx], stride);
  }

  CUTLASS_HOST_DEVICE
  ConstIterator begin() {
    return ConstIterator(*this);
  }
};


} // namespace cutlass