Add more documentation on the internals

Author: timholy

Project.toml

@@ -4,6 +4,7 @@ authors = ["Chris Elrod <[email protected]>"]
 version = "0.7.0"
 
 [deps]
+DocStringExtensions = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 OffsetArrays = "6fe1bfb0-de20-5000-8ca7-80f57d26f881"
 SIMDPirates = "21efa798-c60a-11e8-04d3-e1a92915a26a"
@@ -12,6 +13,7 @@ UnPack = "3a884ed6-31ef-47d7-9d2a-63182c4928ed"
 VectorizationBase = "3d5dd08c-fd9d-11e8-17fa-ed2836048c2f"
 
 [compat]
+DocStringExtensions = "0.8"
 OffsetArrays = "1"
 SIMDPirates = "0.7.13"
 SLEEFPirates = "0.4.4"

docs/make.jl

@@ -21,7 +21,8 @@ makedocs(;
             "devdocs/loopset_structure.md",
             "devdocs/constructing_loopsets.md",
             "devdocs/evaluating_loops.md",
-            "devdocs/lowering.md"
+            "devdocs/lowering.md",
+            "devdocs/reference.md"
         ]
     ],
     # repo="https://github.com/chriselrod/LoopVectorization.jl/blob/{commit}{path}#L{line}",

docs/src/api.md

@@ -22,5 +22,5 @@ vmapntt!
 
 ```@docs
 vfilter
-vfilter!
+LoopVectorization.vfilter!
 ```

docs/src/devdocs/reference.md

@@ -0,0 +1,42 @@
+# Internals reference
+
+## Operation types
+
+```@docs
+LoopVectorization.OperationType
+LoopVectorization.constant
+LoopVectorization.memload
+LoopVectorization.compute
+LoopVectorization.memstore
+LoopVectorization.loopvalue
+```
+
+## Operation
+
+```@docs
+LoopVectorization.Operation
+```
+
+## Instructions and costs
+
+```@docs
+LoopVectorization.Instruction
+LoopVectorization.InstructionCost
+```
+
+## Array references
+
+```@docs
+LoopVectorization.ArrayReference
+LoopVectorization.ArrayReferenceMeta
+```
+
+## Condensed types
+
+These are used when encoding the `@avx` block as a type parameter for passing through
+to the `@generated` function.
+
+```@docs
+LoopVectorization.ArrayRefStruct
+LoopVectorization.OperationStruct
+```

src/LoopVectorization.jl

@@ -16,6 +16,7 @@ using SLEEFPirates: pow
 using Base.Broadcast: Broadcasted, DefaultArrayStyle
 using LinearAlgebra: Adjoint, Transpose
 using Base.Meta: isexpr
+using DocStringExtensions
 
 
 const SUPPORTED_TYPES = Union{Float16,Float32,Float64,Integer}

src/condense_loopset.jl

@@ -5,7 +5,9 @@ Base.:|(u::Unsigned, it::IndexType) = u | UInt8(it)
 Base.:(==)(u::Unsigned, it::IndexType) = (u % UInt8) == UInt8(it)
 
 """
-`ArrayRefStruct` stores a representation of an array-reference expression such as `A[i,j]`.
+    ArrayRefStruct
+
+A condensed representation of an [`ArrayReference`](@ref).
 It supports array-references with up to 8 indexes, where the data for each consecutive index is packed into corresponding 8-bit fields
 of `index_types` (storing the enum `IndexType`), `indices` (the `id` for each index symbol), and `offsets` (currently unused).
 """
