carray-jit 0.1.0 → 0.1.2

data/lib/carray/jit.rb

@@ -216,13 +216,14 @@ class CArray
     JIT.run_stencil(arrays, block, border, type, into)
   end
 
-  # Returns the contraction the block writes: an index that appears twice
-  # is summed.
+  # Returns the contraction the block writes: a repeated index is summed.
   #
   #   CArray.jit_contract { |i, j, k| c[i,j] = a[i,k] * b[k,j] }
   #
   # Every block parameter is an index.  The ones that appear on the left are
-  # the cells written; the rest -- `k` here -- are summed over.
+  # the cells written; the rest -- `k` here -- are summed over.  An index that
+  # repeats is summed however often it repeats: `q[i,i,i]` is one index read
+  # at three positions, and the sum runs along the cube's long diagonal.
   #
   # No extent is given, because every index's extent is fixed by the axes it
   # addresses; an index whose axes disagree is an error, which is the shape
@@ -236,17 +237,61 @@ class CArray
   # Assigning into an array of your own says where to put it, and in what
   # order its axes lie; it does not decide what is summed.
   #
+  # That a repeated index is summed is a statement about *dimensions*,
+  # which is the world the notation comes from: two dimensions met is an inner
+  # product, and there is no other reading.  An index that numbers things --
+  # a point, a sample, a batch -- is not a dimension, and `x[p,k] * y[p,k]`
+  # repeating `p` says "the same point", not "sum over points".  Naming the
+  # result's axes says which is meant:
+  #
+  #   CArray.jit_contract(:p) { |k| x[p,k] * y[p,k] }               # one number per point
+  #   CArray.jit_contract(:a) { square[a,a] }                       # the diagonal, not the trace
+  #   CArray.jit_contract(:b, :i, :j) { |k| u[b,i,k] * v[b,k,j] }   # a batch of products
+  #
+  # The arguments are the result's axes, in that order.  What they say is
+  # which indices are free; what a repetition means is unchanged.  So the rule
+  # is the convention's, with a third clause: an index that repeats is summed,
+  # one that appears once is free, and a named one is free however often it
+  # appears -- which is what puts the diagonal and the per-point quantity
+  # inside the notation instead of outside it.  A free index needs somewhere
+  # to go, so a parameter at a single position is refused once the axes are
+  # named: it is free, and the axes are already stated.  The list is all of
+  # the result's axes rather than some of them -- name one and you have named
+  # them all -- which is what keeps the order it states complete.  Naming is allowed
+  # even where the convention would have reached the same answer, which is how
+  # the result's axes are put in another order.
+  #
+  # The sum is split into partial ones, as `jit_for`'s reduction is: a
+  # contraction says which indices are summed and nothing about the order, so
+  # there is no order here to override.  `CArray::JIT.reassociate = false`
+  # asks for the serial one, which is what a Ruby loop would take.
+  #
   # The block is read and compiled, never called, so it is not yielded to.
   #
+  # @param free_indices [Array<Symbol>] the result's axes, in order; empty to
+  #   let the convention decide, which is an index appearing once.
   # @return [CArray, CArray::JIT::CompiledKernel] the allocated result when the block
   #   assigns into nothing, otherwise the compiled kernel.
   # @raise [CArray::JIT::Unsupported] when the block falls outside the
   #   recognized subset, or an index's axes disagree.
-  def self.jit_contract (&block)
+  def self.jit_contract (*free_indices, &block)
     unless block
       raise JIT::Unsupported, "jit_contract needs a block"
     end
-    JIT.run_contraction(block)
+    unless free_indices.all? { |name| name.is_a?(Symbol) }
+      raise JIT::Unsupported,
+            "jit_contract's arguments are the result's axes, named as symbols, " \
+            "as in `CArray.jit_contract(:p) { |k| x[p,k] * y[p,k] }`"
+    end
+    repeated = free_indices.tally.select { |_, count| count > 1 }.keys
+    unless repeated.empty?
+      raise JIT::Unsupported,
+            "#{repeated.map { |name| "`#{name}`" }.join(', ')} names more than " \
+            "one axis of the result; each axis is one index"
+    end
+    # Nothing named is the convention; naming none of them is a contraction
+    # to a single number, and the two are different statements.
+    JIT.run_contraction(block, free_indices.empty? ? nil : free_indices)
   end
 
   # @!endgroup
