| 73 | namespace x86 { |
| 74 | |
| 75 | void LRNImpl::exec( |
| 76 | _megdnn_tensor_in src, _megdnn_tensor_out dst, _megdnn_workspace workspace) { |
| 77 | check_exec(src.layout, dst.layout, workspace.size); |
| 78 | auto N = src.layout.shape[0], C = src.layout.shape[1], H = src.layout.shape[2], |
| 79 | W = src.layout.shape[3]; |
| 80 | |
| 81 | std::function<void( |
| 82 | const float*, float*, size_t, size_t, size_t, size_t, float, float, float)> |
| 83 | f = nullptr; |
| 84 | if (is_supported(SIMDType::FMA)) { |
| 85 | f = &lrn_single_instance<SIMDType::FMA>; |
| 86 | } else if (is_supported(SIMDType::AVX)) { |
| 87 | f = &lrn_single_instance<SIMDType::AVX>; |
| 88 | } else if (is_supported(SIMDType::SSE)) { |
| 89 | f = &lrn_single_instance<SIMDType::SSE>; |
| 90 | } else { |
| 91 | megdnn_throw("no fma/avx/sse detected"); |
| 92 | } |
| 93 | auto n = param().n; |
| 94 | auto k = param().k; |
| 95 | auto alpha = param().alpha; |
| 96 | auto beta = param().beta; |
| 97 | MEGDNN_DISPATCH_CPU_KERN_OPR(auto sptr = src.ptr<dt_float32>(); |
| 98 | auto dptr = dst.ptr<dt_float32>(); rep(i, N) { |
| 99 | f(sptr, dptr, C, H, W, n, k, alpha, beta); |
| 100 | sptr += C * H * W; |
| 101 | dptr += C * H * W; |
| 102 | }); |
| 103 | } |
| 104 | |
| 105 | } // namespace x86 |
| 106 | } // namespace megdnn |