diff --git a/llvm/include/llvm/Frontend/OpenMP/OMPKinds.def b/llvm/include/llvm/Frontend/OpenMP/OMPKinds.def
--- a/llvm/include/llvm/Frontend/OpenMP/OMPKinds.def
+++ b/llvm/include/llvm/Frontend/OpenMP/OMPKinds.def
@@ -207,6 +207,7 @@
 
 __OMP_RTL(__kmpc_get_hardware_num_blocks, false, Int32, )
 __OMP_RTL(__kmpc_get_hardware_num_threads_in_block, false, Int32, )
+__OMP_RTL(__kmpc_get_warp_size, false, Int32, )
 
 __OMP_RTL(omp_get_thread_num, false, Int32, )
 __OMP_RTL(omp_get_num_threads, false, Int32, )
@@ -455,8 +456,6 @@
 __OMP_RTL(__kmpc_warp_active_thread_mask, false, Int64,)
 __OMP_RTL(__kmpc_syncwarp, false, Void, Int64)
 
-__OMP_RTL(__kmpc_get_warp_size, false, Int32, )
-
 __OMP_RTL(__kmpc_is_generic_main_thread_id, false, Int8, Int32)
 
 __OMP_RTL(__last, false, Void, )
@@ -629,6 +628,7 @@
 
 __OMP_RTL_ATTRS(__kmpc_get_hardware_num_blocks, GetterAttrs, AttributeSet(), ParamAttrs())
 __OMP_RTL_ATTRS(__kmpc_get_hardware_num_threads_in_block, GetterAttrs, AttributeSet(), ParamAttrs())
+__OMP_RTL_ATTRS(__kmpc_get_warp_size, GetterAttrs, AttributeSet(), ParamAttrs())
 
 __OMP_RTL_ATTRS(omp_get_thread_num, GetterAttrs, AttributeSet(), ParamAttrs())
 __OMP_RTL_ATTRS(omp_get_num_threads, GetterAttrs, AttributeSet(), ParamAttrs())
diff --git a/openmp/libomptarget/test/offloading/bug51781.c b/openmp/libomptarget/test/offloading/bug51781.c
--- a/openmp/libomptarget/test/offloading/bug51781.c
+++ b/openmp/libomptarget/test/offloading/bug51781.c
@@ -21,14 +21,29 @@
 // RUN: %fcheck-amdgcn-amd-amdhsa -check-prefix=CUSTOM -input-file=%t.custom
 // RUN: %libomptarget-run-generic 2>&1 | %fcheck-generic
 //
+// Repeat with reduction clause, which has managed to break the custom state
+// machine in the past.
+//
+// RUN: %libomptarget-compile-generic -O1 -Rpass=openmp-opt -DADD_REDUCTION \
+// RUN:   -mllvm -openmp-opt-disable-spmdization > %t.custom 2>&1
+// RUN: %fcheck-nvptx64-nvidia-cuda -check-prefix=CUSTOM -input-file=%t.custom
+// RUN: %fcheck-amdgcn-amd-amdhsa -check-prefix=CUSTOM -input-file=%t.custom
+// RUN: %libomptarget-run-generic 2>&1 | %fcheck-generic
+//
 // CUSTOM: Rewriting generic-mode kernel with a customized state machine.
 
+#if ADD_REDUCTION
+# define REDUCTION(...) reduction(__VA_ARGS__)
+#else
+# define REDUCTION(...)
+#endif
+
 #include <stdio.h>
 int main() {
   int x = 0, y = 1;
-  #pragma omp target teams num_teams(1) map(tofrom:x, y)
+  #pragma omp target teams num_teams(1) map(tofrom:x, y) REDUCTION(+:x)
   {
-    x = 5;
+    x += 5;
     #pragma omp parallel
     y = 6;
   }