transpose-inl.h (7428B)
1 // Copyright (c) the JPEG XL Project Authors. All rights reserved. 2 // 3 // Use of this source code is governed by a BSD-style 4 // license that can be found in the LICENSE file. 5 6 // Block transpose for DCT/IDCT 7 8 #if defined(LIB_JXL_TRANSPOSE_INL_H_) == defined(HWY_TARGET_TOGGLE) 9 #ifdef LIB_JXL_TRANSPOSE_INL_H_ 10 #undef LIB_JXL_TRANSPOSE_INL_H_ 11 #else 12 #define LIB_JXL_TRANSPOSE_INL_H_ 13 #endif 14 15 #include <stddef.h> 16 17 #include <hwy/highway.h> 18 #include <type_traits> 19 20 #include "lib/jxl/base/status.h" 21 #include "lib/jxl/dct_block-inl.h" 22 23 HWY_BEFORE_NAMESPACE(); 24 namespace jxl { 25 namespace HWY_NAMESPACE { 26 namespace { 27 28 #ifndef JXL_INLINE_TRANSPOSE 29 // Workaround for issue #42 - (excessive?) inlining causes invalid codegen. 30 #if defined(__arm__) 31 #define JXL_INLINE_TRANSPOSE HWY_NOINLINE 32 #else 33 #define JXL_INLINE_TRANSPOSE HWY_INLINE 34 #endif 35 #endif // JXL_INLINE_TRANSPOSE 36 37 // Simple wrapper that ensures that a function will not be inlined. 38 template <typename T, typename... Args> 39 JXL_NOINLINE void NoInlineWrapper(const T& f, const Args&... args) { 40 return f(args...); 41 } 42 43 template <bool enabled> 44 struct TransposeSimdTag {}; 45 46 // TODO(veluca): it's not super useful to have this in the SIMD namespace. 47 template <size_t ROWS_or_0, size_t COLS_or_0, class From, class To> 48 JXL_INLINE_TRANSPOSE void GenericTransposeBlock( 49 TransposeSimdTag<false> /* tag */, const From& from, const To& to, 50 size_t ROWSp, size_t COLSp) { 51 size_t ROWS = ROWS_or_0 == 0 ? ROWSp : ROWS_or_0; 52 size_t COLS = COLS_or_0 == 0 ? COLSp : COLS_or_0; 53 for (size_t n = 0; n < ROWS; ++n) { 54 for (size_t m = 0; m < COLS; ++m) { 55 to.Write(from.Read(n, m), m, n); 56 } 57 } 58 } 59 60 // TODO(veluca): AVX3? 61 #if HWY_CAP_GE256 62 constexpr bool TransposeUseSimd(size_t ROWS, size_t COLS) { 63 return ROWS % 8 == 0 && COLS % 8 == 0; 64 } 65 66 template <size_t ROWS_or_0, size_t COLS_or_0, class From, class To> 67 JXL_INLINE_TRANSPOSE void GenericTransposeBlock( 68 TransposeSimdTag<true> /* tag */, const From& from, const To& to, 69 size_t ROWSp, size_t COLSp) { 70 size_t ROWS = ROWS_or_0 == 0 ? ROWSp : ROWS_or_0; 71 size_t COLS = COLS_or_0 == 0 ? COLSp : COLS_or_0; 72 static_assert(MaxLanes(BlockDesc<8>()) == 8, "Invalid descriptor size"); 73 static_assert(ROWS_or_0 % 8 == 0, "Invalid number of rows"); 74 static_assert(COLS_or_0 % 8 == 0, "Invalid number of columns"); 75 for (size_t n = 0; n < ROWS; n += 8) { 76 for (size_t m = 0; m < COLS; m += 8) { 77 const BlockDesc<8> d; 78 auto i0 = from.LoadPart(d, n + 0, m + 0); 79 auto i1 = from.LoadPart(d, n + 1, m + 0); 80 auto i2 = from.LoadPart(d, n + 2, m + 0); 81 auto i3 = from.LoadPart(d, n + 3, m + 0); 82 auto i4 = from.LoadPart(d, n + 4, m + 0); 83 auto i5 = from.LoadPart(d, n + 5, m + 0); 84 auto i6 = from.LoadPart(d, n + 6, m + 0); 85 auto i7 = from.LoadPart(d, n + 7, m + 0); 86 // Surprisingly, this straightforward implementation (24 cycles on port5) 87 // is faster than load128+insert and LoadDup128+ConcatUpperLower+blend. 88 const auto q0 = InterleaveLower(d, i0, i2); 89 const auto q1 = InterleaveLower(d, i1, i3); 90 const auto q2 = InterleaveUpper(d, i0, i2); 91 const auto q3 = InterleaveUpper(d, i1, i3); 92 const auto q4 = InterleaveLower(d, i4, i6); 93 const auto q5 = InterleaveLower(d, i5, i7); 94 const auto q6 = InterleaveUpper(d, i4, i6); 95 const auto q7 = InterleaveUpper(d, i5, i7); 96 97 const auto r0 = InterleaveLower(d, q0, q1); 98 const auto r1 = InterleaveUpper(d, q0, q1); 99 const auto r2 = InterleaveLower(d, q2, q3); 100 const auto r3 = InterleaveUpper(d, q2, q3); 101 const auto r4 = InterleaveLower(d, q4, q5); 102 const auto r5 = InterleaveUpper(d, q4, q5); 103 const auto r6 = InterleaveLower(d, q6, q7); 104 const auto r7 = InterleaveUpper(d, q6, q7); 105 106 i0 = ConcatLowerLower(d, r4, r0); 107 i1 = ConcatLowerLower(d, r5, r1); 108 i2 = ConcatLowerLower(d, r6, r2); 109 i3 = ConcatLowerLower(d, r7, r3); 110 i4 = ConcatUpperUpper(d, r4, r0); 111 i5 = ConcatUpperUpper(d, r5, r1); 112 i6 = ConcatUpperUpper(d, r6, r2); 113 i7 = ConcatUpperUpper(d, r7, r3); 114 to.StorePart(d, i0, m + 0, n + 0); 115 to.StorePart(d, i1, m + 1, n + 0); 116 to.StorePart(d, i2, m + 2, n + 0); 117 to.StorePart(d, i3, m + 3, n + 0); 118 to.StorePart(d, i4, m + 4, n + 0); 119 to.StorePart(d, i5, m + 5, n + 0); 120 to.StorePart(d, i6, m + 6, n + 0); 121 to.StorePart(d, i7, m + 7, n + 0); 122 } 123 } 124 } 125 #elif HWY_TARGET != HWY_SCALAR 126 constexpr bool TransposeUseSimd(size_t ROWS, size_t COLS) { 127 return ROWS % 4 == 0 && COLS % 4 == 0; 128 } 129 130 template <size_t ROWS_or_0, size_t COLS_or_0, class From, class To> 131 JXL_INLINE_TRANSPOSE void GenericTransposeBlock( 132 TransposeSimdTag<true> /* tag */, const From& from, const To& to, 133 size_t ROWSp, size_t COLSp) { 134 size_t ROWS = ROWS_or_0 == 0 ? ROWSp : ROWS_or_0; 135 size_t COLS = COLS_or_0 == 0 ? COLSp : COLS_or_0; 136 static_assert(MaxLanes(BlockDesc<4>()) == 4, "Invalid descriptor size"); 137 static_assert(ROWS_or_0 % 4 == 0, "Invalid number of rows"); 138 static_assert(COLS_or_0 % 4 == 0, "Invalid number of columns"); 139 for (size_t n = 0; n < ROWS; n += 4) { 140 for (size_t m = 0; m < COLS; m += 4) { 141 const BlockDesc<4> d; 142 const auto p0 = from.LoadPart(d, n + 0, m + 0); 143 const auto p1 = from.LoadPart(d, n + 1, m + 0); 144 const auto p2 = from.LoadPart(d, n + 2, m + 0); 145 const auto p3 = from.LoadPart(d, n + 3, m + 0); 146 147 const auto q0 = InterleaveLower(d, p0, p2); 148 const auto q1 = InterleaveLower(d, p1, p3); 149 const auto q2 = InterleaveUpper(d, p0, p2); 150 const auto q3 = InterleaveUpper(d, p1, p3); 151 152 const auto r0 = InterleaveLower(d, q0, q1); 153 const auto r1 = InterleaveUpper(d, q0, q1); 154 const auto r2 = InterleaveLower(d, q2, q3); 155 const auto r3 = InterleaveUpper(d, q2, q3); 156 157 to.StorePart(d, r0, m + 0, n + 0); 158 to.StorePart(d, r1, m + 1, n + 0); 159 to.StorePart(d, r2, m + 2, n + 0); 160 to.StorePart(d, r3, m + 3, n + 0); 161 } 162 } 163 } 164 #else 165 constexpr bool TransposeUseSimd(size_t ROWS, size_t COLS) { return false; } 166 #endif 167 168 template <size_t N, size_t M, typename = void> 169 struct Transpose { 170 template <typename From, typename To> 171 static void Run(const From& from, const To& to) { 172 // This does not guarantee anything, just saves from the most stupid 173 // mistakes. 174 JXL_DASSERT(from.Address(0, 0) != to.Address(0, 0)); 175 TransposeSimdTag<TransposeUseSimd(N, M)> tag; 176 GenericTransposeBlock<N, M>(tag, from, to, N, M); 177 } 178 }; 179 180 // Avoid inlining and unrolling transposes for large blocks. 181 template <size_t N, size_t M> 182 struct Transpose< 183 N, M, typename std::enable_if<(N >= 8 && M >= 8 && N * M >= 512)>::type> { 184 template <typename From, typename To> 185 static void Run(const From& from, const To& to) { 186 // This does not guarantee anything, just saves from the most stupid 187 // mistakes. 188 JXL_DASSERT(from.Address(0, 0) != to.Address(0, 0)); 189 TransposeSimdTag<TransposeUseSimd(N, M)> tag; 190 constexpr void (*transpose)(TransposeSimdTag<TransposeUseSimd(N, M)>, 191 const From&, const To&, size_t, size_t) = 192 GenericTransposeBlock<0, 0, From, To>; 193 NoInlineWrapper(transpose, tag, from, to, N, M); 194 } 195 }; 196 197 } // namespace 198 // NOLINTNEXTLINE(google-readability-namespace-comments) 199 } // namespace HWY_NAMESPACE 200 } // namespace jxl 201 HWY_AFTER_NAMESPACE(); 202 203 #endif // LIB_JXL_TRANSPOSE_INL_H_