Format code

Author: mofeing

src/Compiler.jl

@@ -410,9 +410,9 @@ function compile(f, args; pipeline_options="", client=nothing)
 
             res = :($sym.buffer)
             push!(linearized_args, res)
-            
+
             respaths = ((p for p in arg.paths if p[1] != :args)...,)
-            
+
             resarg = false
             for respath in respaths
                 if respath[1] == :result
@@ -440,7 +440,7 @@ function compile(f, args; pipeline_options="", client=nothing)
                 push!(resarg_syncs, usbuf)
             end
         end
-        
+
         for (idx, result) in enumerate(linear_results)
             paths = ((p for p in result.paths if p[1] != :args)...,)
             for path in paths

src/TracedRArray.jl

@@ -16,7 +16,7 @@ mutable struct TracedRArray{T,N} <: RArray{T,N}
     end
 end
 
-function Base.getindex(a::TracedRArray{T,0}) where T
+function Base.getindex(a::TracedRArray{T,0}) where {T}
     return a
 end
 
@@ -75,13 +75,13 @@ end
 function Base.setindex!(
     a::TracedRArray{T,N}, v, indices::Vararg{Union{Base.AbstractUnitRange,Colon},N}
 ) where {T,N}
-    indices = [(promote_to(TracedRArray{Int, 0}, i isa Colon ? 1 : first(i))-1).mlir_data for i in indices]
+    indices = [
+        (promote_to(TracedRArray{Int,0}, i isa Colon ? 1 : first(i)) - 1).mlir_data for
+        i in indices
+    ]
     v = promote_to(TracedRArray{T,N}, v)
     res = MLIR.IR.result(
-        MLIR.Dialects.stablehlo.dynamic_update_slice(
-           a.mlir_data, v.mlir_data, indices
-        ),
-        1,
+        MLIR.Dialects.stablehlo.dynamic_update_slice(a.mlir_data, v.mlir_data, indices), 1
     )
     a.mlir_data = res
     return v
@@ -162,7 +162,7 @@ end
 
 function promote_to(::Type{TracedRArray{T,N}}, rhs) where {T,N}
     if isa(rhs, TracedRArray)
-        if typeof(rhs) == TracedRArray{T, N}
+        if typeof(rhs) == TracedRArray{T,N}
             return rhs
         end
         return TracedRArray{T,N}(
@@ -185,9 +185,12 @@ function promote_to(::Type{TracedRArray{T,N}}, rhs) where {T,N}
     end
     T0 = eltype(rhs)
     attr = MLIR.IR.DenseElementsAttribute(collect(rhs))
-    return promote_to(TracedRArray{T, N}, TracedRArray{T0,length(size(rhs))}(
-        (), MLIR.IR.result(MLIR.Dialects.stablehlo.constant(; value=attr), 1), size(rhs)
-    ))
+    return promote_to(
+        TracedRArray{T,N},
+        TracedRArray{T0,length(size(rhs))}(
+            (), MLIR.IR.result(MLIR.Dialects.stablehlo.constant(; value=attr), 1), size(rhs)
+        ),
+    )
 end
 
 function promote_to(lhs::TracedRArray{T,N}, rhs) where {T,N}
@@ -491,31 +494,32 @@ for (jlop, hloop, hlocomp, merge) in (
         function elem_apply(
             fn::typeof($jlop), @nospecialize(lhs::TracedRArray{T,N}), @nospecialize(rhs)
         ) where {T,N}
-            elem_apply(fn, lhs, promote_to(lhs, rhs))
+            return elem_apply(fn, lhs, promote_to(lhs, rhs))
         end
 
         function elem_apply(
             ::typeof($jlop), @nospecialize(lhs), @nospecialize(rhs::TracedRArray{T,N})
         ) where {T,N}
-            elem_apply(fn, promote_to(rhs, lhs), rhs)
+            return elem_apply(fn, promote_to(rhs, lhs), rhs)
         end
 
-        function $jlop(@nospecialize(lhs::TracedRArray{T,N}), @nospecialize(rhs)
-        ) where {T, N}
-            $jlop(lhs, promote_to(lhs, rhs))
+        function $jlop(
+            @nospecialize(lhs::TracedRArray{T,N}), @nospecialize(rhs)
+        ) where {T,N}
+            return $jlop(lhs, promote_to(lhs, rhs))
         end
 
-        function $jlop(@nospecialize(lhs), @nospecialize(rhs::TracedRArray{T,N})
-        ) where {T, N}
-            $jlop(promote_to(rhs, lhs), rhs)
+        function $jlop(
+            @nospecialize(lhs), @nospecialize(rhs::TracedRArray{T,N})
+        ) where {T,N}
+            return $jlop(promote_to(rhs, lhs), rhs)
         end
     end
-    
+
     if merge != nothing
         @eval begin
             function $jlop(
-                @nospecialize(lhs::TracedRArray{T,N}),
-                @nospecialize(rhs::TracedRArray{T,N})
+                @nospecialize(lhs::TracedRArray{T,N}), @nospecialize(rhs::TracedRArray{T,N})
             ) where {T,N}
                 elems = elem_apply($jlop, lhs, rhs)
                 if N == 0
@@ -528,10 +532,9 @@ for (jlop, hloop, hlocomp, merge) in (
     else
         @eval begin
             function $jlop(
-                @nospecialize(lhs::TracedRArray{T,0}),
-                @nospecialize(rhs::TracedRArray{T,0})
+                @nospecialize(lhs::TracedRArray{T,0}), @nospecialize(rhs::TracedRArray{T,0})
             ) where {T}
-                elem_apply($jlop, lhs, rhs)
+                return elem_apply($jlop, lhs, rhs)
             end
         end
     end

src/Tracing.jl

@@ -41,7 +41,7 @@ function traced_type(::Type{T}, seen, mode) where {T<:Function}
     changed = false
     traced_fieldtypes = ntuple(Val(N)) do i
         next = traced_type(fieldtype(T, i), seen, mode)
-        changed |=  next != fieldtype(T, i)
+        changed |= next != fieldtype(T, i)
         next
     end
 

test/basic.jl

@@ -246,15 +246,15 @@ end
     # get_view_compiled = Reactant.compile(get_view, (x_concrete,))
 end
 
-tuple_byref(x) = (; a =(; b=x))
+tuple_byref(x) = (; a=(; b=x))
 tuple_byref2(x) = abs2.(x), tuple_byref2(x)
 
 @testset "Tuple byref" begin
     x = Reactant.to_rarray([1.0 -2.0; -3.0 4.0])
     f1 = Reactant.compile(tuple_byref, (x,))
     r1 = f1(x)
     @test r1.a.b.data === x.data
-    
+
     # TODO this seems to hang during compile
     # f2 = Reactant.compile(tuple_byref2, (x,))
     # r2 = f2(x)