diff options
Diffstat (limited to 'arm_compute/runtime/NEON/functions/assembly/NEGEMMInterleavedWrapper.h')
-rw-r--r-- | arm_compute/runtime/NEON/functions/assembly/NEGEMMInterleavedWrapper.h | 48 |
1 files changed, 48 insertions, 0 deletions
diff --git a/arm_compute/runtime/NEON/functions/assembly/NEGEMMInterleavedWrapper.h b/arm_compute/runtime/NEON/functions/assembly/NEGEMMInterleavedWrapper.h index 4c5a8d4e59..26236ffb35 100644 --- a/arm_compute/runtime/NEON/functions/assembly/NEGEMMInterleavedWrapper.h +++ b/arm_compute/runtime/NEON/functions/assembly/NEGEMMInterleavedWrapper.h @@ -44,6 +44,53 @@ class MatrixMultiplyWorkload; class NEGEMMInterleavedTransformAWrapper; class NEGEMMInterleavedMatrixMultiplyWrapper; +/** Buffer manager used when reshaping B on the fly + * + * The typical workflow is: + * - lock_to_reshape_if_needed() + * - If the previous lock was successful: mark_as_reshaped() + * - wait_for_reshaping() wait for the reshaping to be complete + * - mark_as_unused() once the thread is done using this given buffer. + * + * Calls for different indices might be interleaved, however the calls for a given index must always be in that order. + */ +class IBufferManager +{ +public: + /** Lock a buffer for the given index if it's available else return + * + * @param[in] index Index of the buffer to lock + * + * @return True if the buffer has been successfully locked, false if it's already reshaped / being reshaped. + */ + virtual bool lock_to_reshape_if_needed(unsigned int index) = 0; + /** Mark a buffer previously locked as reshaped + * + * @pre The thread calling this function must have locked the given buffer through lock_to_reshape_if_needed() + * + * @param[in] index Index of the buffer to mark as reshaped + */ + virtual void mark_as_reshaped(unsigned int index) = 0; + /** Block until the given buffer is marked as reshaped + * + * @param[in] index Index of the buffer + */ + virtual void wait_for_reshaping(unsigned int index) = 0; + /** Mark a reshaped buffer as unused + * + * Once all the users have marked a buffer as unused then it goes back to being free + */ + virtual void mark_as_unused(unsigned int index) = 0; + + /** Number of buffers used internally + * + * @return The number of buffers used by the manager. + */ + virtual unsigned int num_buffers() const = 0; + /** Default destructor */ + virtual ~IBufferManager() = default; +}; + /** Equivalent to arm_gemm::GemmInterleaved but using Compute Library types. */ class NEGEMMInterleavedWrapper : public IFunction @@ -89,6 +136,7 @@ private: std::unique_ptr<NEGEMMInterleavedPrepareBWrapperKernel> _prepare_b{ nullptr }; std::unique_ptr<NEGEMMInterleavedTransformAWrapper> _transform_a{ nullptr }; std::unique_ptr<NEGEMMInterleavedMatrixMultiplyWrapper> _matrix_multiply{ nullptr }; + std::unique_ptr<IBufferManager> _buffer_manager{ nullptr }; std::vector<TransformAWorkload> _a_workloads{}; std::vector<PrepareBWorkload> _b_workloads{}; std::vector<MatrixMultiplyWorkload> _mm_workloads{}; |