mirror of
https://github.com/zebrajr/pytorch.git
synced 2025-12-07 12:21:27 +01:00
086dba3804
Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/35164 As title Test Plan: CI Reviewed By: jianyuh Differential Revision: D20581853 fbshipit-source-id: 393ddd9487cd965c465eaa49e1509863618a6048
34 lines
655 B
C++
34 lines
655 B
C++
#pragma once
|
|
#ifdef __AVX__
|
|
#include <immintrin.h>
|
|
#endif
|
|
#include "caffe2/utils/math.h"
|
|
|
|
namespace caffe2 {
|
|
namespace internal {
|
|
|
|
// Z=X*Y
|
|
template <typename XT, typename YT, typename ZT>
|
|
void dot(const int N, const XT* x, const YT* y, ZT* z, CPUContext* ctx) {
|
|
CAFFE_THROW("Unsupported, see specialized implementations");
|
|
}
|
|
|
|
template <>
|
|
void dot<float, float, float>(
|
|
const int N,
|
|
const float* x,
|
|
const float* y,
|
|
float* z,
|
|
CPUContext* ctx);
|
|
|
|
template <>
|
|
void dot<float, at::Half, float>(
|
|
const int N,
|
|
const float* x,
|
|
const at::Half* y,
|
|
float* z,
|
|
CPUContext* ctx);
|
|
|
|
} // namespace internal
|
|
} // namespace caffe2
|