adding the cutlass examples

1 files changed, 318 insertions, 0 deletions

diff --git a/cutlass-example/cutlass/iterator_access.h b/cutlass-example/cutlass/iterator_access.h
new file mode 100644
index 0000000..e94beb7
--- /dev/null
+++ b/cutlass-example/cutlass/iterator_access.h
@@ -0,0 +1,318 @@
+/***************************************************************************************************
+ * 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.
+ *
+ **************************************************************************************************/
+/*! \file
+    \brief Free functions for loading and storing to implementations of tile iteartor concepts.
+*/
+#pragma once
+
+#include <cutlass/fragment_load_store.h>
+#include <cutlass/load_store.h>
+#include <cutlass/predicate_vector.h>
+#include <cutlass/shape.h>
+
+namespace cutlass {
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Loads a fragment from an input iterator
+template <typename InputIterator, typename Fragment>
+CUTLASS_HOST_DEVICE void iterator_load(InputIterator &iterator, Fragment &fragment) {
+  typename InputIterator::FragmentIterator frag_iterator(fragment);
+  for (int d = 0; d < InputIterator::Iterations::kD; ++d) {
+    for (int h = 0; h < InputIterator::Iterations::kH; ++h) {
+      for (int w = 0; w < InputIterator::Iterations::kW; ++w) {
+        for (int c = 0; c < InputIterator::Iterations::kC; ++c) {
+          if (iterator.valid(d, h, w, c)) {
+            iterator.get(reinterpret_cast<typename InputIterator::AccessType &>(
+                             frag_iterator.at(d, h, w, c)),
+                         d,
+                         h,
+                         w,
+                         c);
+          }
+        }
+        if (w < InputIterator::Iterations::kW - 1) {
+          iterator.inc_w();
+        }
+      }
+      if (h < InputIterator::Iterations::kH - 1) {
+        iterator.inc_h();
+      }
+    }
+    if (d < InputIterator::Iterations::kD - 1) {
+      iterator.inc_d();
+    }
+  }
+  iterator.inc_advance();
+}
+
+/// Loads a fragment from a shared memory input iterator
+template <typename InputIterator, typename Fragment>
+CUTLASS_DEVICE void shared_iterator_load(InputIterator &iterator, Fragment &fragment) {
+  typename InputIterator::FragmentIterator frag_iterator(fragment);
+  for (int d = 0; d < InputIterator::Iterations::kD; ++d) {
+    for (int h = 0; h < InputIterator::Iterations::kH; ++h) {
+      for (int w = 0; w < InputIterator::Iterations::kW; ++w) {
+        for (int c = 0; c < InputIterator::Iterations::kC; ++c) {
+          int const offset =
+              ComputeOffsetFromStrides<typename InputIterator::ImmediateOffsetStrides>::get(
+                  d, h, w, c);
+
+          FragmentLoad<InputIterator::kIteratorFragment,
+                       InputIterator::Tile::kC,
+                       typename InputIterator::Scalar,
+                       InputIterator::kMemorySpace,
+                       typename InputIterator::FragmentElement,
+                       InputIterator::Tile::kW>::load(frag_iterator.at(d, h, w, c),
+                                                      iterator.data(),
+                                                      offset);
+        }
+      }
+    }
+  }
+}
+
+/// Loads a fragment from a shared memory input iterator
+template <typename InputIterator, typename Fragment>
+CUTLASS_DEVICE void shared_iterator_load(InputIterator &iterator, Fragment &fragment, int d) {
+  typename InputIterator::FragmentIterator frag_iterator(fragment);
+  for (int h = 0; h < InputIterator::Iterations::kH; ++h) {
+    for (int w = 0; w < InputIterator::Iterations::kW; ++w) {
+      for (int c = 0; c < InputIterator::Iterations::kC; ++c) {
+        int const offset =
+            ComputeOffsetFromStrides<typename InputIterator::ImmediateOffsetStrides>::get(
+                d, h, w, c);
+
+        FragmentLoad<InputIterator::kIteratorFragment,
+                     InputIterator::Tile::kC,
+                     typename InputIterator::Scalar,
+                     InputIterator::kMemorySpace,
+                     typename InputIterator::FragmentElement,
+                     InputIterator::Tile::kW>::load(frag_iterator.at(0, h, w, c),
+                                                    iterator.data(),
+                                                    offset);
+      }
+    }
+  }
+}
+
+/// Loads a fragment from an input iterator, masked by a predicate iterator
+template <typename InputIterator, typename Fragment, typename ConstPredicateAdapter>
+CUTLASS_HOST_DEVICE void iterator_load_post_increment(InputIterator &iterator,
+                                                      Fragment &fragment,
+                                                      typename InputIterator::Index offset,
+                                                      ConstPredicateAdapter predicate_adapter) {
+  for (int d = 0; d < InputIterator::Iterations::kD; ++d, iterator.inc_d()) {
+    for (int h = 0; h < InputIterator::Iterations::kH; ++h, iterator.inc_h()) {
+      for (int w = 0; w < InputIterator::Iterations::kW; ++w, iterator.inc_w()) {
+        if (predicate_adapter.at(d, h, w, 0)) {
+          int idx = InputIterator::Tile::kC *
+                    (w + InputIterator::Iterations::kW * (h + InputIterator::Iterations::kH * d));
+
+          Load<typename Fragment::Element, InputIterator::Tile::kC, InputIterator::kMemorySpace>::
+              load(reinterpret_cast<typename InputIterator::AccessType &>(fragment[idx]),
+                   iterator.data(),
+                   offset);
+        }
+      }
+    }
+  }
+}
+
+/// Loads a fragment from an input iterator
+template <typename InputIterator, typename Fragment>
+CUTLASS_HOST_DEVICE void iterator_load_post_increment(InputIterator &iterator,
+                                                      Fragment &fragment,
+                                                      typename InputIterator::Index offset = 0) {
+  TrivialPredicateTileAdapter pred;
+  iterator_load_post_increment(iterator, fragment, offset, pred);
+}
+
+/// Loads a fragment from an input iterator
+template <typename InputIterator, typename Fragment, typename ConstPredicateAdapter>
+CUTLASS_HOST_DEVICE void iterator_load_post_increment(InputIterator &iterator,
+                                                      Fragment &fragment,
+                                                      ConstPredicateAdapter pred_it) {
+  iterator_load_post_increment(iterator, fragment, 0, pred_it);
+}
+
+template <typename InputIterator, typename Fragment, typename ConstPredicateAdapter>
+CUTLASS_HOST_DEVICE void iterator_load(InputIterator const &_iterator,
+                                       Fragment &fragment,
+                                       typename InputIterator::Index offset,
+                                       ConstPredicateAdapter predicate_adapter) {
+  InputIterator iterator(_iterator);
+  iterator_load_post_increment(iterator, fragment, offset, predicate_adapter);
+}
+
+/// Loads a fragment from an input iterator
+template <typename InputIterator, typename Fragment>
+CUTLASS_HOST_DEVICE void iterator_load(InputIterator const &iterator,
+                                       Fragment &fragment,
+                                       typename InputIterator::Index offset = 0) {
+  TrivialPredicateTileAdapter pred;
+  iterator_load(iterator, fragment, offset, pred);
+}
+
+/// Loads a fragment from an input iterator
+template <typename InputIterator, typename Fragment, typename ConstPredicateAdapter>
+CUTLASS_HOST_DEVICE void iterator_load(InputIterator const &iterator,
+                                       Fragment &fragment,
+                                       ConstPredicateAdapter pred_it) {
+  iterator_load(iterator, fragment, 0, pred_it);
+}
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Stores a fragment to an output iterator
+template <typename OutputIterator, typename Fragment>
+CUTLASS_HOST_DEVICE void iterator_store(OutputIterator &iterator, Fragment &fragment) {
+  typename OutputIterator::FragmentIterator frag_iterator(fragment);
+  for (int d = 0; d < OutputIterator::Iterations::kD; ++d) {
+    for (int h = 0; h < OutputIterator::Iterations::kH; ++h) {
+      for (int w = 0; w < OutputIterator::Iterations::kW; ++w) {
+        if (iterator.valid(d, h, w, 0)) {
+          iterator.set(reinterpret_cast<typename OutputIterator::AccessType const &>(
+                           frag_iterator.at(d, h, w, 0)),
+                       d,
+                       h,
+                       w,
+                       0);
+        }
+        if (w < OutputIterator::Iterations::kW - 1) {
+          iterator.inc_w();
+        }
+      }
+      if (h < OutputIterator::Iterations::kH - 1) {
+        iterator.inc_h();
+      }
+    }
+    if (d < OutputIterator::Iterations::kD - 1) {
+      iterator.inc_d();
+    }
+  }
+  iterator.inc_advance();
+}
+
+/// Stores a fragment to a shared memory output iterator
+template <typename OutputIterator, typename Fragment>
+CUTLASS_DEVICE void shared_iterator_store(OutputIterator &iterator, Fragment const &fragment) {
+  typename OutputIterator::FragmentConstIterator frag_iterator(fragment);
+  for (int d = 0; d < OutputIterator::Iterations::kD; ++d) {
+    for (int h = 0; h < OutputIterator::Iterations::kH; ++h) {
+      for (int w = 0; w < OutputIterator::Iterations::kW; ++w) {
+        for (int c = 0; c < OutputIterator::Iterations::kC; ++c) {
+          int const offset =
+              ComputeOffsetFromStrides<typename OutputIterator::ImmediateOffsetStrides>::get(
+                  d, h, w, c);
+
+          FragmentStore<OutputIterator::kIteratorFragment,
+                        OutputIterator::Tile::kC,
+                        typename OutputIterator::Scalar,
+                        OutputIterator::kMemorySpace,
+                        typename OutputIterator::FragmentElement,
+                        OutputIterator::Tile::kW>::store(frag_iterator.at(d, h, w, c),
+                                                         iterator.data(),
+                                                         offset);
+        }
+      }
+    }
+  }
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Stores a fragment to an output iterator, masked by a predicate iterator
+template <typename OutputIterator, typename Fragment, typename ConstPredicateAdapter>
+CUTLASS_HOST_DEVICE void iterator_store_post_increment(OutputIterator &iterator,
+                                                       Fragment const &fragment,
+                                                       typename OutputIterator::Index offset,
+                                                       ConstPredicateAdapter predicate_adapter) {
+  for (int d = 0; d < OutputIterator::Iterations::kD; ++d, iterator.inc_d()) {
+    for (int h = 0; h < OutputIterator::Iterations::kH; ++h, iterator.inc_h()) {
+      for (int w = 0; w < OutputIterator::Iterations::kW; ++w, iterator.inc_w()) {
+        if (predicate_adapter.at(d, h, w, 0)) {
+          int idx = OutputIterator::Tile::kC *
+                    (w + OutputIterator::Iterations::kW * (h + OutputIterator::Iterations::kH * d));
+
+          Store<typename Fragment::Element,
+                OutputIterator::Tile::kC,
+                OutputIterator::kMemorySpace>::
+              store(reinterpret_cast<typename OutputIterator::AccessType const &>(fragment[idx]),
+                    iterator.data(),
+                    offset);
+        }
+      }
+    }
+  }
+}
+
+/// Stores a fragment to an output iterator
+template <typename OutputIterator, typename Fragment>
+CUTLASS_HOST_DEVICE void iterator_store_post_increment(OutputIterator &iterator,
+                                                       Fragment const &fragment,
+                                                       typename OutputIterator::Index offset = 0) {
+  TrivialPredicateTileAdapter pred;
+  iterator_store_post_increment(iterator, fragment, offset, pred);
+}
+
+/// Stores a fragment to an output iterator
+template <typename OutputIterator, typename Fragment, typename ConstPredicateAdapter>
+CUTLASS_HOST_DEVICE void iterator_store_post_increment(OutputIterator &iterator,
+                                                       Fragment const &fragment,
+                                                       ConstPredicateAdapter pred_it) {
+  iterator_store_post_increment(iterator, fragment, 0, pred_it);
+}
+
+/// Stores a fragment to an output iterator, masked by a predicate iterator
+template <typename OutputIterator, typename Fragment, typename ConstPredicateAdapter>
+CUTLASS_HOST_DEVICE void iterator_store(OutputIterator const &_iterator,
+                                        Fragment const &fragment,
+                                        typename OutputIterator::Index offset,
+                                        ConstPredicateAdapter predicate_adapter) {
+  OutputIterator iterator(_iterator);
+  iterator_store_post_increment(iterator, fragment, offset, predicate_adapter);
+}
+
+/// Stores a fragment to an output iterator
+template <typename OutputIterator, typename Fragment>
+CUTLASS_HOST_DEVICE void iterator_store(OutputIterator const &iterator,
+                                        Fragment const &fragment,
+                                        typename OutputIterator::Index offset = 0) {
+  TrivialPredicateTileAdapter pred;
+  iterator_store(iterator, fragment, offset, pred);
+}
+
+/// Stores a fragment to an output iterator
+template <typename OutputIterator, typename Fragment, typename ConstPredicateAdapter>
+CUTLASS_HOST_DEVICE void iterator_store(OutputIterator const &iterator,
+                                        Fragment const &fragment,
+                                        ConstPredicateAdapter pred_it) {
+  iterator_store(iterator, fragment, 0, pred_it);
+}
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+}  // namespace cutlass