mirror of
https://github.com/zebrajr/pytorch.git
synced 2025-12-07 12:21:27 +01:00
eef85f89b9
Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/62428 In this diff we add a broadcast fastpath for reduce utility functions. These functions are used by various elementwise ops, whose tests we update to exercise the new functionality. Test Plan: Added test cases to elementwise ops (which will exercise the new reducer functionality) that will be run by CI. It's worth noting there's still no code (outside of the new test cases) that takes the new code paths added -- the user must explicitly request `allow_broadcast_fastpath=True`, and nothing outside of the added tests currently does so. Differential Revision: D29938264 fbshipit-source-id: 5d5542bd93afb85fd9f7a4073f766adc07eb3b65
27 lines
555 B
C++
27 lines
555 B
C++
#ifndef CAFFE2_UTILS_MATH_BROADCAST_H_
|
|
#define CAFFE2_UTILS_MATH_BROADCAST_H_
|
|
|
|
#include "caffe2/core/common.h"
|
|
#include "caffe2/core/types.h"
|
|
|
|
namespace caffe2 {
|
|
namespace math {
|
|
|
|
bool can_use_broadcast_fastpath(int ndim, const int* dims);
|
|
|
|
template <typename T, class Context, StorageOrder kOrder>
|
|
TORCH_API void AffineChannel(
|
|
const int N,
|
|
const int C,
|
|
const int HxW,
|
|
const T* X,
|
|
const T* scale,
|
|
const T* bias,
|
|
T* Y,
|
|
Context* context);
|
|
|
|
} // namespace math
|
|
} // namespace caffe2
|
|
|
|
#endif // CAFFE2_UTILS_MATH_BROADCAST_H_
|