Merge pull request #553 from willow-ahrens/wma/fix-bc-mmm

misc fixes
finch-tensor · May 13, 2024 · 9052e4b · 9052e4b · mtsokol · May 13, 2024
2 parents c1fab78 + e6a175b
commit 9052e4b
Show file tree

Hide file tree

Showing 16 changed files with 615 additions and 451 deletions.
diff --git a/benchmark/runbenchmarks.jl b/benchmark/runbenchmarks.jl
@@ -1,8 +1,9 @@
 #!/usr/bin/env julia
 if abspath(PROGRAM_FILE) == @__FILE__
     using Pkg
-    Pkg.develop(PackageSpec(path = joinpath(@__DIR__, "..")))
     Pkg.activate(@__DIR__)
+    Pkg.develop(PackageSpec(path = joinpath(@__DIR__, "..")))
+    Pkg.resolve()
     Pkg.instantiate()
 end
 

diff --git a/benchmark/runjudge.jl b/benchmark/runjudge.jl
@@ -1,8 +1,9 @@
 #!/usr/bin/env julia
 if abspath(PROGRAM_FILE) == @__FILE__
     using Pkg
-    Pkg.develop(PackageSpec(path = joinpath(@__DIR__, "..")))
     Pkg.activate(@__DIR__)
+    Pkg.develop(PackageSpec(path = joinpath(@__DIR__, "..")))
+    Pkg.resolve()
     Pkg.instantiate()
 end
 

diff --git a/docs/fix.jl b/docs/fix.jl
@@ -2,6 +2,8 @@
 if abspath(PROGRAM_FILE) == @__FILE__
     using Pkg
     Pkg.activate(@__DIR__)
+    Pkg.develop(PackageSpec(path = joinpath(@__DIR__, "..")))
+    Pkg.resolve()
     Pkg.instantiate()
 end
 

diff --git a/docs/make.jl b/docs/make.jl
@@ -2,6 +2,8 @@
 if abspath(PROGRAM_FILE) == @__FILE__
     using Pkg
     Pkg.activate(@__DIR__)
+    Pkg.develop(PackageSpec(path = joinpath(@__DIR__, "..")))
+    Pkg.resolve()
     Pkg.instantiate()
 end
 

diff --git a/src/Finch.jl b/src/Finch.jl
@@ -197,7 +197,7 @@ end
         ))
 
         if @load_preference("precompile", true)
-            @info "Running enhanced precompilation... (to disable, run `using Preferences; Preferences.set_preference(\"Finch\", \"precompile\"=>false)`"
+            @info "Running enhanced precompilation... (to disable, run `using Preferences; Preferences.set_preferences!(\"Finch\", \"precompile\"=>false)`"
             include("../test/precompile.jl")
         end
     end

diff --git a/src/interface/eager.jl b/src/interface/eager.jl
@@ -115,3 +115,11 @@ function LinearAlgebra.norm(arr::AbstractTensorOrBroadcast, p::Real = 2)
         return root(sum(broadcasted(power, broadcasted(norm, arr, p), p)))
     end
 end
+
+"""
+    expanddims(arr::AbstractTensor, dims)
+
+Expand the dimensions of an array by inserting a new singleton axis or axes that
+will appear at the `dims` position in the expanded array shape.
+"""
+expanddims(arr::AbstractTensor, dims) = compute(expanddims(lazy(arr), dims))
diff --git a/src/interface/einsum.jl b/src/interface/einsum.jl
@@ -125,10 +125,10 @@ function (ctx::EinsumParserVisitor)(ex)
                 $(esc(tns)) = $einsum($(esc(op)), $arg, $(map(QuoteNode, idxs)...);$(map(esc, ctx.opts)...),)
             end
         else
-            throw(FinchSyntaxError("Invalid einsum expression: $ex"))
+            throw(FinchNotation.FinchSyntaxError("Invalid einsum expression: $ex"))
         end
     else
-        throw(FinchSyntaxError("Invalid einsum expression type: $ex"))
+        throw(FinchNotation.FinchSyntaxError("Invalid einsum expression type: $ex"))
     end
 end
 