@@ -328,12 +373,254 @@ class CArray
       RESULT = :__contraction_result
 
       # @private
-      def run_contraction (block)
+      def run_contraction (block, free_indices = nil)
         node, source, origin = read_block(block)
-        names = capture_names(source, node)
+        # An index named at the call site is not a parameter of the block, so
+        # the block reaches for it the way it reaches for a captured value.
+        # It is neither: it is an index, and it is answered here.
+        names = capture_names(source, node) - (free_indices || [])
+        arrays, scalars, c_functions = split_captures(names, binding_of(block))
+        contract(source, arrays, free_indices, node: node, origin: origin,
+                 scalars: scalars, c_functions: c_functions)
+      end
+
+      # Reads a block as a contraction and returns the terms it is a product
+      # of, or nil when it is not one:
+      #
+      #   CArray::JIT.contraction_of(proc { |i, j, k| a[i,k] * b[k,j] })
+      #   #=> { :terms => [[a, [:i, :k]], [b, [:k, :j]]],
+      #   #     :free => [:i, :j], :summed => [:k], :scale => 1 }
+      #
+      # A number multiplied into the product is not a term -- it has no
+      # indices and no cell -- so it comes back as `:scale`, which is 1 where
+      # there is none.  `a[i,k] * b[k,j] * 2.0` is the same contraction scaled,
+      # and a caller that rearranges it has to put the scale back: multiplying
+      # a sum by a number and multiplying each of its terms are the same
+      # arithmetic, but not the same rounding, and they are not the same
+      # computation type either where the number is wider than the cells.
+      #
+      # This is the half of `jit_contract` that decides *what* is being
+      # computed, without compiling anything, for a caller that wants to
+      # rearrange it -- to contract two terms at a time, say, in an order it
+      # chose -- and then reach `contract_terms` with the pieces.
+      #
+      # Nil means "there is nothing here to rearrange", not that the block is
+      # wrong: a summand that is more than a product of element reads
+      # (`Math.exp(a[i,k]) * b[k,j]`, a division, a captured scalar, an index
+      # with an offset), or one that assigns into an array of its own. Such a
+      # block is still a contraction and `jit_contract` still compiles it; it
+      # is just not a product to be taken apart. A block that is not a
+      # contraction at all raises here, as it would there.
+      #
+      # @param block [Proc] the block, read and not called.
+      # @param free_indices [Array<Symbol>] the result's axes, as
+      #   `jit_contract` takes them.
+      # @return [Hash, nil] `{ terms:, free:, summed: }`, or nil.
+      # @raise [CArray::JIT::Unsupported] when the block is not a contraction.
+      def contraction_of (block, *free_indices)
+        node, source, = read_block(block)
+        # As `jit_contract` reads them: naming none is the convention, which
+        # is not the same as naming an empty list of axes.
+        free_indices = nil if free_indices.empty?
+        names = capture_names(source, node) - (free_indices || [])
         arrays, scalars, c_functions = split_captures(names, binding_of(block))
