include/ck_tile/host/reference/reference_batched_gemm.hpp Source File

• include
• ck_tile
• host
• reference

reference_batched_gemm.hpp

Go to the documentation of this file.

1 // SPDX-License-Identifier: MIT

2 // Copyright (c) 2018-2023, Advanced Micro Devices, Inc. All rights reserved.

3 

4 #pragma once

5 

6 #include "ck_tile/core.hpp"

7 #include "ck_tile/host/host_tensor.hpp"

8 #include <thread>

9 

10 namespace ck_tile {

11 

12 template <typename ADataType,

13  typename BDataType,

14  typename AccDataType,

15  typename CDataType,

16  typename AElementOp = ck_tile::identity,

17  typename BElementOp = ck_tile::identity,

18  typename ACCElementOp = ck_tile::identity>

19 CK_TILE_HOST void reference_batched_gemm(const HostTensor<ADataType>& a_b_m_k,

20  const HostTensor<BDataType>& b_b_n_k,

21  HostTensor<CDataType>& c_b_m_n,

22  const AElementOp& a_element_op = {},

23  const BElementOp& b_element_op = {},

24  const ACCElementOp& acc_element_op = {})

25 {

26  const int N = b_b_n_k.mDesc.get_lengths()[1];

27  const int K = b_b_n_k.mDesc.get_lengths()[2];

28 

29  auto f = [&](auto batch, auto m) {

30  for(int n = 0; n < N; ++n)

31  {

32  AccDataType v_acc = 0;

33 

34  for(int k = 0; k < K; ++k)

35  {

36  ADataType v_a = a_element_op(a_b_m_k(batch, m, k));

37  BDataType v_b = b_element_op(b_b_n_k(batch, n, k));

38 

39  v_acc += ck_tile::type_convert<AccDataType>(v_a) *

40  ck_tile::type_convert<AccDataType>(v_b);

41  }

42 

43  c_b_m_n(batch, m, n) = ck_tile::type_convert<CDataType>(acc_element_op(v_acc));

44  }

45  };

46 

47  make_ParallelTensorFunctor(f, c_b_m_n.mDesc.get_lengths()[0], c_b_m_n.mDesc.get_lengths()[1])(

48  std::thread::hardware_concurrency());

49 }

50 } // namespace ck_tile

CK_TILE_HOST

#define CK_TILE_HOST

Definition: config.hpp:39

core.hpp

ck_tile

Definition: cluster_descriptor.hpp:13

ck_tile::make_ParallelTensorFunctor

CK_TILE_HOST auto make_ParallelTensorFunctor(F f, Xs... xs)

Definition: host_tensor.hpp:272

ck_tile::reference_batched_gemm

CK_TILE_HOST void reference_batched_gemm(const HostTensor< ADataType > &a_b_m_k, const HostTensor< BDataType > &b_b_n_k, HostTensor< CDataType > &c_b_m_n, const AElementOp &a_element_op={}, const BElementOp &b_element_op={}, const ACCElementOp &acc_element_op={})

Definition: reference_batched_gemm.hpp:19

ck_tile::HostTensorDescriptor::get_lengths

const std::vector< std::size_t > & get_lengths() const

Definition: host_tensor.hpp:162

ck_tile::HostTensor

Definition: host_tensor.hpp:279

ck_tile::HostTensor::mDesc

Descriptor mDesc

Definition: host_tensor.hpp:678

ck_tile::identity

Definition: functional.hpp:62

host_tensor.hpp