custom_accum.h

Go to the documentation of this file.

/*
 * SPDX-FileCopyrightText: Copyright (c) 2018-2021, NVIDIA CORPORATION.
 * SPDX-License-Identifier: Apache-2.0
 */
 
#pragma once
 
#include <cutlass/cutlass.h>
#include <cutlass/fragment.h>
 
namespace MLCommon {
namespace LinAlg {
 
template <typename AccumulatorsPerThread_,
          typename ThreadsPerWarp_,
          typename ScalarA_,
          typename ScalarB_,
          typename ScalarC_>
struct ThreadDiffSquaredAdd {
  typedef cutlass::Shape<1, 1, 1, 1> InstructionShape;
  typedef AccumulatorsPerThread_ AccumulatorsPerThread;
  typedef ThreadsPerWarp_ ThreadsPerWarp;
  typedef
    typename cutlass::ShapeMul<AccumulatorsPerThread, ThreadsPerWarp>::Shape AccumulatorsPerWarp;
  typedef ScalarA_ ScalarA;
  typedef cutlass::Fragment<ScalarA, AccumulatorsPerThread::kW> FragmentA;
  typedef ScalarB_ ScalarB;
  typedef cutlass::Fragment<ScalarB, AccumulatorsPerThread::kH> FragmentB;
  typedef ScalarC_ ScalarC;
  typedef cutlass::Fragment<ScalarC, AccumulatorsPerThread::kH * AccumulatorsPerThread::kW, 16>
    Accumulators;
 
  CUTLASS_DEVICE ThreadDiffSquaredAdd() {}
 
  CUTLASS_DEVICE void multiply_add(FragmentA const& a,
                                   FragmentB const& b,
                                   Accumulators const& c,
                                   Accumulators& d)
  {
    for (int j = 0; j < AccumulatorsPerThread::kH; ++j) {
      for (int i = 0; i < AccumulatorsPerThread::kW; ++i) {
        auto diff      = a[i] - b[j];
        const auto idx = j * AccumulatorsPerThread::kW + i;
        d[idx]         = diff * diff + c[idx];
      }
    }
  }
};
 
template <typename AccumulatorsPerThread_,
          typename ThreadsPerWarp_,
          typename ScalarA_,
          typename ScalarB_,
          typename ScalarC_>
struct ThreadL1NormAdd {
  typedef cutlass::Shape<1, 1, 1, 1> InstructionShape;
  typedef AccumulatorsPerThread_ AccumulatorsPerThread;
  typedef ThreadsPerWarp_ ThreadsPerWarp;
  typedef
    typename cutlass::ShapeMul<AccumulatorsPerThread, ThreadsPerWarp>::Shape AccumulatorsPerWarp;
  typedef ScalarA_ ScalarA;
  typedef cutlass::Fragment<ScalarA, AccumulatorsPerThread::kW> FragmentA;
  typedef ScalarB_ ScalarB;
  typedef cutlass::Fragment<ScalarB, AccumulatorsPerThread::kH> FragmentB;
  typedef ScalarC_ ScalarC;
  typedef cutlass::Fragment<ScalarC, AccumulatorsPerThread::kH * AccumulatorsPerThread::kW, 16>
    Accumulators;
 
  CUTLASS_DEVICE ThreadL1NormAdd() {}
 
  CUTLASS_DEVICE void multiply_add(FragmentA const& a,
                                   FragmentB const& b,
                                   Accumulators const& c,
                                   Accumulators& d)
  {
    for (int j = 0; j < AccumulatorsPerThread::kH; ++j) {
      for (int i = 0; i < AccumulatorsPerThread::kW; ++i) {
        auto diff      = a[i] < b[j] ? b[j] - a[i] : a[i] - b[j];
        const auto idx = j * AccumulatorsPerThread::kW + i;
        d[idx]         = diff + c[idx];
      }
    }
  }
};
 
};  // end namespace LinAlg
};  // end namespace MLCommon