Mask Sugar

In Finch, expressions like i == j are treated as a sugar for mask tensors, which can be used to encode fancy iteration patterns. For example, the expression i == j is converted to a diagonal boolean mask tensor DiagMask()[i, j], which allows an expression like

    @finch begin
        for i=_, j=_
            if i == j
                s[] += A[i, j]
            end
        end
    end
end

to compile to something like

    for i = 1:n
        s[] += A[i, i]
    end

There are several mask tensors and syntaxes available, summarized in the following table where i, j are indices:

ExpressionTransformed Expression
i < jUpTriMask()[i, j - 1]
i <= jUpTriMask()[i, j]
i > jLoTriMask()[i, j + 1]
i >= jLoTriMask()[i, j]
i == jDiagMask()[i, j]
i != j!(DiagMask()[i, j])

Note that either i or j may be expressions, so long as the expression is constant with respect to the loop over the index.

The mask tensors are described below:

Finch.uptrimaskConstant
uptrimask

A mask for an upper triangular tensor, uptrimask[i, j] = i <= j. Note that this specializes each column for the cases where i <= j and i > j.

Finch.lotrimaskConstant
lotrimask

A mask for an upper triangular tensor, lotrimask[i, j] = i >= j. Note that this specializes each column for the cases where i < j and i >= j.

Finch.diagmaskConstant
diagmask

A mask for a diagonal tensor, diagmask[i, j] = i == j. Note that this specializes each column for the cases where i < j, i == j, and i > j.

Finch.bandmaskConstant
bandmask

A mask for a banded tensor, bandmask[i, j, k] = j <= i <= k. Note that this specializes each column for the cases where i < j, j <= i <= k, and k < i.

Finch.chunkmaskFunction
chunkmask(b)

A mask for a chunked tensor, chunkmask[i, j] = b * (j - 1) < i <= b * j. Note that this specializes each column for the cases where i < b * (j - 1), `b * (j

    1. < i <= b * j, andb * j < i`.