modify barrier for x86

ztxz16 · May 8, 2024 · 7812651 · 7812651
1 parent cf1f6b1
commit 7812651
Showing 1 changed file with 11 additions and 4 deletions.
diff --git a/include/devices/cpu/alivethreadpool.h b/include/devices/cpu/alivethreadpool.h
@@ -11,6 +11,13 @@
 #include <cstring>
 
 namespace fastllm {
+    static void barrier() {
+#ifdef __aarch64__
+        asm volatile("dmb ish");
+#else
+        __asm__ __volatile__("": : :"memory");
+#endif
+    }
     struct MultiThreadBaseOp {
         virtual void Run() = 0;
     };
@@ -38,11 +45,11 @@ namespace fastllm {
         void operator()() {
             auto lastRunTime = std::chrono::system_clock::now();
             while (true) {
-                asm volatile("dmb ish");
+                barrier();
                 if (task->signal == 1) {
                     task->op->Run();
                     task->signal = 0;
-                    asm volatile("dmb ish");
+                    barrier();
                     lastRunTime = std::chrono::system_clock::now();
                 }
 
@@ -56,9 +63,9 @@ namespace fastllm {
 
         void PushOp(MultiThreadBaseOp *op) {
             this->task->op = op;
-            asm volatile("dmb ish");
+            barrier();
             this->task->signal = 1;
-            asm volatile("dmb ish");
+            barrier();
         }
 
         void Wait() {