carray-jit 0.1.2 → 0.1.3

data/lib/carray/jit/kernel.rb

@@ -19,7 +19,13 @@ class CArray
       #   @return [Boolean] whether this call invoked the compiler, rather
       #     than reusing a cached object.
       attr_reader :source, :c_source, :arrays, :storage_types, :reals,
-                  :integers, :complexes, :rank, :index_names, :written_arrays,
+                  :integers, :complexes, :unsigned_integers,
+                  :rank, :index_names, :written_arrays,
+                  # The arrays handed to a C function whole, by address,
+                  # rather than walked a cell at a time.  A caller that lines
+                  # operands up has to leave these alone: their shape is the
+                  # C function's business, not the expression's.
+                  :address_arrays,
                   :compiled, :masked, :directions, :contracted_names,
                   :index_axes,
                   # How far a stencil's windows reach on each axis, as
@@ -38,6 +44,7 @@ class CArray
         @reals = generator.reals
         @integers = generator.integers
         @complexes = generator.complexes
+        @unsigned_integers = generator.unsigned_integers
         # Per array, the axes read at an index only the running kernel knows,
         # and the extents it checks them against.
         @extent_slots = generator.extent_slots
@@ -48,6 +55,9 @@ class CArray
         @written_arrays = analyzer.written_arrays
         @array_ranks = analyzer.array_ranks
         @inner_ranges = analyzer.inner_ranges
+        # An index the block wrote twice has two identifiers here; the
+        # messages speak the name it was written with.
+        @index_sources = analyzer.index_sources
         @axis_uses = @arrays.to_h { |array|
           [array, (0...array_rank(array)).map { |axis| analyzer.axis_use(array, axis) }]
         }
@@ -61,6 +71,17 @@ class CArray
           end
         end
 
+        # Which captured scalars decide where the kernel reaches, rather than
+        # what it computes.  A pinned subscript, an offset and an inner loop's
+        # range are worked out here, in Ruby, before the first cell; every
+        # other captured number is packed into a buffer and read by the C.  So
+        # these are the ones a prepared call has to be keyed on -- and a
+        # coefficient or a clock that changes every step is not one of them,
+        # which is what keeps a time-stepping caller on the prepared path.
+        @shape_scalar_names = shape_scalar_names
+        # Calls prepared earlier, most recently used first.
+        @plans = []
+
         @masked = generator.masked
         @window_reach = analyzer.window_reach
         @window_reaches = analyzer.window_reaches
@@ -68,6 +89,15 @@ class CArray
         # writes into the error slot.  The message does not travel: C has
         # nothing to carry it in, and it was known when this was compiled.
         @raise_messages = generator.raise_messages
+        # The arrays this kernel allocates at its entry, as
+        # [name, storage, shape] -- so that a shape it could not use, or an
+        # allocation that failed, can be named with the lengths it came to on
+        # this call.  The shapes are worked out here as they are there: from
+        # the same expressions over the same captured integers.
+        @heap_arrays = analyzer.local_array_declarations.filter_map { |entry|
+          _scope, name, storage, shape, heap = entry
+          [name, storage, shape] if heap
+        }
         # Kept so the sweep entry point can be built from it if one is ever
         # asked for.  Compiling it eagerly would pay for a road most kernels
         # never take.
@@ -109,16 +139,15 @@ class CArray
       def call (array_values, scalar_values, bounds, c_function_values = {},
                 border: false)
         arrays = @arrays.map { |name| array_values.fetch(name) }
-        ranges = index_ranges(bounds, scalar_values)
-        verify(arrays, bounds, ranges, scalar_values, reach: !border)
+        plan = prepared_call(arrays, array_values, bounds, scalar_values, border)
 
-        writable = @arrays.map { |name| @written_arrays.include?(name) }
+        writable = plan[:writable]
         error = [0].pack("l")