+        # A compiled function in the summand is not something the structure
+        # carries; a captured number is, as the scale.
+        return nil unless c_functions.empty?
+
+        analyzer = Analyzer.new(source, node: node, array_names: arrays.keys,
+                                contract: :probe, free_indices: free_indices,
+                                cell_names: cell_names(arrays))
+        statements = analyzer.body.statements
+        # Locals before the summand are computation the structure cannot
+        # carry either; an ElementWrite last is the assigned form.
+        return nil unless statements.size == 1
+        factors = product_terms(statements.last, arrays, scalars)
+        return nil unless factors
+        terms, scale = factors
+        return nil if terms.empty?
+
+        free = analyzer.probe_free_names
+        { :terms => terms, :free => free,
+          :summed => terms.flat_map(&:last).uniq - free, :scale => scale }
+      end
+
+      # The product's terms and what multiplies them, or nil where the tree is
+      # anything but a product of cells and numbers.
+      def product_terms (node, arrays, scalars)
+        case node
+        when BinaryOperation
+          return nil unless node.operator == :*
+          left = product_terms(node.left, arrays, scalars)
+          right = left && product_terms(node.right, arrays, scalars)
+          right && [left.first + right.first, left.last * right.last]
+        when ElementRead
+          subscripts = node.subscripts
+          return nil unless subscripts.all? { |index, offset|
+            index && offset.is_a?(Integer) && offset.zero?
+          }
+          [[[arrays.fetch(node.array), subscripts.map(&:first)]], 1]
+        when IntegerLiteral, FloatLiteral
+          [[], node.value]
+        when CaptureRead
+          value = scalars[node.name]
+          value.is_a?(Numeric) ? [[], value] : nil
+        end
+      end
+
+      # @private
+      TERM_PREFIX = "term"
+      # @private
+      TARGET = :target
+
+      # Runs the contraction the terms describe, where a term is an array and
+      # the indices it is read at:
+      #
+      #   CArray::JIT.contract_terms([[a, [:i, :k]], [b, [:k, :j]]],
+      #                              free: [:i, :j])            # a matrix product
+      #
+      # It is `jit_contract` with the block already taken apart -- the same
+      # rules, the same errors, the same kernels -- for a caller that has the
+      # structure rather than a block to read.  A contraction that was decided
+      # rather than written has no source, which is what this is for.
+      #
+      # `free:` is the result's axes in order, and is required: a structure has
+      # no parameter list, so there is nowhere else for the order to be said.
+      # Every index that is not named is summed, and must appear at more than
+      # one position, exactly as in a block whose axes are named.
+      #
+      # The terms are a product of element reads and nothing else. A summand
+      # that is more than that -- `Math.exp(a[i,k]) * b[k,j]`, a division, a
+      # captured scalar, an index with an offset -- is written as a block and
+      # compiled by `jit_contract`; there is no structure here that says it.
+      #
+      # @param terms [Array<Array>] `[array, [index, ...]]` pairs.
+      # @param free [Array<Symbol>] the result's axes, in order.
+      # @param into [CArray, nil] an array of yours to write into, which then
+      #   decides the axis order and is what comes back.
+      # @return [CArray] `into` when it is given, otherwise a new array.
+      # @raise [CArray::JIT::Unsupported] when the terms are malformed, or the
+      #   contraction they describe is one the compiler refuses.
+      def contract_terms (terms, free:, into: nil)
+        names = check_terms(terms)
+        indices = terms.flat_map { |_, subscripts| subscripts }
+        free = check_free_indices(free, indices)
+        arrays = names.zip(terms.map(&:first)).to_h
+
+        factors = terms.each_with_index.map { |(_, subscripts), position|
+          "#{names[position]}[#{subscripts.join(',')}]"
+        }
+        body = factors.join(" * ")
+        if into
+          check_destination(into, free)
+          arrays[TARGET] = into
+          # With every index summed the result is one number, which lives in a
+          # one-cell array at a fixed subscript -- as it is written in a block.
+          body = "#{TARGET}[#{free.empty? ? '0' : free.join(',')}] = #{body}"
+        end
+        summed = indices.uniq - free
+        parameters = summed.empty? ? "" : "|#{summed.join(', ')}| "
+
+        # The source the block would have had.  It is what the kernel cache is
+        # keyed by, so the same terms under the same shapes reach the same
+        # kernel however they were arrived at -- and the string is canonical,
+        # which is the same service einsum's subscripts perform.
+        result = contract("proc { #{parameters}#{body} }", arrays, free)
+        into || result
+      end
 
