1 files changed, 305 insertions, 0 deletions
diff --git a/cutlass-example/cutlass/shape.h b/cutlass-example/cutlass/shape.h
new file mode 100644
index 0000000..4f6b222
--- /dev/null
+++ b/cutlass-example/cutlass/shape.h
@@ -0,0 +1,305 @@
+/***************************************************************************************************
+ * 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 Defines Shape implementing the Layout concept for representing a 4D hypercube of objects.
+*/
+#pragma once
+
+#include <cutlass/cutlass.h>
+
+namespace cutlass {
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+/*!@defgroup layout_concept Layout Concept
+* @{
+* @par Implementations of \ref layout_concept are used to describe a cube with DxHxW elements and C
+scalars per element.
+ A HxW slice of a cube is called an image and a cube consists of D images.
+*
+* @par Notations
+*   Let Layout be an implementation of the \ref layout_concept.
+*
+* @par Valid Expressions
+* - <b>Layout::D</b> specifies the depth of a cube
+* - <b>Layout::H</b> specifies the height of a cube
+* - <b>Layout::W</b> specifies the height of a cube
+* - <b>Layout::C</b> specifies the number of channels of each element in a cube
+* - <b>Layout::W_c</b> specifies the number of scalars of each row in one image of a cube.
+* - <b>Layout::H_w</b> specifies the number of elements in an image slice.
+* - <b>Layout::H_w_c</b>_specifies the number of scalars in an image slice.
+* - <b>Layout::D_h_w</b> specifies the number of elements in a cube.
+* - <b>Layout::D_h_w_c</b> specifies the number of scalars in a cube.
+* - <b>Layout::Strides</b> is a \ref layout_concept specifying the strides.
+* @}
+*/
+
+/**
+* @brief A Shape implementing \ref layout_concept describing the dimensions of a cube.
+* @concept{layout_concept}
+*/
+template <int kD_ = 1, int kH_ = 1, int kW_ = 1, int kC_ = 1>
+struct Shape {
+  /// The depth of the cube.
+  static int const kD = kD_;
+  /// The height of the cube.
+  static int const kH = kH_;
+  /// The width of the cube.
+  static int const kW = kW_;
+  /// The number of scalars per element.
+  static int const kC = kC_;
+};
+
+/**
+* @brief Compute derived counted of a \ref layout_concept based class
+*/
+template <typename Shape>
+struct ShapeCount {
+  /// The number of elements per row.
+  static int const kWc = Shape::kW * Shape::kC;
+  /// The number of pixels per image.
+  static int const kHw = Shape::kH * Shape::kW;
+  /// The number of elements per image.
+  static int const kHwc = Shape::kH * kWc;
+  /// The number of pixels per cube.
+  static int const kDhw = Shape::kD * kHw;
+  /// The number of elements in the 4D space.
+  static int const kDhwc = Shape::kD * kHwc;
+  /// The number of elements in the 4D space.
+  static int const kCount = kDhwc;
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <typename A_, int kScale_>
+struct ShapeScale {
+  typedef Shape<A_::kD * kScale_, A_::kH * kScale_, A_::kW * kScale_, A_::kC * kScale_> Shape;
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <typename A_, typename B_>
+struct ShapeAdd {
+  typedef Shape<A_::kD + B_::kD, A_::kH + B_::kH, A_::kW + B_::kW, A_::kC + B_::kC> Shape;
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <typename A_, typename B_>
+struct ShapeSub {
+  typedef Shape<A_::kD - B_::kD, A_::kH - B_::kH, A_::kW - B_::kW, A_::kC - B_::kC> Shape;
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <typename A_, typename B_>
+struct ShapeMul {
+  typedef Shape<A_::kD * B_::kD, A_::kH * B_::kH, A_::kW * B_::kW, A_::kC * B_::kC> Shape;
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <typename A_, typename B_>
+struct ShapeDiv {
+  typedef Shape<A_::kD / B_::kD, A_::kH / B_::kH, A_::kW / B_::kW, A_::kC / B_::kC> Shape;
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <typename A_, typename B_>
+struct ShapeMax {
+  typedef Shape<(A_::kD > B_::kD ? A_::kD : B_::kD),
+                (A_::kH > B_::kH ? A_::kH : B_::kH),
+                (A_::kW > B_::kW ? A_::kW : B_::kW),
+                (A_::kC > B_::kC ? A_::kC : B_::kC)>
+      Shape;
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <typename A_, typename B_>
+struct ShapeMin {
+  typedef Shape<(A_::kD < B_::kD ? A_::kD : B_::kD),
+                (A_::kH < B_::kH ? A_::kH : B_::kH),
+                (A_::kW < B_::kW ? A_::kW : B_::kW),
+                (A_::kC < B_::kC ? A_::kC : B_::kC)>
+      Shape;
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <typename Shape_, int kElementsPerAccess>
+struct ShapeStrides {
+  typedef Shape<Shape_::kH * Shape_::kW * Shape_::kC,
+                Shape_::kW * Shape_::kC,
+                Shape_::kC,
+                kElementsPerAccess>
+      Shape;
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+/**
+* @brief Compute the offset for the given coordinates in a cube
+* @tparam A \ref layout_concept where each dimension of the cube specifies the corresponding stride.
+*/
+template <typename Shape_>
+struct ComputeOffsetFromShape {
+  static CUTLASS_DEVICE int get(int d, int h, int w, int c) {
+    // clang-format off
+    return d * Shape_::kH * Shape_::kW * Shape_::kC +
+           h * Shape_::kW * Shape_::kC +
+           w * Shape_::kC +
+           c;
+    // clang-format on
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+/**
+* @brief Compute the offset for the given coordinates in a cube with a depth of 1
+* @tparam kSh Elements in the H dimension
+* @tparam kSw Elements in the W dimension
+* @tparam kSc Separation between two elements in "elements"
+*/
+template <int kSh_, int kSw_, int kSc_>
+struct ComputeOffsetFromShape<Shape<1, kSh_, kSw_, kSc_> > {
+  static CUTLASS_DEVICE int get(int d, int h, int w, int c) {
+    return h * kSw_ * kSc_ + w * kSc_ + c;
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+/**
+* @brief Compute the offset for the given coordinates in a cube with one channel and a depth of 1
+* @tparam kSh Elements in the H dimension
+* @tparam kSw Elements in the W dimension
+*/
+template <int kSh_, int kSw_>
+struct ComputeOffsetFromShape<Shape<1, kSh_, kSw_, 1> > {
+  static CUTLASS_DEVICE int get(int d, int h, int w, int c) { return h * kSw_ + w; }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+/**
+* @brief Compute the offset for the given coordinates in a cube
+* @tparam A \ref layout_concept where each dimension of the cube specifies the corresponding stride.
+*/
+template <typename Strides_>
+struct ComputeOffsetFromStrides {
+  static CUTLASS_DEVICE int get(int d, int h, int w, int c) {
+    return d * Strides_::kD + h * Strides_::kH + w * Strides_::kW + c * Strides_::kC;
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+/**
+* @brief Compute the offset for the given coordinates in a cube with a depth of 1
+* @tparam S_h Stride in the H dimension in scalars
+* @tparam S_w Stride in the W dimension in scalars
+* @tparam S_c Stride between two scalars.
+*/
+template <int S_h_, int S_w_, int S_c_>
+struct ComputeOffsetFromStrides<Shape<1, S_h_, S_w_, S_c_> > {
+  static CUTLASS_DEVICE int get(int d, int h, int w, int c) {
+    return h * S_h_ + w * S_w_ + c * S_c_;
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+/**
+* @brief Compute the offset for the given coordinates in a cube with one channel and a depth of 1
+* @tparam S_h Stride in the H dimension in scalars
+* @tparam S_w Stride in the W dimension in scalars
+*/
+template <int S_h_, int S_w_>
+struct ComputeOffsetFromStrides<Shape<1, S_h_, S_w_, 1> > {
+  static CUTLASS_DEVICE int get(int d, int h, int w, int c) { return h * S_h_ + w * S_w_; }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+/**
+* @brief Decompose threadId.x into coordinate of a cube whose dimensions are specified by Threads_.
+* Afterwards compute the offset of those coordinates using Strides_
+* @tparam Threads_ The dimension of the cube the threadIdx.x value is mapped on
+* @tparam Strides_ The strides to use when compute the offsets based on the coordinates of the cube.
+*/
+template <typename Threads_, typename Strides_>
+struct ComputeThreadOffsetFromStrides {
+  static CUTLASS_DEVICE int get() {
+    // Decompose the thread index.
+    int c = threadIdx.x % Threads_::kC;
+    int w = threadIdx.x / Threads_::kC % Threads_::kW;
+    int h = threadIdx.x / Threads_::kC / Threads_::kW % Threads_::kH;
+    int d = threadIdx.x / Threads_::kC / Threads_::kW / Threads_::kH;
+
+    // Compute the offset.
+    return d * Strides_::kD + h * Strides_::kH + w * Strides_::kW + c * Strides_::kC;
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+/**
+*@brief Specialization for D=1
+*/
+template <int T_h_, int T_w_, int T_c_, int S_h_, int S_w_, int S_c_>
+struct ComputeThreadOffsetFromStrides<Shape<1, T_h_, T_w_, T_c_>, Shape<1, S_h_, S_w_, S_c_> > {
+  static CUTLASS_DEVICE int get() {
+    // Decompose the thread index.
+    int c = threadIdx.x % T_c_;
+    int w = threadIdx.x / T_c_ % T_w_;
+    int h = threadIdx.x / T_c_ / T_w_ % T_h_;
+
+    // Compute the offset.
+    return h * S_h_ + w * S_w_ + c * S_c_;
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+/**
+*@brief Specialization for D=1 and C=1
+*/
+template <int T_h_, int T_w_, int S_h_, int S_w_>
+struct ComputeThreadOffsetFromStrides<Shape<1, T_h_, T_w_, 1>, Shape<1, S_h_, S_w_, 1> > {
+  static CUTLASS_DEVICE int get() {
+    // Decompose the thread index.
+    int w = threadIdx.x % T_w_;
+    int h = threadIdx.x / T_w_;
+
+    // Compute the offset.
+    return h * S_h_ + w * S_w_;
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+}  // namespace cutlass