format barrier.py with ruff

Author: reubenconducts

flash_attn/cute/barrier.py

@@ -4,8 +4,9 @@
 from cutlass.cutlass_dsl import T, dsl_user_op
 from cutlass._mlir.dialects import llvm
 
+
 @dsl_user_op
-def ld_acquire(lock_ptr : cute.Pointer, *, loc=None, ip=None) -> cutlass.Int32:
+def ld_acquire(lock_ptr: cute.Pointer, *, loc=None, ip=None) -> cutlass.Int32:
     lock_ptr_i64 = lock_ptr.toint(loc=loc, ip=ip).ir_value()
     state = llvm.inline_asm(
         T.i32(),
@@ -18,8 +19,11 @@ def ld_acquire(lock_ptr : cute.Pointer, *, loc=None, ip=None) -> cutlass.Int32:
     )
     return cutlass.Int32(state)
 
+
 @dsl_user_op
-def red_relaxed(lock_ptr : cute.Pointer, val: cutlass.Constexpr[Int32], *, loc=None, ip=None) -> None:
+def red_relaxed(
+    lock_ptr: cute.Pointer, val: cutlass.Constexpr[Int32], *, loc=None, ip=None
+) -> None:
     lock_ptr_i64 = lock_ptr.toint(loc=loc, ip=ip).ir_value()
     llvm.inline_asm(
         None,
@@ -31,8 +35,11 @@ def red_relaxed(lock_ptr : cute.Pointer, val: cutlass.Constexpr[Int32], *, loc=N
         asm_dialect=llvm.AsmDialect.AD_ATT,
     )
 
+
 @dsl_user_op
-def red_release(lock_ptr : cute.Pointer, val: cutlass.Constexpr[Int32], *, loc=None, ip=None) -> None:
+def red_release(
+    lock_ptr: cute.Pointer, val: cutlass.Constexpr[Int32], *, loc=None, ip=None
+) -> None:
     lock_ptr_i64 = lock_ptr.toint(loc=loc, ip=ip).ir_value()
     llvm.inline_asm(
         None,
@@ -43,28 +50,22 @@ def red_release(lock_ptr : cute.Pointer, val: cutlass.Constexpr[Int32], *, loc=N
         is_align_stack=False,
         asm_dialect=llvm.AsmDialect.AD_ATT,
     )
-    
+
+
 @cute.jit
-def wait_eq(
-    lock_ptr : cute.Pointer,
-    thread_idx : int | Int32,
-    flag_offset : int,
-    val : Int32
-) -> None:
+def wait_eq(lock_ptr: cute.Pointer, thread_idx: int | Int32, flag_offset: int, val: Int32) -> None:
     flag_ptr = lock_ptr + flag_offset
     if thread_idx == 0:
         read_val = Int32(0)
         while read_val != val:
             read_val = ld_acquire(flag_ptr)
 
+
 @cute.jit
 def arrive_inc(
-    lock_ptr : cute.Pointer,
-    thread_idx : int | Int32,
-    flag_offset : int,
-    val : cutlass.Constexpr[Int32]
+    lock_ptr: cute.Pointer, thread_idx: int | Int32, flag_offset: int, val: cutlass.Constexpr[Int32]
 ) -> None:
     flag_ptr = lock_ptr + flag_offset
     if thread_idx == 0:
         red_release(flag_ptr, val)
-        # red_relaxed(flag_ptr, val)
+        # red_relaxed(flag_ptr, val)