-        result = allocate_result(source, node, arrays, scalars)
+      # The names the synthesized source gives the terms.  They are the
+      # source's own, so nothing outside chose them -- but an index may
+      # collide with one, and a collision would silently read an array as an
+      # index.
+      def check_terms (terms)
+        unless terms.is_a?(Array) && !terms.empty?
+          raise Unsupported,
+                "contract_terms takes the terms of a product, as in " \
+                "`[[a, [:i, :k]], [b, [:k, :j]]]`"
+        end
+        terms.each do |term|
+          unless term.is_a?(Array) && term.size == 2 && term.first.is_a?(CArray)
+            raise Unsupported,
+                  "a term is an array and the indices it is read at, as in " \
+                  "`[a, [:i, :k]]`; got #{term.inspect}"
+          end
+          array, subscripts = term
+          unless subscripts.is_a?(Array) && subscripts.all? { |name| index_name?(name) }
+            raise Unsupported,
+                  "the indices of a term are symbols naming its axes, as in " \
+                  "`[a, [:i, :k]]`; got #{subscripts.inspect}"
+          end
+          unless subscripts.size == array.rank
+            raise Unsupported,
+                  "a term names one index per axis: this array has rank " \
+                  "#{array.rank} and #{subscripts.size} " \
+                  "#{subscripts.size == 1 ? 'index' : 'indices'} were given"
+          end
+        end
+        names = terms.each_index.map { |position| :"#{TERM_PREFIX}#{position}" }
+        reserved = (names + [TARGET]) & terms.flat_map { |_, subscripts| subscripts }
+        unless reserved.empty?
+          raise Unsupported,
+                "#{reserved.map { |name| "`#{name}`" }.join(', ')} names a term " \
+                "here and cannot also be an index"
+        end
+        names
+      end
+
+      def check_free_indices (free, indices)
+        unless free.is_a?(Array) && free.all? { |name| index_name?(name) }
+          raise Unsupported,
+                "`free:` is the result's axes in order, named as symbols"
+        end
+        repeated = free.tally.select { |_, count| count > 1 }.keys
+        unless repeated.empty?
+          raise Unsupported,
+                "#{repeated.map { |name| "`#{name}`" }.join(', ')} names more " \
+                "than one axis of the result; each axis is one index"
+        end
+        missing = free - indices
+        unless missing.empty?
+          raise Unsupported,
+                "#{missing.map { |name| "`#{name}`" }.join(', ')} " \
+                "#{missing.size == 1 ? 'names no axis' : 'name no axis'} here"
+        end
+        free
+      end
+
+      def check_destination (into, free)
+        unless into.is_a?(CArray)
+          raise Unsupported, "`into:` is an array to write into"
+        end
+        expected = free.empty? ? 1 : free.size
+        unless into.rank == expected
+          raise Unsupported,
+                "`into:` has rank #{into.rank}, and the result's axes are " \
+                "#{free.empty? ? 'none, which is one cell' : free.map { |name| "`#{name}`" }.join(', ')}"
+        end
+      end
+
+      # An index is interpolated into the source this writes, so it has to be
+      # a name Ruby reads back as a local variable.  `:end` and `:do` would
+      # not parse at all, and `:nil` and `:self` would come back as something
+      # else -- all of them as an error about a source the caller never wrote.
+      def index_name? (name)
+        return false unless name.is_a?(Symbol)
+        text = name.to_s
+        return false unless text.match?(/\A[a-z_][A-Za-z0-9_]*\z/)
+        parsed = Prism.parse("#{text} = 1")
+        parsed.errors.empty? &&
+          parsed.value.statements.body.first.is_a?(Prism::LocalVariableWriteNode)
+      end
+
+      # What a contraction is once its block has been read: a source, the
+      # arrays that source names, and which indices are free.  Everything
+      # before this is about recovering those from a block; everything after
+      # is the same whatever recovered them, which is what lets
+      # `contract_terms` reach it with a source it wrote itself.
+      def contract (source, arrays, free_indices,
+                    node: nil, origin: nil, scalars: {}, c_functions: {})
+        result = allocate_result(source, node, arrays, scalars, free_indices)
         arrays = arrays.merge(RESULT => result) if result
 
         kernel = compile(source,
@@ -346,8 +633,15 @@ class CArray
                          masked: arrays.each_value.any? { |array| array.has_mask? },
                          contract: true,
                          result: RESULT,
+                         # A contraction sums an index; which order it sums it
+                         # in is not something the caller wrote, so splitting
+                         # the sum into partial ones does not change what the
+                         # contraction means.  Same licence `jit_for` takes.
+                         reassociate: JIT.reassociate,
+                         free_indices: free_indices,
                          cell_names: cell_names(arrays))
 
+        refuse_aliased_result(kernel, arrays)
         extents = contraction_extents(kernel, arrays)
         if kernel.masked
           kernel.written_arrays.each do |name|
@@ -359,12 +653,41 @@ class CArray
         result || kernel
       end
 