-        packed_bounds = bounds.flatten.pack("q*")
+        packed_bounds = plan[:bounds]
         reals = packed_reals(scalar_values)
-        integers = (@integers.map { |name| Integer(scalar_values.fetch(name)) } +
-                    @extent_slots.map { |name, axis| array_values.fetch(name).dim[axis] })
-                   .pack("q*")
+        # The captured integers change from call to call and are packed here;
+        # the extents behind them are the arrays' own and travel in the plan.
+        integers = packed_scalar_integers(scalar_values) + plan[:extents]
         functions = @c_function_names.map { |name|
           c_function_values.fetch(name).pointer.to_i
         }.pack("Q*")
@@ -136,14 +165,12 @@ class CArray
           }
         }.pack("Q*")
 
-        box = region_box(ranges, scalar_values, array_values)
-        Access.open(arrays, writable, box[0], box[1]) do |bases|
-          pointers = bases.map { |basis| basis[:pointer] }.pack("Q*")
-          strides = bases.flat_map { |basis| basis[:strides] }.pack("q*")
-          mask_pointers = bases.map { |basis| basis[:mask_pointer] || 0 }.pack("Q*")
-          mask_strides = bases.flat_map { |basis|
-            basis[:mask_strides] || Array.new(@rank, 0)
-          }.pack("q*")
+        box = plan[:box]
+        # Opened for a kernel rather than for a reader: the four buffers come
+        # back packed, instead of a hash and two arrays per operand that this
+        # would pack back into bytes and throw away.
+        Access.open(arrays, writable, box[0], box[1], true) do
+                   |pointers, strides, mask_pointers, mask_strides|
           entry = border ? border_function : @function
           entry.call(buffer(pointers), buffer(strides), buffer(packed_bounds),
                          buffer(reals), buffer(integers), buffer(functions),
@@ -157,9 +184,17 @@ class CArray
           array[] = buffer
         end
 
-        report(error)
+        report(error, scalar_values)
       end
 
+      # How many prepared calls one kernel keeps.  A stencil asks under its
+      # interior's bounds and one pair per axis of frame, so a two-dimensional
+      # one needs five before it repeats; a caller alternating between two
+      # model states needs two of whatever it was already using.  Past this
+      # the oldest goes and the call prepares itself again, which is what
+      # every call did before.
+      PLAN_LIMIT = 8
+
       # @private
       SLAB_NAME = "carray_jit_slab"
 
@@ -178,10 +213,7 @@ class CArray
 
         error = [0].pack("l")
         reals = packed_reals(scalar_values)
-        integers = (@integers.map { |name| Integer(scalar_values.fetch(name)) } +
-                    @extent_slots.map { |name, axis|
-                      array_values.fetch(name).dim[axis]
-                    }).pack("q*")
+        integers = packed_integers(scalar_values, array_values)
         functions = @c_function_names.map { |name|
           c_function_values.fetch(name).pointer.to_i
         }.pack("Q*")
@@ -211,7 +243,7 @@ class CArray
           array[] = item
         end
 
-        report(error)
+        report(error, scalar_values)
       end
 
       # The wrapper CArray's sweep calls, for the same reason #c_source is
@@ -230,16 +262,138 @@ class CArray
 
       private
 
+      # The parts of a call that the operands' shapes, the bounds and the
+      # kernel's own scalars decide, rather than the data or the views.
+      #
+      # Everything here answers the same for two calls whose key agrees, so it
+      # is worked out once and kept: the bounds check, the box the xfer tier
+      # transfers, the extents the kernel is handed, and which operands it
+      # writes.  What is deliberately *not* here is anything read off a view:
+      # two arrays of one shape and one type can be laid out differently, and
+      # a stride read from the one is not a stride into the other.  Those come
+      # from Access.open on every call, as they always did.
+      #
+      # Nor is the refusal of a masked array handed whole to a C function
+      # (#address_buffer): a mask is runtime state that leaves the shape and
+      # the type alone, so no key here would notice one appearing.  It is
+      # asked on every call, which costs a predicate and keeps the refusal.
+      def prepared_call (arrays, array_values, bounds, scalars, border)
+        key = plan_key(arrays, bounds, scalars, border)
+        if key
+          found = @plans.assoc(key)
+          if found
+            unless @plans.first.equal?(found)
+              @plans.unshift(@plans.delete(found))
+            end
+            return found[1]
+          end
+        end
+        ranges = index_ranges(bounds, scalars)
+        verify(arrays, bounds, ranges, scalars, reach: !border)
+        plan = {
+          :writable => @arrays.map { |name| @written_arrays.include?(name) },
+          :bounds   => bounds.flatten.pack("q*"),
+          :extents  => packed_extents(array_values),
+          :box      => region_box(ranges, scalars, array_values),
+        }
+        freeze_throughout(plan)
+        if key
+          @plans.unshift([key, plan])
+          @plans.pop if @plans.size > PLAN_LIMIT
+        end
+        plan
+      end
+
+      # What two calls have to agree on for one to be prepared like the other.
+      #
+      # Shapes and types, because that is all `verify` and `region_box` and
+      # the extents read of an array; the bounds and the border, because they
+      # are what the check is against; and the scalars that decide where the
+      # kernel reaches, but not the ones it merely computes with.
+      #
+      # No identity: two arrays of one shape and type prepare the same way, so
+      # keying on which object it was would only make a caller that rebuilds
+      # its fields every step miss every time -- and would have this kernel,
+      # which lives as long as the process, hold their memory.
+      #
+      # nil where an operand is not a CArray at all, so that the call falls
+      # through to `verify` and is refused by name rather than by whatever
+      # the next method call happens to raise.
+      def plan_key (arrays, bounds, scalars, border)
+        shapes = []
+        arrays.each do |array|
+          return nil unless array.is_a?(CArray)
+          shapes << array.data_type << array.dim
+        end
+        [shapes, bounds, border,
+         @shape_scalar_names.map { |name| scalars[name] }]
+      end
+
+      # A kernel is shared, and Fiddle lets go of the GVL for the call, so two
+      # threads can be reading one plan while a third prepares another.  That
+      # is sound only while nothing in a plan is written after it is published
+      # -- and a box is arrays inside an array, so freezing the top of it
+      # would leave the parts a later edit could still reach.
+      def freeze_throughout (value)
+        case value
+        when Hash  then value.each_value { |item| freeze_throughout(item) }
+        when Array then value.each { |item| freeze_throughout(item) }
+        end
+        value.freeze
+      end
+
+      # The captured scalars a pinned subscript, an offset or an inner loop's
+      # range is worked out from.
+      def shape_scalar_names
+        names = []
+        @axis_uses.each_value do |uses|
+          uses.each do |walkers, pinned|
+            pinned.each { |node| collect_captures(node, names) }
+            walkers.each { |walker| walker[1].each { |node|
+              collect_captures(node, names)
+            } }
+          end
+        end
+        @inner_ranges.each_value do |from, to, _step|
+          collect_captures(from, names)
+          collect_captures(to, names)
+        end
+        names.uniq.freeze
+      end
+
+      def collect_captures (node, names)
+        return names unless node.is_a?(Node)
+        names << node.name if node.is_a?(CaptureRead)
+        node.children.each { |child| collect_captures(child, names) }
+        names
+      end
+
       # What the cell that stopped said, raised now that Ruby has control
       # back.  C cannot raise, so a failure is a code in a slot and a loop
       # that stops; this is where it becomes the exception the same body run
       # in Ruby would have raised.
-      def report (error)
+      def report (error, scalars = {})
         code = error.unpack1("l")
+        if (failure = CGenerator::FIXED_FAILURES[code])
+          raise failure[0], failure[1]
+        end
         case code
         when 0 then nil
+        when CGenerator::SHAPE_CODE then raise_a_bad_shape(scalars)
+        when CGenerator::MEMORY_CODE then raise_out_of_memory(scalars)
         when 1 then raise ZeroDivisionError, "divided by 0"
         when 2 then raise IndexError, "index out of range"
+        when 3 then raise ArgumentError,
+                            "min argument must be less than or equal to " \
+                            "max argument"
+        # Ruby's own wording names the value it could not compare, and the
+        # slot carries a code rather than a number, so the reason is named
+        # instead.  The class and the failure are Ruby's.
+        when 4 then raise ArgumentError,
+                            "comparison with a NaN failed, so `clamp` has " \
+                            "no answer"
+        when 5 then raise Math::DomainError,
+                            "Numerical argument is out of domain - gamma"
         else
           message = @raise_messages[code]
           # A code with no message behind it is this compiler's bug, not the
@@ -251,6 +405,53 @@ class CArray
         end
       end
 
+      # A length the kernel worked out that is no length: zero cells, a
+      # negative count, or so many that the bytes would not be a `size_t`.
+      # The kernel reports which of its arrays it was by reporting at all --
+      # it stops before the first cell -- and the lengths are worked out
+      # again here, where there is a message to put them in.
+      def raise_a_bad_shape (scalars)
+        offending = @heap_arrays.filter_map { |name, storage, shape|
+          lengths = shape_lengths(shape, scalars)
+          next unless lengths&.any? { |length| length < 1 }
+          "`#{name}` (#{storage}) asks for " \
+          "#{lengths.join(' x ')} #{lengths.size == 1 ? 'cells' : 'cells per axis'}"
+        }
+        raise ArgumentError,
+              "a local array holds at least one cell: " \
+              "#{offending.join('; ')}" unless offending.empty?
+        raise ArgumentError,
+              "a local array of this kernel asks for more cells than its " \
+              "bytes could be counted in"
+      end
+
+      def raise_out_of_memory (scalars)
+        asked = @heap_arrays.map { |name, storage, shape|
+          lengths = shape_lengths(shape, scalars)
+          cells = lengths&.inject(1, :*)
+          bytes = cells && cells * LOCAL_ARRAY_STORAGE_BYTES.fetch(storage, 1)
+          "`#{name}` (#{storage}, #{lengths ? lengths.join(' x ') : 'shape'}" +
+            (bytes ? ", #{bytes} bytes" : "") + ")"
+        }
+        raise NoMemoryError,
+              "a kernel could not allocate the local arrays it works in: " \
+              "#{asked.join(', ')}"
+      end
+
+      # @private
+      LOCAL_ARRAY_STORAGE_BYTES = Analyzer::LOCAL_ARRAY_STORAGE_BYTES
+
+      # The shape as numbers, or nil where one of its axes cannot be worked
+      # out here -- which is not a failure worth a second exception, the one
+      # being raised having happened already.
+      def shape_lengths (shape, scalars)
+        shape.map { |extent|
+          extent.is_a?(Integer) ? extent : evaluate(extent.node, scalars)
+        }
+      rescue StandardError, Unsupported
+        nil
+      end
+
       def slab_pointer
         @slab_pointer ||= @handle[SLAB_NAME]
       end
@@ -280,6 +481,35 @@ class CArray
          }).pack("d*")
       end
 
