CUTLASS 2.5

2021-02-26 09:58:26 -05:00
parent ccb697bac7
commit 0e13748649
771 changed files with 15474 additions and 1715 deletions
@@ -1,5 +1,5 @@
 /***************************************************************************************************
- * Copyright (c) 2017-2020, NVIDIA CORPORATION.  All rights reserved.
+ * Copyright (c) 2017-2021, 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:
@@ -30,9 +30,12 @@
 #else
 #include <cstdint>
 #include <cmath>
+#include <type_traits>
 #endif

 #include "cutlass/cutlass.h"
+#include "cutlass/uint128.h"
+#include "cutlass/coord.h"

 /**
 * \file
@@ -151,6 +154,7 @@ constexpr int ceil_div(int a, int b) {
  return (a + b - 1) / b;
 }

+/////////////////////////////////////////////////////////////////////////////////////////////////

 /**
 * log2 computation, what's the
@@ -221,6 +225,8 @@ void fast_divmod(int& quo, int64_t& rem, int64_t src, int div, unsigned int mul,
  rem = src - (quo * div);
 }

+/////////////////////////////////////////////////////////////////////////////////////////////////
+
 /// Object to encapsulate the fast division+modulus operation.
 ///
 /// This object precomputes two values used to accelerate the computation and is best used
@@ -272,9 +278,159 @@ struct FastDivmod {
  }
 };

-/******************************************************************************
- * Min/Max
- ******************************************************************************/
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Object to encapsulate the fast division+modulus operation for 64b integer division.
+///
+/// This object precomputes two values used to accelerate the computation and is best used
+/// when the divisor is a grid-invariant. In this case, it may be computed in host code and
+/// marshalled along other kernel arguments using the 'Params' pattern.
+///
+/// Example:
+///
+///
+///   uint64_t quotient, remainder, dividend, divisor;
+///
+///   FastDivmodU64 divmod(divisor);
+///
+///   divmod(quotient, remainder, dividend);  
+///
+///   // quotient = (dividend / divisor)
+///   // remainder = (dividend % divisor)
+///
+struct FastDivmodU64 {
+
+  uint64_t divisor;
+  uint64_t multiplier;
+  unsigned int shift_right;
+  unsigned int round_up;
+
+  //
+  // Static methods
+  //
+
+  /// Computes b, where 2^b is the greatest power of two that is less than or equal to x
+  CUTLASS_HOST_DEVICE
+  static uint32_t integer_log2(uint64_t x) {
+    uint32_t n = 0;
+    while (x >>= 1) {
+      ++n;
+    }
+    return n;
+  }
+
+  /// Default ctor
+  CUTLASS_HOST_DEVICE
+  FastDivmodU64(): divisor(0), multiplier(0), shift_right(0), round_up(0) { }
+
+  /// Construct the FastDivmod object, in host code ideally.
+  ///
+  /// This precomputes some values based on the divisor and is computationally expensive.
+  CUTLASS_HOST_DEVICE
+  FastDivmodU64(uint64_t divisor_): divisor(divisor_), multiplier(1), shift_right(0), round_up(0) {
+
+    if (divisor) {
+      shift_right = integer_log2(divisor);
+
+      if ((divisor & (divisor - 1)) == 0) {
+        multiplier = 0;
+      }
+      else {
+        uint64_t power_of_two = (uint64_t(1) << shift_right);
+        uint64_t multiplier_lo = uint128_t(0, power_of_two) / divisor;
+        multiplier = uint128_t(power_of_two, power_of_two) / divisor;
+        round_up = (multiplier_lo == multiplier ? 1 : 0);
+      }
+    }
+  }
+
+  /// Returns the quotient of floor(dividend / divisor)
+  CUTLASS_HOST_DEVICE
+  uint64_t divide(uint64_t dividend) const {
+    uint64_t quotient = 0;
+
+    #ifdef __CUDA_ARCH__
+      uint64_t x = dividend;
+      if (multiplier) {
+        x = __umul64hi(dividend + round_up, multiplier);
+      }
+      quotient = (x >> shift_right);
+    #else
+      // TODO - use proper 'fast' division here also. No reason why x86-code shouldn't be optimized.
+      quotient = dividend / divisor;
+    #endif
+
+    return quotient;
+  }
+
+  /// Computes the remainder given a computed quotient and dividend
+  CUTLASS_HOST_DEVICE
+  uint64_t modulus(uint64_t quotient, uint64_t dividend) const {
+    return uint32_t(dividend - quotient * divisor);
+  }
+
+  /// Returns the quotient of floor(dividend / divisor) and computes the remainder
+  CUTLASS_HOST_DEVICE
+  uint64_t divmod(uint64_t &remainder, uint64_t dividend) const {
+    uint64_t quotient = divide(dividend);
+    remainder = modulus(quotient, dividend);
+    return quotient;
+  }
+
+  /// Computes integer division and modulus using precomputed values. This is computationally
+  /// inexpensive.
+  CUTLASS_HOST_DEVICE
+  void operator()(uint64_t &quotient, uint64_t &remainder, uint64_t dividend) const {
+    quotient = divmod(remainder, dividend);
+  }
+};
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Computes the coordinate decomposition from a linear index.
+///
+/// This decomposition is accelerated by the FastDivmodU64 object. It is assumed that
+/// a coordinate of <Rank> indices can be decomposed by <Rank - 1> div/mod operations.
+/// Note, is assumed that element divmod[0] divides by extent[1].
+///
+/// For example, assume 4-D coordinate (n, p, q, c) is mapped to a linear index `npqc`. This
+/// can be decomposed via three divide and modulus operations:
+///
+///      c = npqc % C;         |  divmod[2] = FastDivmodU64(C)
+///    npq = npqc / C;         |   coord[3] = c
+///
+///      q =  npq % Q;         |  divmod[1] = FastDivmodU64(Q)
+///     np =  npq / Q;         |   coord[2] = q
+///
+///      p =   np % P;         |  divmod[0] = FastDivmodU64(P)
+///      n =   np / P;         |   coord[1] = p
+///
+///                            |   coord[0] = n
+///
+template <int Rank>
+CUTLASS_HOST_DEVICE Coord<Rank> CoordinateDecomposition(
+  uint64_t linear_idx,                    ///< Linear index to decompose
+  FastDivmodU64 const *divmod) {          ///< Pointer to array of Rank-1 FastDivmodU64 objects
+
+  static_assert(Rank > 0, "CoordinateDecomposition requires Rank=1 or greater.");
+
+  Coord<Rank> coord;
+
+  CUTLASS_PRAGMA_UNROLL
+  for (int i = Rank; i > 1; --i) {
+    uint64_t remainder;
+    linear_idx = divmod[i - 2].divmod(remainder, linear_idx);
+    coord[i - 1] = int(remainder);
+  }
+
+  coord[0] = int(linear_idx);
+
+  return coord;
+}
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+// Min/Max
+/////////////////////////////////////////////////////////////////////////////////////////////////

 template <int A, int B>
 struct Min {
@@ -296,6 +452,30 @@ constexpr int const_max(int a, int b) {
    return (b > a ? b : a);
 }

+template <typename T>
+CUTLASS_HOST_DEVICE
+T fast_min(T a, T b) {
+  return (b < a ? b : a);
+}
+
+template <>
+CUTLASS_HOST_DEVICE
+float fast_min(float a, float b) {
+  return fminf(a, b);
+}
+
+template <typename T>
+CUTLASS_HOST_DEVICE
+T fast_max(T a, T b) {
+  return (a < b ? b : a);
+}
+
+template <>
+CUTLASS_HOST_DEVICE
+float fast_max(float a, float b) {
+  return fmaxf(a, b);
+}
+
 CUTLASS_HOST_DEVICE
 float fast_cos(float theta) {
  #if defined(__CUDA_ARCH__)
@@ -404,6 +584,24 @@ double fast_log(double x) {
  #endif
 }

+CUTLASS_HOST_DEVICE
+float fast_tanh(float x) {
+  #if defined(__CUDA_ARCH__)
+  return ::tanhf(x);
+  #else
+  return std::tanh(x);
+  #endif
+}
+
+CUTLASS_HOST_DEVICE
+double fast_tanh(double x) {
+  #if defined(__CUDA_ARCH__)
+  return ::tanh(x);
+  #else
+  return std::tanh(x);
+  #endif
+}
+
 /////////////////////////////////////////////////////////////////////////////////////////////////

 }  // namespace cutlass