arm64: assembler: add cond_yield macro

Add a macro cond_yield that branches to a specified label when called if
the TIF_NEED_RESCHED flag is set and decreasing the preempt count would
make the task preemptible again, resulting in a schedule to occur. This
can be used by kernel mode SIMD code that keeps a lot of state in SIMD
registers, which would make chunking the input in order to perform the
cond_resched() check from C code disproportionately costly.

Signed-off-by: Ard Biesheuvel <ardb@kernel.org>
Link: https://lore.kernel.org/r/20210203113626.220151-2-ardb@kernel.org
Signed-off-by: Will Deacon <will@kernel.org>

diff --git a/arch/arm64/include/asm/assembler.h b/arch/arm64/include/asm/assembler.h
index bf125c59111688840bf3f3e84fca12b74fb1098b..27b1ea721c2df86a60088ec40b97a99f42d26bce 100644 (file)
--- a/arch/arm64/include/asm/assembler.h
+++ b/arch/arm64/include/asm/assembler.h
@@ -745,6 +745,22 @@ USER(\label, ic    ivau, \tmp2)                    // invalidate I line PoU
 .Lyield_out_\@ :
        .endm
 
+       /*
+        * Check whether preempt-disabled code should yield as soon as it
+        * is able. This is the case if re-enabling preemption a single
+        * time results in a preempt count of zero, and the TIF_NEED_RESCHED
+        * flag is set. (Note that the latter is stored negated in the
+        * top word of the thread_info::preempt_count field)
+        */
+       .macro          cond_yield, lbl:req, tmp:req
+#ifdef CONFIG_PREEMPTION
+       get_current_task \tmp
+       ldr             \tmp, [\tmp, #TSK_TI_PREEMPT]
+       sub             \tmp, \tmp, #PREEMPT_DISABLE_OFFSET
+       cbz             \tmp, \lbl
+#endif
+       .endm
+
 /*
  * This macro emits a program property note section identifying
  * architecture features which require special handling, mainly for