+      # The `integers` slots the kernel reads: the signed captures, then the
+      # unsigned ones, then the extents an index is checked against.  A uint64
+      # is packed as the bits it is -- `pack("q")` would take 2**63 modulo the
+      # width and say nothing -- and the kernel casts the slot back, so the
+      # two directives are what makes one buffer carry both widths.
+      #
+      # One method because both loops pack it, and the same buffer packed in
+      # two places is what `packed_reals` is one method about.
+      def packed_integers (scalar_values, array_values)
+        packed_scalar_integers(scalar_values) + packed_extents(array_values)
+      end
+
+      # The captured half, which is different every time a caller steps a
+      # counter, and the arrays' half, which is the shape and travels in the
+      # plan.  Two methods because `#call` packs the first and takes the
+      # second from what it prepared earlier; `#sweep` still wants both.
+      def packed_scalar_integers (scalar_values)
+        @integers.map { |name| Integer(scalar_values.fetch(name)) }.pack("q*") +
+          @unsigned_integers.map { |name|
+            Integer(scalar_values.fetch(name))
+          }.pack("Q*")
+      end
+
+      def packed_extents (array_values)
+        @extent_slots.map { |name, axis|
+          array_values.fetch(name).dim[axis]
+        }.pack("q*")
+      end
+
       def address_buffer (name, array)
         @address_parameters.fetch(name).each do |parameter|
           wanted = CDeclaration::DATA_TYPES.fetch(parameter.element.fiddle)
