pytorch/caffe2/operators/minmax_ops.cc
Orion Reblitz-Richardson 7f33ec55b2 Fix Eigen issue on OS X with CUDA and nvcc compile (#9350)
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/9350

Re-apply #9270

Breaking this out of #8338

This takes care of the Eigen failure we saw on Mac CUDA builds when BUILD_CAFFE2 and BUILD_ATEN were removed. Fix is to isolate Eigen from headers included by cu files and processed by nvcc. This was worked on with smessmer.

Reviewed By: mingzhe09088

Differential Revision: D8794431

fbshipit-source-id: de656334af46c697802073f8e8d9a6aeb9ca65a7
2018-07-11 14:00:05 -07:00

176 lines
4.4 KiB
C++

#include "caffe2/operators/minmax_ops.h"
#include "caffe2/utils/eigen_utils.h"
namespace caffe2 {
REGISTER_CPU_OPERATOR(Max, MaxOp<float, CPUContext>);
REGISTER_CPU_OPERATOR(Min, MinOp<float, CPUContext>);
OPERATOR_SCHEMA(Max)
.NumInputs(1, INT_MAX)
.NumOutputs(1)
.IdenticalTypeAndShapeOfInput(0)
.AllowInplace({{0, 0}})
.SetDoc(R"DOC(
Element-wise max of an arbitrary number of input tensors. This operation can be
performed in-place, by using the first input blob as the output blob. All inputs
must have the same shape and data type, and the output will have the same shape
as the inputs.
Github Link:
- https://github.com/pytorch/pytorch/blob/master/caffe2/operators/minmax_ops.cc
<details>
<summary> <b>Example</b> </summary>
**Code**
```
workspace.ResetWorkspace()
op = core.CreateOperator(
"Max",
["X", "Y", "Z"],
["X"],
)
workspace.FeedBlob("X", (np.random.rand(3,3)).astype(np.float32))
workspace.FeedBlob("Y", (np.random.rand(3,3)).astype(np.float32))
workspace.FeedBlob("Z", (np.random.rand(3,3)).astype(np.float32))
print("X:", workspace.FetchBlob("X"))
print("Y:", workspace.FetchBlob("Y"))
print("Z:", workspace.FetchBlob("Z"))
workspace.RunOperatorOnce(op)
print("Max:", workspace.FetchBlob("X"))
```
**Result**
```
X:
[[0.4496477 0.07061381 0.7139333 ]
[0.83203 0.05970785 0.72786295]
[0.75988126 0.04601283 0.32820013]]
Y:
[[0.05683139 0.16872478 0.671098 ]
[0.70739156 0.09878621 0.03416285]
[0.34087983 0.94986707 0.67263436]]
Z:
[[0.48051122 0.07141234 0.85264146]
[0.77086854 0.22082241 0.13154659]
[0.42401117 0.995431 0.4263775 ]]
Max:
[[0.48051122 0.16872478 0.85264146]
[0.83203 0.22082241 0.72786295]
[0.75988126 0.995431 0.67263436]]
```
</details>
)DOC")
.Input(0, "X, Y, ...", "*(type: Tensor`<Ord>`)* List of input tensors with the same shape.")
.Output(0, "M", "*(type: Tensor`<Ord>`)* Output tensor with same dimensions as input(s)."
"Contains the maximum valued element at each location.")
.InheritOnnxSchema("Max");
OPERATOR_SCHEMA(Min)
.NumInputs(1, INT_MAX)
.NumOutputs(1)
.IdenticalTypeAndShapeOfInput(0)
.AllowInplace({{0, 0}})
.SetDoc(R"DOC(
Element-wise min of an arbitrary number of input tensors. This operation can be performed in-place, by using the first input blob as the output blob. All inputs must have the same shape and data type, and the output will have the same shape as the inputs.
Github Link:
- https://github.com/pytorch/pytorch/blob/master/caffe2/operators/minmax_ops.cc
<details>
<summary> <b>Example</b> </summary>
**Code**
```
workspace.ResetWorkspace()
op = core.CreateOperator(
"Min",
["X", "Y", "Z"],
["X"],
)
workspace.FeedBlob("X", (np.random.rand(2,2)).astype(np.float32))
workspace.FeedBlob("Y", (np.random.rand(2,2)).astype(np.float32))
workspace.FeedBlob("Z", (np.random.rand(2,2)).astype(np.float32))
print("X:", workspace.FetchBlob("X"))
print("Y:", workspace.FetchBlob("Y"))
print("Z:", workspace.FetchBlob("Z"))
workspace.RunOperatorOnce(op)
print("Min:", workspace.FetchBlob("X"))
```
**Result**
```
X:
[[0.32731926 0.4939747 ]
[0.29242373 0.43460014]]
Y:
[[0.40928316 0.916115 ]
[0.77526504 0.29339448]]
Z:
[[0.7899794 0.90335774]
[0.82599413 0.2843068 ]]
Min:
[[0.32731926 0.4939747 ]
[0.29242373 0.2843068 ]]
```
</details>
)DOC")
.Input(0, "X, Y, ...", "*(type: Tensor`<Ord>`)* List of input tensors with the same shape.")
.Output(0, "M", "*(type: Tensor`<Ord>`)* Output tensor with same dimensions as input(s)."
"Contains the minimum valued element at each location.")
.InheritOnnxSchema("Min");
template <typename T, class Context>
bool MaxOp<T, Context>::Compute() {
auto& input0 = Input(0);
const int N = input0.size();
T* output_data = Output(0)->template mutable_data<T>();
for (int i = 1; i < InputSize(); i++) {
auto input_data = Input(i).template data<T>();
EigenVectorMap<T> output_vec(output_data, N);
output_vec = output_vec.cwiseMax(ConstEigenVectorMap<T>(input_data, N));
}
return true;
}
template <typename T, class Context>
bool MinOp<T, Context>::Compute() {
auto& input0 = Input(0);
const int N = input0.size();
T* output_data = Output(0)->template mutable_data<T>();
for (int i = 1; i < InputSize(); i++) {
auto input_data = Input(i).template data<T>();
EigenVectorMap<T> output_vec(output_data, N);
output_vec = output_vec.cwiseMin(ConstEigenVectorMap<T>(input_data, N));
}
return true;
}
} // namespace caffe2