1 #pragma once 2 3 #ifdef USE_XNNPACK 4 5 #include <xnnpack.h> 6 #include <caffe2/utils/threadpool/pthreadpool-cpp.h> 7 #include <c10/util/ArrayRef.h> 8 #include <limits> 9 #include <memory> 10 11 namespace at::native::xnnpack { 12 13 struct Deleter final { operatorfinal14 void operator()(const xnn_operator_t op) const { 15 xnn_delete_operator(op); 16 } 17 }; 18 19 using Operator = std::unique_ptr<xnn_operator, Deleter>; 20 21 struct ContextLinear final { 22 Operator op; 23 int64_t output_channels; 24 25 ContextLinear() = delete; 26 ContextLinearfinal27 ContextLinear(Operator&& o, int64_t o_channels) : op(std::move(o)), output_channels(o_channels) {} 28 static constexpr float kMin = -std::numeric_limits<float>::infinity(); 29 static constexpr float kMax = std::numeric_limits<float>::infinity(); 30 }; 31 32 // This contains information for both the transpose and non-transpose cases. 33 struct ContextConv2D final { 34 Operator op; 35 std::array<int64_t, 4> weight_size_; 36 std::array<int64_t, 2> padding_; 37 std::array<int64_t, 2> output_padding_; 38 std::array<int64_t, 2> stride_; 39 std::array<int64_t, 2> dilation_; 40 bool transposed_; 41 int64_t groups_; 42 43 ContextConv2D() = delete; 44 ContextConv2Dfinal45 ContextConv2D( 46 Operator&& o, 47 std::array<int64_t, 4> weight_size, 48 std::array<int64_t, 2> padding, 49 std::array<int64_t, 2> output_padding, 50 std::array<int64_t, 2> stride, 51 std::array<int64_t, 2> dilation, 52 bool transposed, 53 int64_t groups) 54 : op(std::move(o)), 55 weight_size_(weight_size), 56 padding_(padding), 57 output_padding_(output_padding), 58 stride_(stride), 59 dilation_(dilation), 60 transposed_(transposed), 61 groups_(groups) {} 62 static constexpr float kMin = -std::numeric_limits<float>::infinity(); 63 static constexpr float kMax = std::numeric_limits<float>::infinity(); 64 }; 65 66 67 namespace internal { 68 69 struct Layout final { 70 // 4D Activation Maps 71 struct Activation4D final { 72 static constexpr size_t batch = 0u; 73 static constexpr size_t channels = 1u; 74 static constexpr size_t height = 2u; 75 static constexpr size_t width = 3u; 76 }; 77 78 // ND Activation Maps 79 struct ActivationND final { 80 // Some operators may not be limited to 4 dimensional tensors. In that scenario, 81 // XNNPACK denotes that operator with an _nc suffix and expects all dimensions, 82 // except channels, to be flattened into one argument: batch_size. batchfinal::final83 static int64_t batch(const IntArrayRef tensor) { 84 if (C10_UNLIKELY(tensor.empty())) { 85 return -1; 86 } 87 88 // Handle the case where batch size is zero. 89 int64_t batch = tensor[0]; 90 91 for (size_t index = 1u; index < (tensor.size() - 1u); ++index) { 92 batch *= tensor[index]; 93 } 94 95 return batch; 96 }; 97 channelfinal::final98 static int64_t channel(const IntArrayRef tensor) { 99 if (C10_UNLIKELY(tensor.empty())) { 100 return -1; 101 } 102 103 return tensor.back(); 104 }; 105 }; 106 107 // Convolution Filters 108 struct Filter final { 109 static constexpr size_t output = 0u; 110 static constexpr size_t input = 1u; 111 static constexpr size_t height = 2u; 112 static constexpr size_t width = 3u; 113 }; 114 115 // Parameters (Pooling Kernels, Dilation, Padding, Stride, etc.) 116 struct Parameter final { 117 static constexpr size_t height = 0u; 118 static constexpr size_t width = 1u; 119 }; 120 }; 121 } // namespace internal 122 } // namespace at::native::xnnpack 123 124 #endif /* USE_XNNPACK */ 125 126 namespace at::native::xnnpack { 127 bool available(); 128 } // namespace at::native::xnnpack 129