@@ -318,7 +548,7 @@ class CArray
         packed.empty? ? "\0" * 8 : packed
       end
 
-      # The region tier transfers a box rather than the whole view.  The box
+      # The xfer tier transfers a box rather than the whole view.  The box
       # is the loop range grown by how far the kernel reaches from the cell it
       # is on -- per array, because two arrays in one kernel need not be read
       # at the same offsets, and per axis, because they need not be read at
@@ -376,13 +606,117 @@ class CArray
           ranges[name] = covered_span(bounds[axis])
         end
         flat = bounds.flatten
-        @inner_ranges.each do |name, (from, to)|
-          ranges[name] = [evaluate(from, scalar_values, flat),
-                          evaluate(to, scalar_values, flat)]
+        # In source order, outermost first, so an index a range is written
+        # over has its own range settled before it is read.
+        @inner_ranges.each do |name, (from, to, step)|
+          step ||= 1
+          first = span_of(from, scalar_values, flat, ranges, name)
+          last = span_of(to, scalar_values, flat, ranges, name)
+          # The widest pass the range could describe: counting up, that is the
+          # earliest start against the latest end, and counting down it is the
+          # other way round.
+          start, limit = step.positive? ? [first.first, last.last]
+                                        : [first.last, last.first]
+          # A step that skips cells lands on the ones a pass started on, and
+          # where the start moves so do they: `(p...8).step(3)` visits 0, 3, 6
+          # from one start and 1, 4, 7 from the next.  Stepped out from the
+          # earliest start alone, the reach stopped at 6 and a pass wrote at
+          # 7.  So with a start that moves, the span is every cell between
+          # the widest start and the widest end, which every pass lies within.
+          step = step <=> 0 if first.first != first.last
+          ranges[name] = covered_span([start, limit, step])
         end
         ranges
       end
 