diff --git a/src/interface/lazy.jl b/src/interface/lazy.jl
@@ -26,12 +26,20 @@ function Base.getindex(arr::LazyTensor{T, N}, idxs::Vararg{Union{Nothing, Colon}
     if length(idxs) - count(isnothing, idxs) != N
         throw(ArgumentError("Cannot index a lazy tensor with more or fewer `:` dims than it had original dims."))
     end
-    fields = [field(gensym(:i)) for _ in 1:length(idxs)]
-    original_fields = fields[findall(!isnothing, idxs)]
-    data = reorder(relabel(arr.data, original_fields...), fields...)
-    extrude = [true for _ in 1:length(idxs)]
-    extrude[findall(!isnothing, idxs)] .= arr.extrude
-    return LazyTensor{T}(data, (extrude...,), arr.default)
+    return expanddims(arr, findall(isnothing, idxs))
+end
+
+function expanddims(arr::LazyTensor{T}, dims) where {T}
+    @assert allunique(dims)
+    @assert issubset(dims,1:ndims(arr) + length(dims))
+    antidims = setdiff(1:ndims(arr) + length(dims), dims)
+    idxs_1 = [field(gensym(:i)) for _ in 1:ndims(arr)]
+    idxs_2 = [field(gensym(:i)) for _ in 1:ndims(arr) + length(dims)]
+    idxs_2[antidims] .= idxs_1
+    data_2 = reorder(relabel(arr.data, idxs_1...), idxs_2...)
+    extrude_2 = [false for _ in 1:ndims(arr) + length(dims)]
+    extrude_2[antidims] .= arr.extrude
+    return LazyTensor{T, ndims(arr) + length(dims)}(data_2, tuple(extrude_2...), default(arr))
 end
 
 function identify(data)
@@ -183,7 +191,7 @@ function broadcast_to_query(tns::LazyTensor{T, N}, idxs) where {T, N}
 end
 
 function broadcast_to_extrude(bc::Broadcast.Broadcasted, n)
-    any(map(arg -> broadcast_to_extrude(arg, n), bc.args))
+    all(map(arg -> broadcast_to_extrude(arg, n), bc.args))
 end
 
 function broadcast_to_extrude(tns::LazyTensor, n)
@@ -423,29 +431,60 @@ is fused with the execution of `z + 1`.
 """
 lazy(arg) = LazyTensor(arg)
 
+default_scheduler(;verbose=false) = LogicExecutor(DefaultLogicOptimizer(LogicCompiler()), verbose=verbose)
+
 """
-    fused(f, args...; [optimizer=DefaultOptimizer()])
+    fused(f, args...; kwargs...)
 
 This function decorator modifies `f` to fuse the contained array
 operations and optimize the resulting program. The function must return a single
-array or tuple of arrays. The `optimizer` keyword argument specifies the
-optimizer to use.
+array or tuple of arrays. `kwargs` are passed to [`compute`](@ref)
 """
-function fused(f, args...; optimizer=DefaultOptimizer())
-    compute(f(map(LazyTensor, args...)), optimizer)
+function fused(f, args...; kwargs...)
+    compute(f(map(LazyTensor, args...)), kwargs...)
 end
 
-default_scheduler(;verbose=false) = LogicExecutor(DefaultLogicOptimizer(LogicCompiler()), verbose=verbose)
+current_scheduler = Ref{Any}(default_scheduler())
+
+"""
+    set_scheduler!(scheduler)
+
+Set the current scheduler to `scheduler`. The scheduler is used by `compute` to
+execute lazy tensor programs.
+"""
+set_scheduler!(scheduler) = current_scheduler[] = scheduler
+
+"""
+    get_scheduler()
+
+Get the current Finch scheduler used by `compute` to execute lazy tensor programs.
+"""
+get_scheduler() = current_scheduler[]
+
+"""
+    with_scheduler(f, scheduler)
+
+Execute `f` with the current scheduler set to `scheduler`.
+"""
+function with_scheduler(f, scheduler)
+    old_scheduler = get_scheduler()
+    set_scheduler!(scheduler)
+    try
+        return f()
+    finally
+        set_scheduler!(old_scheduler)
+    end
+end
 
 """
     compute(args..., ctx=default_scheduler()) -> Any
 
 Compute the value of a lazy tensor. The result is the argument itself, or a
 tuple of arguments if multiple arguments are passed.
 """
-compute(args...; ctx=default_scheduler(), kwargs...) = compute_parse(set_options(ctx; kwargs...), map(lazy, args))
-compute(arg; ctx=default_scheduler(), kwargs...) = compute_parse(set_options(ctx; kwargs...), (lazy(arg),))[1]
-compute(args::Tuple; ctx=default_scheduler(), kwargs...) = compute_parse(set_options(ctx; kwargs...), map(lazy, args))
+compute(args...; ctx=get_scheduler(), kwargs...) = compute_parse(set_options(ctx; kwargs...), map(lazy, args))
+compute(arg; ctx=get_scheduler(), kwargs...) = compute_parse(set_options(ctx; kwargs...), (lazy(arg),))[1]
+compute(args::Tuple; ctx=get_scheduler(), kwargs...) = compute_parse(set_options(ctx; kwargs...), map(lazy, args))
 function compute_parse(ctx, args::Tuple)
     args = collect(args)
     vars = map(arg -> alias(gensym(:A)), args)

diff --git a/src/interface/traits.jl b/src/interface/traits.jl
@@ -130,22 +130,31 @@ Return a storage trait object representing the result of mapping `f` over
 storage traits `args`. Assumes representation is collapsed.
 """
 function map_rep(f, args...)
-    map_rep_def(f, map(arg -> pad_data_rep(arg, maximum(ndims, args)), args))
+    map_rep_def(f, map(arg -> paddims_rep(arg, maximum(ndims, args)), args))
 end
 
-pad_data_rep(rep, n) = ndims(rep) < n ? pad_data_rep(ExtrudeData(rep), n) : rep
+paddims_rep(rep, n) = ndims(rep) < n ? paddims_rep(ExtrudeData(rep), n) : rep
 
 """
-    extrude_rep(tns, dims)
-Expand the representation of `tns` to the dimensions `dims`, which must have length(ndims(tns)) and be in ascending order.
+    expanddims_rep(tns, dims)
+Expand the representation of `tns` by inserting singleton dimensions `dims`.
 """
-extrude_rep(tns, dims) = extrude_rep_def(tns, reverse(dims)...)
-extrude_rep_def(tns) = tns
-extrude_rep_def(tns::HollowData, dims...) = HollowData(extrude_rep_def(tns.lvl, dims...))
-extrude_rep_def(tns::SparseData, dim, dims...) = SparseData(pad_data_rep(extrude_rep_def(tns.lvl, dims...), dim - 1))
-extrude_rep_def(tns::RepeatData, dim, dims...) = RepeatData(pad_data_rep(extrude_rep_def(tns.lvl, dims...), dim - 1))
-extrude_rep_def(tns::DenseData, dim, dims...) = DenseData(pad_data_rep(extrude_rep_def(tns.lvl, dims...), dim - 1))
-extrude_rep_def(tns::ExtrudeData, dim, dims...) = ExtrudeData(pad_data_rep(extrude_rep_def(tns.lvl, dims...), dim - 1))
+function expanddims_rep(tns, dims)
+    @assert allunique(dims)
+    @assert issubset(dims,1:ndims(tns) + length(dims))
+    expanddims_rep_def(tns, ndims(tns) + length(dims), dims)
+end
+expanddims_rep_def(tns::HollowData, dim, dims) = HollowData(expanddims_rep_def(tns.lvl, dim, dims))
+expanddims_rep_def(tns::ElementData, dim, dims) =
+    dim in dims ? ExtrudeData(expanddims_rep_def(tns, dim-1, dims)) : tns
+expanddims_rep_def(tns::SparseData, dim, dims) =
+    dim in dims ? ExtrudeData(expanddims_rep_def(tns, dim-1, dims)) : SparseData(expanddims_rep_def(tns.lvl, dim-1, dims))
+expanddims_rep_def(tns::RepeatData, dim, dims) =
+    dim in dims ? ExtrudeData(expanddims_rep_def(tns, dim-1, dims)) : RepeatData(expanddims_rep_def(tns.lvl, dim-1, dims))
+expanddims_rep_def(tns::DenseData, dim, dims) =
+    dim in dims ? ExtrudeData(expanddims_rep_def(tns, dim-1, dims)) : DenseData(expanddims_rep_def(tns.lvl, dim-1, dims))
+expanddims_rep_def(tns::ExtrudeData, dim, dims) =
+    dim in dims ? ExtrudeData(expanddims_rep_def(tns, dim-1, dims)) : ExtrudeData(expanddims_rep_def(tns.lvl, dim-1, dims))
 
 struct MapRepExtrudeStyle end
 struct MapRepSparseStyle end
@@ -309,12 +318,12 @@ function permutedims_rep(tns, perm)
         end
         n += 1
     end
-    src = extrude_rep(tns, src_dims)
+    src = expanddims_rep(tns, setdiff(1:maximum(src_dims, init=0), src_dims))
     for mask_dims in diags
-        mask = extrude_rep(DenseData(SparseData(ElementData(false, Bool))), mask_dims)
-        src = map_rep(filterop(default(src)), pad_data_rep(mask, ndims(src)), src)
+        mask = expanddims_rep(DenseData(SparseData(ElementData(false, Bool))), setdiff(1:ndims(src), mask_dims))
+        src = map_rep(filterop(default(src)), mask, src)
     end
-    aggregate_rep(initwrite(default(tns)), default(tns), src, setdiff(src_dims, dst_dims))
+    res = aggregate_rep(initwrite(default(tns)), default(tns), src, setdiff(src_dims, dst_dims))
 end
 
 """

diff --git a/src/scheduler/LogicCompiler.jl b/src/scheduler/LogicCompiler.jl
@@ -42,10 +42,22 @@ end
     mode = :fast
 end
 
-function finch_pointwise_logic_to_code(ex)
+@kwdef struct PointwiseLowerer
+    bound_idxs = []
+end
+
+function compile_pointwise_logic(ex)
+    ctx = PointwiseLowerer()
+    code = ctx(ex)
+    bound_idxs = ctx.bound_idxs
+    (code, bound_idxs)
+end
+
+function (ctx::PointwiseLowerer)(ex)
     if @capture ex mapjoin(~op, ~args...)
-        :($(op.val)($(map(arg -> finch_pointwise_logic_to_code(arg), args)...)))
+        :($(op.val)($(map(ctx, args)...)))
     elseif (@capture ex reorder(relabel(~arg::isalias, ~idxs_1...), ~idxs_2...))
+        append!(ctx.bound_idxs, idxs_1)
         :($(arg.name)[$(map(idx -> idx in idxs_2 ? idx.name : 1, idxs_1)...)]) #TODO need a trait for the first index
     elseif (@capture ex reorder(~arg::isimmediate, ~idxs...))
         arg.val
@@ -82,12 +94,18 @@ function (ctx::LogicLowerer)(ex)
     elseif @capture ex query(~lhs::isalias, reformat(~tns, reorder(relabel(~arg::isalias, ~idxs_1...), ~idxs_2...)))
         loop_idxs = map(idx -> idx.name, withsubsequence(intersect(idxs_1, idxs_2), idxs_2))
         lhs_idxs = map(idx -> idx.name, idxs_2)
-        rhs = finch_pointwise_logic_to_code(reorder(relabel(arg, idxs_1...), idxs_2...))
+        (rhs, rhs_idxs) = compile_pointwise_logic(reorder(relabel(arg, idxs_1...), idxs_2...))
         body = :($(lhs.name)[$(lhs_idxs...)] = $rhs)
         for idx in loop_idxs
-            body = :(for $idx = _
-                $body
-            end)
+            if field(idx) in rhs_idxs
+                body = :(for $idx = _
+                    $body
+                end)
+            elseif idx in lhs_idxs
+                body = :(for $idx = 1:1
+                    $body
+                end)
+            end
         end
         quote
             $(lhs.name) = $(compile_logic_constant(tns))
@@ -102,13 +120,19 @@ function (ctx::LogicLowerer)(ex)
         ctx(query(lhs, reformat(tns, aggregate(initwrite(z), immediate(z), mapjoin(args...)))))
     elseif @capture ex query(~lhs, reformat(~tns, aggregate(~op, ~init, ~arg, ~idxs_1...)))
         idxs_2 = map(idx -> idx.name, getfields(arg))
-        idxs_3 = map(idx -> idx.name, setdiff(getfields(arg), idxs_1))
-        rhs = finch_pointwise_logic_to_code(arg)
-        body = :($(lhs.name)[$(idxs_3...)] <<$(compile_logic_constant(op))>>= $rhs)
+        lhs_idxs = map(idx -> idx.name, setdiff(getfields(arg), idxs_1))
+        (rhs, rhs_idxs) = compile_pointwise_logic(arg)
+        body = :($(lhs.name)[$(lhs_idxs...)] <<$(compile_logic_constant(op))>>= $rhs)
         for idx in idxs_2
-            body = :(for $idx = _
-                $body
-            end)
+            if field(idx) in rhs_idxs
+                body = :(for $idx = _
+                    $body
+                end)
+            elseif idx in lhs_idxs
+                body = :(for $idx = 1:1
+                    $body
+                end)
+            end
         end
         quote
             $(lhs.name) = $(compile_logic_constant(tns))

diff --git a/src/scheduler/LogicExecutor.jl b/src/scheduler/LogicExecutor.jl
@@ -68,12 +68,13 @@ end
 
 codes = Dict()
 function (ctx::LogicExecutor)(prgm)
-    f = get!(codes, get_structure(prgm)) do
-        eval(logic_executor_code(ctx.ctx, prgm))
+    (f, code) = get!(codes, get_structure(prgm)) do
+        thunk = logic_executor_code(ctx.ctx, prgm)
+        (eval(thunk), thunk)
     end
     if ctx.verbose
         println("Executing:")
-        display(logic_executor_code(ctx.ctx, prgm))
+        display(code)
     end
     return Base.invokelatest(f, prgm)
 end