+      # An array a contraction writes and also reads is a recurrence, which
+      # the analyzer refuses -- but it compares the names a block gave them,
+      # and two names may be one array.  `x = a` is one, and so is a view of
+      # something being read: cells reached before the write see the old value
+      # and cells reached after see the new one, so the answer depends on the
+      # order and is not the contraction that was asked for.
+      #
+      # This is where the arrays themselves are known, so it is where the
+      # question can be asked of them rather than of their names.  Views are
+      # followed to what they are views of, since that is the memory two names
+      # would share.
+      def refuse_aliased_result (kernel, arrays)
+        kernel.written_arrays.each do |written|
+          target = root_of(arrays.fetch(written))
+          arrays.each do |name, array|
+            next if name == written
+            next unless root_of(array).equal?(target)
+            raise Unsupported,
+                  "`#{written}` and `#{name}` are the same array, which this " \
+                  "both writes and reads; that is a recurrence rather than a " \
+                  "contraction, and is written with jit_for"
+          end
+        end
+      end
+
+      def root_of (array)
+        array.respond_to?(:root_array) ? array.root_array : array
+      end
+
       # A contraction with nothing to assign into needs its result sized and
       # typed before there is a kernel to ask, so the block is analyzed once
       # without being compiled.  Returns nil when the block assigns into an
       # array of its own.
-      def allocate_result (source, node, arrays, scalars)
-        probe = probe_contraction(source, node, arrays, scalars)
+      def allocate_result (source, node, arrays, scalars, free_indices = nil)
+        probe = probe_contraction(source, node, arrays, scalars, free_indices)
         return nil unless probe
         free, index_axes, type = probe
         shape = free.map do |index|
@@ -385,13 +708,17 @@ class CArray
       end
 
       # @private
-      def probe_contraction (source, node, arrays, scalars)
+      def probe_contraction (source, node, arrays, scalars, free_indices = nil)
         key = [source, arrays.transform_values(&:data_type_name),
                scalars.transform_values { |value| TypeAssignment.scalar_type(value) },
+               # The result's axes are named at the call site rather than in
+               # the block, so the same source under another naming is another
+               # kernel -- and another probe.
+               free_indices,
                cell_names(arrays)]
         cached = probe_cache[key]
         return cached unless cached.nil?
-        probe_cache[key] = build_probe(source, node, arrays, scalars)
+        probe_cache[key] = build_probe(source, node, arrays, scalars, free_indices)
       end
 
       # @private
@@ -400,10 +727,10 @@ class CArray
       end
 
       # @private
-      def build_probe (source, node, arrays, scalars)
+      def build_probe (source, node, arrays, scalars, free_indices = nil)
         storage_types = arrays.transform_values(&:data_type_name)
         analyzer = Analyzer.new(source, node: node, array_names: arrays.keys,
-                                contract: :probe,
+                                contract: :probe, free_indices: free_indices,
                                 cell_names: cell_names(arrays))
         return false if analyzer.body.statements.last.is_a?(ElementWrite)
 
@@ -417,9 +744,13 @@ class CArray
             end
           end
         end
+        # The type the value was computed in decides the array it is collected
+        # into, and that table lives with the computation types rather than
+        # here -- a copy kept at the allocation site answers after a type is
+        # added, and answers wrongly.  This one did: float32 and cmplx64 fell
+        # through its default and a float contraction came back int64.
         [analyzer.probe_free_names, axes,
-         { :double => :float64, :complex => :cmplx128 }
-           .fetch(summand.type, :int64)]
+         TypeAssignment.result_storage_type(summand.type)]
       end
 
       # Each index's extent comes from the axes it addresses.  Where it
@@ -621,6 +952,26 @@ class CArray
                 "#{shape.inspect}, so there is no cell where the window is " \
                 "inside the array"
         end
+        # Writing into an array a window reads is not a pass over the array:
+        # a cell written here is a neighbour a later cell reads, so what comes
+        # back depends on the order the cells were reached in.  The question
+        # is about that window alone -- a window that reaches nowhere reads
+        # only the cell the loop is on, and may be written in place however
+        # far the other windows in the same kernel reach.
+        if into
+          aliased = given.find { |name, array|
+            next false unless root_of(array).equal?(root_of(into))
+            kernel.window_reaches[name].any? { |low, high|
+              !low.zero? || !high.zero?
+            }
+          }
+          if aliased
+            raise Unsupported,
+                  "`into:` is the array `#{aliased.first}` reaches its window " \
+                  "into, and a cell written there is one a later cell reads; " \
+                  "a stencil writes into an array of its own"
+          end
+        end
         result.mask = 0 if (kernel.masked || border == :mask) && !result.has_mask?
         mark_frame(result, bounds, shape) if border == :mask
         operands = given.merge(MAP_RESULT => result)