+      # A range written over another index is read at its widest (see
+      # #span_of), so the number in the message above can be one no pass
+      # actually starts or ends at.  Saying which index it was written over is
+      # what tells a reader whether to narrow the range or the array's use --
+      # without it the message names a bound the loop never reaches.
+      def widest_reading (index, spelled)
+        over = range_indices(index)
+        return "" if over.empty?
+        ". The range on `#{spelled}` is written over " \
+        "#{over.map { |name| "`#{name}`" }.join(' and ')}, and a range over " \
+        "another index is read here at its widest -- as far as any value of " \
+        "#{over.size == 1 ? 'it' : 'them'} could take it, which may be " \
+        "further than it goes on any one pass"
+      end
+
+      # The indices an inner loop's range mentions, by the name the block
+      # wrote for each.
+      def range_indices (index)
+        range = @inner_ranges[index]
+        return [] unless range
+        from, to, = range
+        [from, to].flat_map { |node| mentioned_indices(node) }.uniq
+          .map { |name| @index_sources.fetch(name, name) }
+      end
+
+      def mentioned_indices (node)
+        return [] unless node.is_a?(Node)
+        return [node.name] if node.is_a?(IndexVariable)
+        node.children.flat_map { |child| mentioned_indices(child) }
+      end
+
+      # What an index range covers, as [lowest, highest] inclusive.
+      #
+      # A range written over another index -- `(p+1...3)`, the shape a
+      # triangular loop takes -- has no single pair of numbers to be: `k`
+      # starts somewhere else for every `p`.  What the two callers of this
+      # want is where the index can reach, not what it is on a given pass:
+      # `verify_bounds` asks whether a subscript can leave its array, and
+      # `region_box` asks what a stencil touches.  So the answer is an
+      # interval, and a range over another index is read at its widest.
+      #
+      # Widest is the safe direction.  It covers every pass the loop can take
+      # and then some, so a subscript this says is inside really is; what it
+      # costs is that a reach nothing actually makes can still be refused.
+      #
+      # A range that mentions no index gives an interval of one number, which
+      # is what `evaluate` gave before, so those kernels are unchanged.
+      def span_of (node, scalars, flat_bounds, ranges, index)
+        case node
+        when IntegerLiteral then [node.value, node.value]
+        when CaptureRead    then value = Integer(scalars.fetch(node.name))
+                                 [value, value]
+        when BoundsValue    then value = flat_bounds.fetch(node.slot)
+                                 [value, value]
+        when IndexVariable
+          # Held half-open, as the loops are; the last index it visits is one
+          # below the end.
+          low, high = ranges.fetch(node.name) do
+            raise Unsupported,
+                  "the range on `#{index}` is written over `#{node.name}`, " \
+                  "which is not an index around it"
+          end
+          [low, high - 1]
+        when UnaryMinus
+          low, high = span_of(node.operand, scalars, flat_bounds, ranges, index)
+          [-high, -low]
+        when BinaryOperation
+          left = span_of(node.left, scalars, flat_bounds, ranges, index)
+          right = span_of(node.right, scalars, flat_bounds, ranges, index)
+          case node.operator
+          when :+ then [left.first + right.first, left.last + right.last]
+          when :- then [left.first - right.last, left.last - right.first]
+          when :*
+            # Four corners, because either interval may straddle zero.
+            corners = [left.first * right.first, left.first * right.last,
+                       left.last * right.first, left.last * right.last]
+            [corners.min, corners.max]
+          else
+            raise Unsupported,
+                  "an inner loop's range is built from `+`, `-` and `*` only"
+          end
+        else
+          raise Unsupported,
+                "an inner loop's range is an integer expression over " \
+                "literals, captured scalars and the indices around it"
+        end
+      end
+
       # The half-open span an axis actually touches.  A step may skip cells,
       # but every cell it lands on has to exist, so the span runs from the
       # first index to the last one visited.
@@ -463,19 +797,22 @@ class CArray
           walkers.each do |index, offsets|
             low, high = ranges.fetch(index)
             next if low >= high
+            spelled = @index_sources.fetch(index, index)
             minimum, maximum = offset_span(offsets, scalars)
             if low + minimum < 0
               raise Unsupported,
                     "`#{name}` is indexed at " \
-                    "#{offset_text(name, index, minimum)}, so " \
-                    "the range on `#{index}` cannot start at #{low}"
+                    "#{offset_text(name, spelled, minimum)}, so " \
+                    "the range on `#{spelled}` cannot start at #{low}" +
+                    widest_reading(index, spelled)
             end
             if high - 1 + maximum > extent - 1
               raise Unsupported,
                     "`#{name}` is indexed at " \
-                    "#{offset_text(name, index, maximum)}, so " \
-                    "the range on `#{index}` cannot end at #{high} " \
-                    "for an extent of #{extent}"
+                    "#{offset_text(name, spelled, maximum)}, so " \
+                    "the range on `#{spelled}` cannot end at #{high} " \
+                    "for an extent of #{extent}" +
+                    widest_reading(index, spelled)
             end
           end
         end