Class: Cumo::CUDA::ReductionKernel

Inherits:
Object
  • Object
show all
Includes:
UserKernel
Defined in:
lib/cumo/cuda/reduction_kernel.rb

Overview

A kernel that maps every element, reduces the mapped values along the given axes and maps the result, the way CuPy's ReductionKernel does.

l2norm = Cumo::CUDA::ReductionKernel.new(
"T x", "T y", "x * x", "a + b", "y = sqrt(a)", "0", "l2norm")
l2norm.call(x, axis: 1)

The map expression sees the inputs by name, the reduce expression sees the two values a and b, and the post expression sees the reduced value a and writes the output. The identity starts every reduction.

Constant Summary collapse

BLOCK = 
512
MAX_GRID = 
65_535
WANT_BLOCKS =

A second pass is worth it once this many blocks would otherwise be idle, and each chunk keeps this many elements per thread.

256
MIN_PER_THREAD = 
16
DTYPE_OF_CTYPE = 
CTYPE.invert.freeze

Constants included from UserKernel

UserKernel::CTYPE, UserKernel::INT_RANGE, UserKernel::PACK, UserKernel::RESERVED, UserKernel::TYPES

Instance Attribute Summary collapse

Instance Method Summary collapse

Constructor Details

#initialize(in_params, out_params, map_expr, reduce_expr, post_map_expr, identity, name, preamble: "", reduce_type: nil) ⇒ ReductionKernel

Returns a new instance of ReductionKernel.

Raises:

  • (ArgumentError) 


29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
 
# File 'lib/cumo/cuda/reduction_kernel.rb', line 29

def initialize(in_params, out_params, map_expr, reduce_expr, post_map_expr, identity, name,
               preamble: "", reduce_type: nil)
  check_name(name)
  @in_params = parse_params(in_params)
  @out_params = parse_params(out_params)
  if (raw = (@in_params + @out_params).find(&:raw))
    raise ArgumentError, "#{raw.name} is raw, and a reduction indexes every argument itself"
  end
  raise ArgumentError, "a reduction needs an output" if @out_params.empty?
  if (taken = (@in_params + @out_params).find { |p| %w[a b].include?(p.name) })
    raise ArgumentError, "#{taken.name} is what the reduce expression calls its operands"
  end
  @map_expr = map_expr
  @reduce_expr = reduce_expr
  @post_map_expr = post_map_expr
  @identity = identity_source(identity)
  @name = name
  @preamble = preamble
  @reduce_type = check_reduce_type(reduce_type)
  @functions = {}
  @finals = {}
end

Instance Attribute Details

#name ⇒ Object (readonly)

Returns the value of attribute name.



27
28
29
 
# File 'lib/cumo/cuda/reduction_kernel.rb', line 27

def name
  @name
end

Instance Method Details

#call(*args, axis: nil, keepdims: false) ⇒ Object

Reduces along axis:, every axis when it is nil, and answers the output or the outputs as an Array. keepdims: keeps the reduced axes as length one. Outputs may be given after the inputs.

Raises:

  • (ArgumentError) 


55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
 
# File 'lib/cumo/cuda/reduction_kernel.rb', line 55

def call(*args, axis: nil, keepdims: false)
  ins, outs = split_args(args, @in_params, @out_params)
  types = resolve_types(@in_params, ins, @out_params, outs)
  arrays = ins.select { |a| a.is_a?(Cumo::NArray) }
  raise ArgumentError, "a reduction needs an NArray to reduce" if arrays.empty?
  in_shape = broadcast_shapes(arrays.map(&:shape))
  axes = normalize_axes(axis, in_shape.size)
  kept = (0...in_shape.size).to_a - axes
  out_shape = kept.map { |d| in_shape[d] }
  full_shape = in_shape.each_with_index.map { |s, d| axes.include?(d) ? 1 : s }
  want = keepdims ? full_shape : out_shape

  outs = @out_params.each_with_index.map { |p, k| outs[k] || types[p.type].new(*want) }
  check_outputs(outs, @out_params, want)

  out_size = out_shape.inject(1, :*)
  red_size = axes.map { |d| in_shape[d] }.inject(1, :*)
  launch(ins, outs, types, in_shape, axes + kept, out_size, red_size) if out_size > 0
  @out_params.size == 1 ? outs[0] : outs
end