@@ -53,6 +55,11 @@ function ArrayRefStruct(ls::LoopSet, mref::ArrayReferenceMeta, arraysymbolinds::
     ArrayRefStruct{mref.ref.array,mref.ptr}( index_types, indices, offsets )
 end
 
+"""
+    OperationStruct
+
+A condensed representation of an [`Operation`](@ref).
+"""
 struct OperationStruct <: AbstractLoopOperation
     # instruction::Instruction
     loopdeps::UInt64

src/constructors.jl

@@ -90,6 +90,29 @@ julia> b ≈ c
 true
 ```
 
+# Extended help
+
+Advanced users can customize the implementation of the `@avx`-annotated block
+using keyword arguments:
+
+```
+@avx inline=false unroll=2 body
+```
+
+where `body` is the code of the block (e.g., `for ... end`).
+
+`inline` is a Boolean. When `true` (the default), `body` will be directly inlined
+into the function (via a forced-inlining call to `_avx_!`).
+When `false`, it will call `__avx__!` instead, letting Julia's own inlining engine
+determine whether the call to `__avx__!` should be inlined. (Typically, it won't.)
+In priniciple, first calling `__avx__!` (which itself calls `_avx_!`) can sometimes
+allow better code generation.
+One can find some circumstances where `inline=true` is faster, and other circumstances
+where `inline=false` is faster, so the best setting may require experimentation.
+
+`unroll` is an integer that specifies the loop unrolling factor, or a
+tuple `(4, 2)` signaling that the generated code should unroll more than
+one loop.
 """
 macro avx(q)
     q = macroexpand(__module__, q)
@@ -121,7 +144,7 @@ function check_unroll(arg)
             T = convert(Int8, tup.args[2])
         else
             return nothing
-        end    
+        end
     else
         return nothing
     end
@@ -179,4 +202,3 @@ macro avx_debug(q)
     q = macroexpand(__module__, q)
     esc(LoopVectorization.setup_call_debug(LoopSet(q, __module__)))
 end
-

src/costs.jl

@@ -6,7 +6,12 @@ else
 end
 
 
+"""
+    Instruction
 
+`Instruction` represents a function via its module and symbol. It is
+similar to a `GlobalRef` and may someday be replaced by `GlobalRef`.
+"""
 struct Instruction
     mod::Symbol
     instr::Symbol
@@ -36,10 +41,24 @@ Base.isequal(ins1::Instruction, ins2::Instruction) = (ins1.instr === ins2.instr)
 
 const LOOPCONSTANT = Instruction(Symbol("LOOPCONSTANTINSTRUCTION"))
 
+"""
+    InstructionCost
+
+Store parameters related to performance for individual CPU instructions.
+
+$(TYPEDFIELDS)
+"""
 struct InstructionCost
+    "A flag indicating how instruction cost scales with vector width (128, 256, or 512 bits)"
     scaling::Float64 # sentinel values: -3 == no scaling; -2 == offset_scaling, -1 == linear scaling, >0 ->  == latency == reciprocal throughput
+    """The number of clock cycles per operation when many of the same operation are repeated in sequence.
+    Think of it as the inverse of the flow rate at steady-state. It is typically ≤ the `scalar_latency`."""
     scalar_reciprocal_throughput::Float64
+    """The minimum delay, in clock cycles, associated with the instruction.
+    Think of it as the delay from turning on a faucet to when water starts coming out the end of the pipe.
+    See also `scalar_reciprocal_throughput`."""
     scalar_latency::Int
+    "Number of floating-point registered used"
     register_pressure::Int
 end
 InstructionCost(sl::Int, srt::Float64, scaling::Float64 = -3.0) = InstructionCost(scaling, srt, sl, 0)

src/operations.jl

@@ -1,7 +1,21 @@
+"""
+    ArrayReference
+
+A type for encoding an array reference `A[i,j]` occurring inside an `@avx` block.
+
+# Fields
 
+$(TYPEDFIELDS)
+"""
 struct ArrayReference
+    "The array variable"
     array::Symbol
+    "The list of indices (e.g., `[:i, :j]`), or `name(op)` for computed indices."
     indices::Vector{Symbol}
+    """Index offset, e.g., `a[i+7]` would store the `7`. `offsets` is also used
+    to help identify opportunities for avoiding reloads, for example in `y[i] = x[i] - x[i-1]`,
+    the previous load `x[i-1]` can be "carried over" to the next iteration.
+    Only used for small (`Int8`) offsets."""    
     offsets::Vector{Int8}
 end
 ArrayReference(array, indices) = ArrayReference(array, indices, zeros(Int8, length(indices)))
@@ -26,9 +40,21 @@ function Base.isequal(x::ArrayReference, y::ArrayReference)
     end
     true
 end
+"""
+    ArrayReferenceMeta
+
+A type similar to [`ArrayReference`](@ref) but holding additional information.
+
+# Fields
+
+$(TYPEDFIELDS)
+"""
 struct ArrayReferenceMeta
+    "The `ArrayReference`"
     ref::ArrayReference
+    "A vector of Bools indicating whether each index is a loop variable (`false` for operation-computed indices)"
     loopedindex::Vector{Bool}
+    "Variable holding the pointer to the array's underlying storage"
     ptr::Symbol
 end
 function ArrayReferenceMeta(ref::ArrayReference, loopedindex, ptr = vptr(ref))
@@ -56,29 +82,89 @@ Base.:(==)(x::ArrayReferenceMeta, y) = false
 
 abstract type AbstractLoopOperation end
 
+"""
+`OperationType` is an `@enum` for classifying supported operations that can appear in
+`@avx` blocks. Type `LoopVectorization.OperationType` to see the different types.
+"""
 @enum OperationType begin
     constant
     memload
     compute
     memstore
     loopvalue
 end
+"An operation setting a variable to a constant value (e.g., `a = 0.0`)" constant
+"An operation setting a variable from a memory location (e.g., `a = A[i,j]`)" memload
+"An operation computing a new value from one or more variables (e.g., `a = b + c`)" compute
+"An operation storing a value to a memory location (e.g., `A[i,j] = a`)" memstore
+"""
+`loopvalue` indicates an loop variable (`i` in `for i in ...`). These are the "parents" of `compute`
+operations that involve the loop variables.
+"""
+loopvalue
 
 # TODO: can some computations be cached in the operations?
 """
+    Operation
+
+A structure to encode a particular action occuring inside an `@avx` block.
+
+# Fields
+
+$(TYPEDFIELDS)
+
+# Example
+
+```jldoctest Operation; filter = r"\\"##.*\\""
+julia> using LoopVectorization
+
+julia> AmulBq = :(for m ∈ 1:M, n ∈ 1:N
+           C[m,n] = zero(eltype(B))
+           for k ∈ 1:K
+               C[m,n] += A[m,k] * B[k,n]
+           end
+       end);
+
+julia> lsAmulB = LoopVectorization.LoopSet(AmulBq);
+
+julia> LoopVectorization.operations(lsAmulB)
+6-element Array{LoopVectorization.Operation,1}:
+ var"##RHS#253" = var"##zero#254"
+ C[m, n] = var"##RHS#253"
+ var"##tempload#255" = A[m, k]
+ var"##tempload#256" = B[k, n]
+ var"##RHS#253" = LoopVectorization.vfmadd_fast(var"##tempload#255", var"##tempload#256", var"##RHS#253")
+ var"##RHS#253" = LoopVectorization.identity(var"##RHS#253")
+```
+Each one of these lines is a pretty-printed `Operation`.
 """
 mutable struct Operation <: AbstractLoopOperation
+    """A unique identifier for this operation.
+    `identifer(op::Operation)` returns the index of this operation within `operations(ls::LoopSet)`."""
     identifier::Int
+    """The name of the variable storing the result of this operation.
+    For `a = val` this would be `:a`. For array assignments `A[i,j] = val` this would be `:A`."""
     variable::Symbol
+    "Intended to be the size of the result, in bytes. Often inaccurate, not to be relied on."
     elementbytes::Int
+    "The specific operator, e.g., `identity` or `+`"
     instruction::Instruction
+    "The [`OperationType`](@ref) associated with this operation"
     node_type::OperationType
+    "The loop variables this operation depends on"
     dependencies::Vector{Symbol}
+    "Additional loop dependencies that must execute before this operation can be performed successfully (often needed in reductions)"
     reduced_deps::Vector{Symbol}
+    "Operations whose result this operation depends on"
     parents::Vector{Operation}
+    "For `memload` or `memstore`, encodes the array location"
     ref::ArrayReferenceMeta
+    "`gensymmed` name of result."
     mangledvariable::Symbol
+    """Loop variables that *consumers* of this operation depend on.
+    Often used in reductions to replicate assignment of initializers when unrolling."""
     reduced_children::Vector{Symbol}
+
     function Operation(
         identifier::Int,
         variable,
@@ -129,7 +215,7 @@ const NOPARENTS = Operation[]
 function Base.show(io::IO, op::Operation)
     if isconstant(op)
         if op.instruction === LOOPCONSTANT
-            
+
             print(io, Expr(:(=), op.variable, 0))
         else
             print(io, Expr(:(=), op.variable, op.instruction.instr))

src/reconstruct_loopset.jl

@@ -417,8 +417,26 @@ function _avx_loopset(OPSsv, ARFsv, AMsv, LPSYMsv, LBsv, vargs)
         AMsv, LPSYMsv, LBsv, vargs
     )
 end
+const _body_ = Ref{Any}(nothing)
+"""
+    _avx_!(ut, ops, arf, am, lpsym, lb, vargs...)
+
+Execute an `@avx` block. The block's code is represented via the arguments:
+- `ut` is `Val((U,T))`, where `U` is the unrolling factor and `T` ?has something to do with tiling?
+- `ops` is `Tuple{mod1, sym1, op1, mod2, sym2, op2...}` encoding the operations of the loop.
+  `mod` and `sym` encode the module and symbol of the called function; `op` is an [`OperationStruct`](@ref)
+  encoding the details of the operation.
+- `arf` is `Tuple{arf1, arf2...}`, where each `arfi` is an [`ArrayRefStruct`](@ref) encoding
+  an array reference.
+- `am` contains miscellaneous data about the LoopSet (see `process_metadata!`)
+- `lpsym` is `Tuple{:i,:j,...}`, a Tuple of the "loop symbols", i.e. the item variable `i` in `for i ∈ iter`
+- `lb` is `Tuple{RngTypei,RngTypej,...}`, a Tuple encoding syntactically-knowable information about
+  the iterators corresponding to `lpsym`. For example, in `for i ∈ 1:n`, the `1:n` would be encoded with
+  `StaticLowerUnitRange(1)` because the lower bound of the iterator can be determined to be 1.
+- `vargs...` holds the encoded pointers of all the arrays (see `VectorizationBase`'s various pointer types).
+"""
 @generated function _avx_!(::Val{UT}, ::Type{OPS}, ::Type{ARF}, ::Type{AM}, ::Type{LPSYM}, lb::LB, vargs...) where {UT, OPS, ARF, AM, LPSYM, LB}
     1 + 1 # Irrelevant line you can comment out/in to force recompilation...
     ls = _avx_loopset(OPS.parameters, ARF.parameters, AM.parameters, LPSYM.parameters, LB.parameters, vargs)
-    avx_body(ls, UT)
+    return _body_[] = copy(avx_body(ls, UT))
 end