@@ -896,7 +1247,8 @@ class CArray
       def compile (source, node: nil, origin: nil, array_names:, storage_types:,
                    scalar_values:, c_functions: {}, masked: false, rank: nil,
                    steps: nil, contract: false, result: nil, map: false,
-                   reassociate: false, cell_names: [], windows: [], border: nil)
+                   reassociate: false, cell_names: [], windows: [], border: nil,
+                   free_indices: nil)
         # A kernel that mentions UNDEF is a masked one whatever its arrays
         # carry, and deciding that here means no caller has to remember it.
         masked ||= mentions_undef(source, node)
@@ -909,6 +1261,10 @@ class CArray
                # symbol, which stands for its body -- see `CFunction#kernel_key`.
                c_functions.transform_values(&:kernel_key),
                masked, rank, steps, contract, result, map,
+               # A contraction whose free indices were named is not the kernel
+               # the same source is without them, nor with them in another
+               # order: the naming decides what is summed and what comes out.
+               free_indices,
                # Which names are read at their one cell rather than walked:
                # the same source over a CScalar is a different kernel from
                # the same source over a one-cell CArray.
@@ -929,7 +1285,7 @@ class CArray
         registry[key] = build(source, node, array_names, storage_types,
                               scalar_values, c_functions, masked, rank, steps,
                               contract, result, origin, map, reassociate,
-                              cell_names, windows, border)
+                              cell_names, windows, border, free_indices)
       end
 
       # A kernel that mentions UNDEF is a masked kernel whatever its arrays
@@ -1004,6 +1360,27 @@ class CArray
         Compiler.cache_root
       end
 
+      # Puts this application's kernels somewhere of its own, rather than in
+      # the cache shared under the home directory.  Say it before the first
+      # kernel is compiled -- at the top of the program, beside the other
+      # requires:
+      #
+      #     CArray::JIT.cache_root = File.expand_path("../.jit-cache", __dir__)
+      #
+      # Kernels already loaded keep working and what is already on disk stays
+      # where it is; this says where the next one is looked for and written.
+      # `CARRAY_JIT_CACHE` and `CARRAY_JIT_NO_CACHE` still come first, so
+      # whoever runs the program can put the cache somewhere writable or do
+      # without one. A directory inside a project wants to be ignored by the
+      # version control it sits in.
+      #
+      # @param path [String, nil] the directory, relative to where the
+      #   program starts; `nil` restores the default.
+      # @return [void]
+      def cache_root= (path)
+        Compiler.cache_root = path
+      end
+
       # @return [Array<String>] the environment directories no longer in use --
       #   another version, or another architecture.
       def stale_cache_environments
@@ -1036,11 +1413,11 @@ class CArray
       def build (source, node, array_names, storage_types, scalar_values, c_functions,
                  masked, rank = nil, steps = nil, contract = false, result = nil,
                  origin = nil, map = false, reassociate = false,
-                 cell_names = [], windows = [], border = nil)
+                 cell_names = [], windows = [], border = nil, free_indices = nil)
         analyzer = Analyzer.new(source, node: node, array_names: array_names,
                                 c_functions: c_functions,
                                 rank: rank, steps: steps, contract: contract,
-                                result: result, map: map,
+                                result: result, map: map, free_indices: free_indices,
                                 cell_names: cell_names, windows: windows)
         assignment = TypeAssignment.new(analyzer.body, storage_types,
                                         scalar_values, c_functions)

metadata

@@ -1,7 +1,7 @@
 --- !ruby/object:Gem::Specification
 name: carray-jit
 version: !ruby/object:Gem::Version
-  version: 0.1.0
+  version: 0.1.2
 platform: ruby
 authors:
 - himotoyoshi
@@ -133,7 +133,7 @@ required_rubygems_version: !ruby/object:Gem::Requirement
     - !ruby/object:Gem::Version
       version: '0'
 requirements: []
-rubygems_version: 4.0.17
+rubygems_version: 4.0.8
 specification_version: 4
 summary: JIT compilation of CArray kernels written in Ruby
 test_files: []