1 files changed, 182 insertions, 0 deletions
diff --git a/cutlass-example/cutlass/gemm/gemm_global_stream.h b/cutlass-example/cutlass/gemm/gemm_global_stream.h
new file mode 100644
index 0000000..ec675a3
--- /dev/null
+++ b/cutlass-example/cutlass/gemm/gemm_global_stream.h
@@ -0,0 +1,182 @@
+/***************************************************************************************************
+ * 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 Implements efficient loading of the thread block-level tile from global memory and
+   storing
+      to shared memory.
+*/
+#pragma once
+
+#include <cutlass/convert.h>
+#include <cutlass/gemm/gemm_global_tile.h>
+#include <cutlass/iterator_access.h>
+
+namespace cutlass {
+namespace gemm {
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <
+    /// The load iterator.
+    typename LoadIterator_,
+    /// The store iterator to copy to shared memory.
+    typename StoreIterator_,
+    /// The transformer to be applied after the data has been copied from global memory.
+    typename Transformer_>
+
+struct GlobalLoadStreamBase {
+  /// The load iterator.
+  typedef LoadIterator_ LoadIterator;
+  /// The transformer.
+  typedef Transformer_ Transformer;
+  /// The store iterator to write to shared memory.
+  typedef StoreIterator_ StoreIterator;
+
+  /// The fragment that is copied from shared memory.
+  typedef typename LoadIterator::Fragment FetchedFragment;
+  /// The fragment that is obtained after the transformation by the transformer.
+  typedef typename Transformer::OutputFragment TransformedFragment;
+  /// Make sure the fragments match.
+  static_assert((platform::is_same<FetchedFragment, typename Transformer::InputFragment>::value),
+                "");
+  /// The output fragment.
+  typedef TransformedFragment Fragment;
+  /// Make sure the transformed fragment is the same as the store fragment.
+  static_assert((platform::is_same<TransformedFragment, typename StoreIterator::Fragment>::value),
+                "");
+
+  /// The layout.
+  static MatrixLayout::Kind const kLayout = LoadIterator::kLayout;
+  /// The scalar type of the iterator.
+  typedef typename LoadIterator::Scalar Scalar;
+  /// The pointer.
+  typedef typename LoadIterator::Pointer Pointer;
+  /// The index.
+  typedef typename LoadIterator::Index Index;
+
+  /// The params.
+  struct Params {
+    // The load iterator.
+    typename LoadIterator::Params load_iterator;
+    // The store iterator.
+    typename StoreIterator::Params store_iterator;
+
+    /// Setup the params.
+    template <typename GemmDesc_>
+    CUTLASS_HOST_DEVICE int initialize(GemmDesc_ const& desc, Pointer pointer, Index ld) {
+      int error_code = load_iterator.initialize(desc, pointer, ld);
+      if (error_code) {
+        return error_code;
+      }
+
+      return store_iterator.initialize();
+    }
+  };
+
+  /// The amount of storage in shared memory needed to store the tile.
+  typedef typename StoreIterator::SharedStorage SharedStoreStorage;
+
+  /// The storage in shared memory needed by that stream.
+  union SharedStorage {
+    // The load iterator.
+    typename LoadIterator::SharedStorage load_iterator;
+    // The store iterator.
+    SharedStoreStorage store_iterator;
+  };
+
+  /// Ctor.
+  CUTLASS_DEVICE GlobalLoadStreamBase(Params const& params,
+                                      SharedStorage& shared_storage,
+                                      Coord<3> const bounds,
+                                      Coord<3> const& block)
+      : load_iterator(params.load_iterator, bounds, block),
+        transformer(),
+        store_iterator(params.store_iterator, shared_storage.store_iterator)
+
+  {
+    fetched_fragment.clear();
+  }
+
+  /// Load the data from shared memory to the fetch fragment.
+  CUTLASS_DEVICE void copy() { iterator_load(load_iterator, fetched_fragment); }
+
+  /// Commit the data.
+  CUTLASS_DEVICE void commit() {
+    transformer.transform(fetched_fragment, transformed_fragment);
+    iterator_store(store_iterator, transformed_fragment);
+    store_iterator.inc_stage();
+  }
+
+  /// Move to the beginning of the residue code. That's a new code path in CUTLASS 1.0.1.
+  CUTLASS_DEVICE void move_to_residue(Index k) { load_iterator.move_to_residue(k); }
+
+  /// Execute the residue code.
+  CUTLASS_DEVICE void residue(Index k, bool skip_clear = false) {
+    load_iterator.residue(k);
+    if (!skip_clear) {
+      fetched_fragment.clear();
+    }
+  }
+
+  /// Rollback to the beginning of the GEMM-k dimension.
+  CUTLASS_DEVICE void rollback() { load_iterator.rollback(); }
+
+  /// The iterator.
+  LoadIterator load_iterator;
+  /// The fragment to fetch from shared memory.
+  FetchedFragment fetched_fragment;
+  /// The transformer.
+  Transformer transformer;
+  /// The fragment to convert the data after it has been fetched from shared memory.
+  TransformedFragment transformed_fragment;
+  /// The store iterator.
+  StoreIterator store_iterator;
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <
+    /// The load iterator.
+    typename LoadIterator_,
+    /// The store iterator to copy to shared memory.
+    typename StoreIterator_,
+    /// The transformer to be applied after the data has been copied from global memory.
+    typename Transformer_ = Copy<typename LoadIterator_::Fragment> >
+
+struct GlobalLoadStream : public GlobalLoadStreamBase<LoadIterator_, StoreIterator_, Transformer_> {
+  /// The base class.
+  typedef GlobalLoadStreamBase<LoadIterator_, StoreIterator_, Transformer_> Base;
+
+  /// Ctor.
+  CUTLASS_DEVICE GlobalLoadStream(typename Base::Params const& params,
+                                  typename Base::SharedStorage& shared_storage,
+                                  Coord<3> const& bounds,
+                                  Coord<3> const& block)
+      : Base(params, shared_storage, bounds, block) {}
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+}  // namespace gemm
+}  // namespace cutlass