streamk example and performance tuning (#760)

* streamk example and performance tuning * one missing file Co-authored-by: Haicheng Wu <haichengw@nvidia.com>
2023-01-10 16:10:02 -05:00
parent a1046d49c1
commit 764b840d6f
10 changed files with 1071 additions and 266 deletions
@@ -124,7 +124,7 @@ public:

    GemmUniversalMode mode;
    GemmCoord problem_size;
-    int batch_count;
+    int batch_count;        // Either (mode == GemmUniversalMode::kBatched) the batch count, or (mode == GemmUniversalMode::kGemm) the tile-splitting factor

    typename EpilogueOutputOp::Params epilogue;

@@ -148,7 +148,7 @@ public:
    typename LayoutC::Stride::LongIndex ldc;
    typename LayoutC::Stride::LongIndex ldd;

-    int sm_limit;    /// Carvout override: when the above are defaulted, the number of SMs that dispatch heuristics will attempt to load-balance
+    int avail_sms;          /// The number of SMs that StreamK dispatch heuristics will attempt to load-balance across (-1 defaults to device width, 1 implies classic data-parallel scheduling)


    //
@@ -159,15 +159,18 @@ public:
    Arguments():
      mode(GemmUniversalMode::kGemm),
      batch_count(1),
-      ptr_A(nullptr), ptr_B(nullptr), ptr_C(nullptr), ptr_D(nullptr),
-      sm_limit(-1)
+      ptr_A(nullptr),
+      ptr_B(nullptr),
+      ptr_C(nullptr),
+      ptr_D(nullptr),
+      avail_sms(-1)
    {}

    /// Constructor
    Arguments(
      GemmUniversalMode mode,
      GemmCoord problem_size,
-      int batch_count,
+      int batch_split,                              /// Either (mode == GemmUniversalMode::kBatched) the batch count, or (mode == GemmUniversalMode::kGemm) the tile-splitting factor (1 defaults to StreamK, >1 emulates Split-K)
      typename EpilogueOutputOp::Params epilogue,
      void const * ptr_A,
      void const * ptr_B,
@@ -181,15 +184,15 @@ public:
      typename LayoutB::Stride stride_b,
      typename LayoutC::Stride stride_c,
      typename LayoutC::Stride stride_d,
-      int sm_limit = -1                    /// Carvout override: when the above are defaulted, the number of SMs that dispatch heuristics will attempt to load-balance
+      int avail_sms = -1                            /// The number of SMs that StreamK dispatch heuristics will attempt to load-balance across (-1 defaults to device width, 1 implies classic data-parallel scheduling)
    ):
      mode(mode),
      problem_size(problem_size),
-      batch_count(batch_count),
+      batch_count(batch_split),
      epilogue(epilogue),
      ptr_A(ptr_A), ptr_B(ptr_B), ptr_C(ptr_C), ptr_D(ptr_D),
      batch_stride_A(batch_stride_A), batch_stride_B(batch_stride_B), batch_stride_C(batch_stride_C), batch_stride_D(batch_stride_D),
-      stride_a(stride_a), stride_b(stride_b), stride_c(stride_c), stride_d(stride_d), sm_limit(sm_limit)
+      stride_a(stride_a), stride_b(stride_b), stride_c(stride_c), stride_d(stride_d), avail_sms(avail_sms)
    {
      CUTLASS_TRACE_HOST("GemmUniversalStreamk::Arguments::Arguments() - problem_size: " << problem_size);
    }
@@ -198,7 +201,7 @@ public:
    Arguments(
      GemmUniversalMode mode,
      GemmCoord problem_size,
-      int batch_count,
+      int batch_split,                              /// Either (mode == GemmUniversalMode::kBatched) the batch count, or (mode == GemmUniversalMode::kGemm) the tile-splitting factor (1 defaults to StreamK, >1 emulates Split-K)
      typename EpilogueOutputOp::Params epilogue,
      void const * ptr_A,
      void const * ptr_B,
@@ -212,15 +215,15 @@ public:
      typename LayoutB::Stride::LongIndex ldb,
      typename LayoutC::Stride::LongIndex ldc,
      typename LayoutC::Stride::LongIndex ldd,
-      int sm_limit = -1                    /// Carvout override: when the above are defaulted, the number of SMs that dispatch heuristics will attempt to load-balance
+      int avail_sms = -1                            /// The number of SMs that StreamK dispatch heuristics will attempt to load-balance across (-1 defaults to device width, 1 implies classic data-parallel scheduling)
    ):
      mode(mode),
      problem_size(problem_size),
-      batch_count(batch_count),
+      batch_count(batch_split),
      epilogue(epilogue),
      ptr_A(ptr_A), ptr_B(ptr_B), ptr_C(ptr_C), ptr_D(ptr_D),
      batch_stride_A(batch_stride_A), batch_stride_B(batch_stride_B), batch_stride_C(batch_stride_C), batch_stride_D(batch_stride_D),
-      lda(lda), ldb(ldb), ldc(ldc), ldd(ldd), sm_limit(sm_limit)
+      lda(lda), ldb(ldb), ldc(ldc), ldd(ldd), avail_sms(avail_sms)
    {
      stride_a = make_Coord(lda);
      stride_b = make_Coord(ldb);
@@ -254,29 +257,36 @@ public:
    // Data members
    //

-    ThreadblockSwizzle block_mapping;
+    void * ptr_A;
+    void * ptr_B;

    typename Mma::IteratorA::Params params_A;
    typename Mma::IteratorB::Params params_B;
-    typename Epilogue::OutputTileIterator::Params params_C;
-    typename Epilogue::OutputTileIterator::Params params_D;
-    typename EpilogueOutputOp::Params output_op;
-
-    GemmUniversalMode mode;
-
-    void * ptr_A;
-    void * ptr_B;
-    void * ptr_C;
-    void * ptr_D;

    int64_t batch_stride_A;
    int64_t batch_stride_B;
-    int64_t batch_stride_C;
-    int64_t batch_stride_D;
+
+    GemmUniversalMode mode;
+
+    ThreadblockSwizzle block_mapping;
+
+    bool quick_dp;

    void *barrier_workspace;
    void *partials_workspace;

+    typename EpilogueOutputOp::Params output_op;
+
+    void * ptr_D;
+    void * ptr_C;
+
+    typename Epilogue::OutputTileIterator::Params params_D;
+    typename Epilogue::OutputTileIterator::Params params_C;
+
+    int64_t batch_stride_D;
+    int64_t batch_stride_C;
+
+
  protected:

    //
@@ -295,7 +305,7 @@ public:
    {
      // For atomic reduction, each SK-block needs a synchronization flag.  For parallel reduction,
      // each reduction block needs its own synchronization flag.
-      int sk_blocks = block_mapping.sk_regions * block_mapping.sk_blocks_per_region;
+      int sk_blocks = block_mapping.sk_regions() * block_mapping.sk_blocks_per_region();
      int num_flags = fast_max(sk_blocks, block_mapping.reduction_blocks);

      return cacheline_align_up(sizeof(typename Barrier::T) * num_flags);
@@ -304,7 +314,7 @@ public:
    /// Get the workspace size needed for intermediate partial sums
    size_t get_partials_workspace_size() const
    {
-      int sk_blocks = block_mapping.sk_regions * block_mapping.sk_blocks_per_region;
+      int sk_blocks = block_mapping.sk_regions() * block_mapping.sk_blocks_per_region();
      return cacheline_align_up(kWorkspaceBytesPerBlock * sk_blocks);
    }

@@ -343,9 +353,9 @@ public:
      partials_workspace(nullptr)
    {
      // Number of SMs to make available for StreamK decomposition
-      int avail_sms = (args.sm_limit == -1) ?
+      int avail_sms = (args.avail_sms == -1) ?
                        device_sms :
-                        fast_min(args.sm_limit, device_sms);
+                        fast_min(args.avail_sms, device_sms);

      // Initialize the block mapping structure
      block_mapping = ThreadblockSwizzle(
@@ -355,7 +365,15 @@ public:
        {ThreadblockShape::kM, ThreadblockShape::kN, ThreadblockShape::kK},
        args.batch_count,
        sm_occupancy,
+        device_sms,
        avail_sms);
+
+      quick_dp =
+        (block_mapping.sk_waves == 0) &&
+        (mode == GemmUniversalMode::kGemm) &&
+        !block_mapping.cohort_raster &&
+        !EpilogueOutputOp(output_op).is_source_needed();
+
    }


@@ -426,7 +444,7 @@ public:
    /// Returns the GEMM volume in thread block tiles
    cutlass::gemm::GemmCoord get_tiled_shape() const
    {
-      return block_mapping.tiled_shape;
+      return block_mapping.tiled_shape();
    }


@@ -533,9 +551,6 @@ protected:
  /// ID of each thread within a warp
  int lane_idx;

-  /// Block index
-  int block_idx;
-
  /// Threadblock scoped epilogue
  Epilogue epilogue;

@@ -640,16 +655,18 @@ protected:

  /// Iterator for fetching tile fragments from A
  CUTLASS_DEVICE
-  typename Mma::IteratorA init_iterator_A(TileWorkDesc &tile_work)
+  typename Mma::IteratorA init_iterator_A(
+    TileWorkDesc &tile_work,
+    GemmUniversalMode mode)
  {
    // The input A matrix
    ElementA *ptr_A = static_cast<ElementA *>(params.ptr_A);

    // Update input pointers based on batched/array mode
-    if (params.mode == GemmUniversalMode::kBatched) {
+    if (mode == GemmUniversalMode::kBatched) {
      ptr_A += tile_work.tiled_coord.k() * params.batch_stride_A;
    }
-    if (params.mode == GemmUniversalMode::kArray) {
+    if (mode == GemmUniversalMode::kArray) {
      ptr_A = static_cast<ElementA * const *>(params.ptr_A)[tile_work.tiled_coord.k()];
    }

@@ -667,16 +684,18 @@ protected:

  /// Iterator for fetching tile fragments from B
  CUTLASS_DEVICE
-  typename Mma::IteratorB init_iterator_B(TileWorkDesc &tile_work)
+  typename Mma::IteratorB init_iterator_B(
+    TileWorkDesc &tile_work,
+    GemmUniversalMode mode)
  {
    // The input B matrix
    ElementB *ptr_B = static_cast<ElementB *>(params.ptr_B);

    // Update input pointers based on batched/array mode
-    if (params.mode == GemmUniversalMode::kBatched) {
+    if (mode == GemmUniversalMode::kBatched) {
      ptr_B += tile_work.tiled_coord.k() * params.batch_stride_B;
    }
-    if (params.mode == GemmUniversalMode::kArray) {
+    if (mode == GemmUniversalMode::kArray) {
      ptr_B = static_cast<ElementB * const *>(params.ptr_B)[tile_work.tiled_coord.k()];
    }

@@ -700,10 +719,10 @@ protected:
    tile_work.tile_idx = tile_idx;

    // The first global-scoped MAC-iteration this threadblock will perform for this tile
-    tile_work.iter_begin = tile_idx * params.block_mapping.iters_per_tile;
+    tile_work.iter_begin = tile_idx * params.block_mapping.iters_per_tile();

    // The number of MAC-iterations this threadblock will perform for this tile
-    tile_work.k_iters_remaining = params.block_mapping.iters_per_tile;
+    tile_work.k_iters_remaining = params.block_mapping.iters_per_tile();

    // The starting index in the k-domain for MAC-iterations this threadblock will perform for this tile
    tile_work.k_begin = 0;
@@ -727,7 +746,7 @@ protected:
    tile_work.tile_idx = tile_idx;

    // The first global-scoped MAC-iteration for this tile
-    int tile_iter_begin = tile_idx * params.block_mapping.iters_per_tile;
+    int tile_iter_begin = tile_idx * params.block_mapping.iters_per_tile();

    // The first global-scoped MAC-iteration this threadblock will perform for this tile
    tile_work.iter_begin = max(block_iter_begin, tile_iter_begin);
@@ -756,7 +775,10 @@ protected:

  /// Share accumulators with peers
  CUTLASS_DEVICE
-  void share_accumulators(AccumulatorTile const &accumulator_tile, int first_block_idx)
+  void share_accumulators(
+    AccumulatorTile const &accumulator_tile,
+    int block_idx,
+    int first_block_idx)
  {
    AccumulatorTile *accum_tile_workspace = reinterpret_cast<AccumulatorTile *>(params.partials_workspace);

@@ -795,6 +817,7 @@ protected:
  CUTLASS_DEVICE
  void acquire_accumulators(
    AccumulatorTile &accumulator_tile,
+    int block_idx,
    int first_block_idx)
  {
    AccumulatorTile *accum_tile_workspace = reinterpret_cast<AccumulatorTile *>(params.partials_workspace);
@@ -868,8 +891,8 @@ protected:
    reduce_tile_idx = reduce_idx / Epilogue::kAccumulatorFragments;
    reduce_fragment_idx = reduce_idx % Epilogue::kAccumulatorFragments;

-    int iter_tile_first = reduce_tile_idx * params.block_mapping.iters_per_tile;
-    int iter_tile_last = iter_tile_first + params.block_mapping.iters_per_tile - 1;
+    int iter_tile_first = reduce_tile_idx * params.block_mapping.iters_per_tile();
+    int iter_tile_last = iter_tile_first + params.block_mapping.iters_per_tile() - 1;

    peer_idx_begin = params.block_mapping.get_sk_block_idx(iter_tile_first);
    peer_idx_last = params.block_mapping.get_sk_block_idx(iter_tile_last);
@@ -895,16 +918,6 @@ protected:
    ElementC *ptr_C = static_cast<ElementC *>(params.ptr_C);
    ElementC *ptr_D = static_cast<ElementC *>(params.ptr_D);

-    // Update pointers for batched/array mode(s)
-    if (params.mode == GemmUniversalMode::kBatched) {
-      ptr_C += tiled_coord.k() * params.batch_stride_C;
-      ptr_D += tiled_coord.k() * params.batch_stride_D;
-    }
-    if (params.mode == GemmUniversalMode::kArray) {
-      ptr_C = static_cast<ElementC * const *>(params.ptr_C)[tiled_coord.k()];
-      ptr_D = static_cast<ElementC * const *>(params.ptr_D)[tiled_coord.k()];
-    }
-
    // Tile iterator loading from source tensor.
    typename Epilogue::OutputTileIterator iterator_C(
        params.params_C,
@@ -936,12 +949,13 @@ protected:
  CUTLASS_DEVICE
  void process_tile(
    TileWorkDesc tile_work,
+    int block_idx,
    int dp_start_block_idx,
    int block_iter_begin)
  {
    // Initialize input iterators
-    typename Mma::IteratorA iterator_A = init_iterator_A(tile_work);
-    typename Mma::IteratorB iterator_B = init_iterator_B(tile_work);
+    typename Mma::IteratorA iterator_A = init_iterator_A(tile_work, params.mode);
+    typename Mma::IteratorB iterator_B = init_iterator_B(tile_work, params.mode);

    // Initialize accumulators
    AccumulatorTile accumulator_tile;
@@ -968,7 +982,7 @@ protected:

      if (!tile_work.tile_finished(params)) {
        // Non "finishing" SK blocks must share their partial accumulator sums through global scratch workspace
-        share_accumulators(accumulator_tile, first_block_idx);
+        share_accumulators(accumulator_tile, block_idx, first_block_idx);
      }
      else
      {
@@ -976,7 +990,7 @@ protected:
        if (!tile_work.tile_started())
        {
          // A "finishing" SK block must first aggregate its accumulator partial sums with those shared by peer threadblocks
-          acquire_accumulators(accumulator_tile, first_block_idx);
+          acquire_accumulators(accumulator_tile, block_idx, first_block_idx);
        }

        do_epilogue(tile_work, accumulator_tile);
@@ -1008,11 +1022,12 @@ protected:
    // Initialize block's iteration range
    int tile_idx, block_iter_begin, block_iters_remaining;

-    int sk_padding_start_block_idx =  params.block_mapping.sk_regions * params.block_mapping.sk_blocks_per_region;
+    int sk_padding_start_block_idx =  params.block_mapping.sk_regions() * params.block_mapping.sk_blocks_per_region();
    int dp_start_block_idx = params.block_mapping.sk_waves * params.block_mapping.avail_sms;
    int reduce_start_block_idx = dp_start_block_idx + params.block_mapping.dp_blocks;
    int grid_padding_start_block_idx = reduce_start_block_idx + params.block_mapping.reduction_blocks;

+    int block_idx = params.block_mapping.get_block_idx();
    if (block_idx < sk_padding_start_block_idx)
    {
      // This is a SK block
@@ -1044,8 +1059,9 @@ protected:
      }

      block_iter_begin = 0;
-      block_iters_remaining = params.block_mapping.iters_per_tile * tile_allottment;
+      block_iters_remaining = params.block_mapping.iters_per_tile() * tile_allottment;
    }
+
    else if ((ThreadblockSwizzle::kReductionStrategy == ThreadblockSwizzle::kMixed) &&
             (block_idx < grid_padding_start_block_idx))
    {
@@ -1072,8 +1088,8 @@ protected:

        // DP blocks exit if out of bounds or overlap an SK tile (only possible during cohort rasterization, where dp_first_wave_tiles must be 1)
        if ((tile_idx < params.block_mapping.sk_tiles) ||
-          (tile_work.tiled_coord.m() >= params.block_mapping.tiled_shape.m()) ||
-          (tile_work.tiled_coord.n() >= params.block_mapping.tiled_shape.n()))
+          (tile_work.tiled_coord.m() >= params.block_mapping.tiled_shape().m()) ||
+          (tile_work.tiled_coord.n() >= params.block_mapping.tiled_shape().n()))
        {
          break;
        }
@@ -1084,7 +1100,7 @@ protected:
      }

      // Perform this block's share of work for this tile
-      process_tile(tile_work, dp_start_block_idx, block_iter_begin);
+      process_tile(tile_work, block_idx, dp_start_block_idx, block_iter_begin);

      // Update remaining work for this block
      block_iters_remaining -= tile_work.k_iters_remaining;
@@ -1110,6 +1126,64 @@ protected:

  }

+
+  /// Executes one DP-only GEMM
+  CUTLASS_DEVICE
+  void gemm_dp()
+  {
+    int block_idx = blockIdx.x;
+    int tile_idx = block_idx;
+
+    TileWorkDesc tile_work;
+    tile_work.tile_idx = tile_idx;
+    tile_work.iter_begin = tile_idx * params.block_mapping.iters_per_tile();
+    tile_work.k_iters_remaining = params.block_mapping.iters_per_tile();
+    tile_work.k_begin = 0;
+    tile_work.k_end = params.block_mapping.problem_size.k();
+    tile_work.tiled_coord = params.block_mapping.get_tile_offset_row_major(tile_work.tile_idx);
+
+    // Initialize input iterators
+    typename Mma::IteratorA iterator_A = init_iterator_A(tile_work, params.mode);
+    typename Mma::IteratorB iterator_B = init_iterator_B(tile_work, params.mode);
+
+    // Initialize accumulators
+    AccumulatorTile accumulator_tile;
+    accumulator_tile.clear();
+
+    // Perform this tile's range of multiply-accumulate (MAC) iterations
+    Mma mma(
+      shared_storage.main_loop,
+      thread_idx,
+      warp_idx,
+      lane_idx);
+
+    mma(tile_work.k_iters_remaining, accumulator_tile, iterator_A, iterator_B, accumulator_tile);
+
+    ElementC *ptr_D = static_cast<ElementC *>(params.ptr_D);
+
+    // Location of this tile in item-coords
+    MatrixCoord threadblock_item_begin(
+      tile_work.tiled_coord.m() * Mma::Shape::kM,
+      tile_work.tiled_coord.n() * Mma::Shape::kN
+    );
+
+    // Tile iterator writing to destination tensor.
+    typename Epilogue::OutputTileIterator iterator_D(
+        params.params_D,
+        ptr_D,
+        params.block_mapping.problem_size.mn(),
+        thread_idx,
+        threadblock_item_begin);
+
+    // Execute the epilogue operator to update the destination tensor.
+    epilogue(
+        EpilogueOutputOp(params.output_op),
+        iterator_D,
+        accumulator_tile);
+  }
+
+
+
 public:

  //
@@ -1138,7 +1212,6 @@ public:
      thread_idx(threadIdx.x),
      warp_idx(__shfl_sync(0xffffffff, threadIdx.x / 32, 0)),   // broadcast the warp_id computed by lane 0 to ensure dependent code
      lane_idx(threadIdx.x % 32),
-      block_idx(params.block_mapping.get_block_idx()),
      epilogue(
        shared_storage.epilogue,
        thread_idx,
@@ -1151,7 +1224,17 @@ public:
  CUTLASS_DEVICE
  void operator()()
  {
-    // Do the GEMM
+#if (__CUDACC_VER_MAJOR__ > 10)
+    if (params.quick_dp)
+    {
+      // Simple (low-bootstrap latency) GEMM code path for data-parallel only.  (kBatched and kArray
+      // modes will only be launched using a data-parallel configurations)
+      gemm_dp();
+      return;
+    }
+#endif
+
+    // Generic SK code path
    gemm();

  }