cumo 0.6.0 → 0.7.0

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: 0d6b3e0decf69d07199a7930dc6a23d1a05a90ed044cf007203c78c28a8e9ea1
-  data.tar.gz: dec0ca1f195b302919d30085fcc139911299ba075193b67cab9965907210092f
+  metadata.gz: 97c7f575209616cb9fa484e5dd4096bac11170aef18e99f1f6f733e599daa9af
+  data.tar.gz: 5efabd7aface13fac85dcef81fef54dd9bb22d3748666177ff221741a8419486
 SHA512:
-  metadata.gz: e4f8d6204fca59288efad1ab5dec228d474de02ee3465f89afe7fc7440b456b0d46b78c92fb1e48e92063a2b3c385dedb8955aed3ddfa4411c13beb2f32a5417
-  data.tar.gz: a397b427534b829b13f8bd4a4d6f8616d918bde834dc5c69a5611671ec942796d8615bd74d2564a2c3cd6cdaf9e572f4599e0704071f21a3c40394d0236b148d
+  metadata.gz: 314d2b224cb753761618b6bc833af4715d8d1cb2e1fd8de05b0ad65483ec940c31167997abe7cc76e2f056b12b442885731d8e6b4d4b47ec174c19c70c2326e9
+  data.tar.gz: 5a4909d952dbe926b52ccaf04010b3ffcf08617a2af710ff49a70764ff7e626ffddcef04ebcb4ca6b5aab8c7b14170b9b3328768a203dd7961d2fbab5c973002

data/CHANGELOG.md

@@ -1,3 +1,45 @@
+# 0.7.0 (2026/09/13)
+
+Breaking changes:
+
+* `cumo_narray_data_t` no longer carries `owned` and `cumo_narray_t` no longer carries `elmsz`. Both ship in `narray.h`, so an extension compiled against 0.6.0 has to be rebuilt. `NArray#free` also refuses a frozen array, as the other release paths already did (PR #421)
+* `gemm` raises `RangeError` for the sizes cuBLAS takes as `int`, where they wrapped and the call read outside its operands (PR #399)
+* Taking a pointer refuses an `allocate` that leaves the array without a buffer, or hands back one already holding data, and `to_binary` measures the array by its type rather than by the `ELEMENT_BYTE_SIZE` its class names (PR #393, PR #392)
+* `store_binary` refuses a non-contiguous view, and a `Cumo::Bit` view that does not begin and end on a byte. It copies the string it is given, so `NArray#free` answers true after a frozen one where it answered false (PR #390, PR #389)
+* `Cumo::Int8.upcast(Cumo::UInt8)` answers `Cumo::Int16` where it answered `Cumo::Int8`, and `Cumo::Bit.upcast(Cumo::RObject)` answers rather than `nil`, the tables having been built before the classes they name existed (PR #378)
+
+Fixes:
+
+* Fix a view constructor that raises partway leaving `stridx` unreadable for the free hook, and leaking the index array it had just built (PR #407)
+* Fix `Cumo::Bit#each`, `map` and `map_with_index` reading what the block queued before the device had run it, 7 of the 18 dtype and method pairs having answered differently from numo (PR #401)
+* Fix `gemm` reading and writing far outside its operands when a matrix holds more than 2**31 elements, the batch strides reaching cuBLAS as `int` where cuBLAS takes `long long int` (PR #399)
+* Fix `median(nan: true)` answering with a number when a NaN is present, the host quicksort it used leaving NaN unordered (PR #398)
+* Fix `where`, `where2` and a masked `aref` writing past the index array they sized, `allocate` being free to resize it between the count and the walk (PR #397)
+* Fix `Cumo::Bit#each_with_index` writing in front of the counter array on a zero-dimensional array, which ordinary code reaches since indexing a Bit answers a zero-dimensional Bit (PR #396)
+* Fix the segfaults an `allocate` that takes no buffer, reshapes the array, or claims its own element size leaves behind, and `to_binary` handing back a 16KB String from a 32 byte allocation (PR #393, PR #392)
+* Fix `store_binary` writing at the start of the base rather than through the view, and pointing the array at a frozen String's bytes, which a compaction moves out from under it (PR #390, PR #389)
+* Fix storing a Ruby Array of sub-narrays: the rows came back as the first one repeated, a row short in an outer axis copied past its own end, and a row with no elements stored the loop's flag as a value, 8.5us a row down to 1.6 (PR #384, PR #383, PR #382)
+* Fix flattening a `Cumo::Bit` view that walks an axis backwards reading near address zero, which `to_a`, `where`, `where2` and `mask` all reach (PR #380)
+* Fix a length-1 axis costing a view its contiguity, and `reshape!` on one carrying a negative stride walking out of the allocation; of 288 views measured, 92 more answer `contiguous?` and none fewer (PR #379, PR #376)
+* Fix `:new` in front of an index-backed dimension reading and writing outside the array, only a trailing new axis having been recognized as one (PR #375)
+* Fix `min`, `max`, `minmax` and `ptp` answering a finite number for an array of infinities, the nan-aware forms seeding with the largest finite value the dtype holds (PR #372)
+
+Changes:
+
+* Add `Cumo::HFloat`, the 16-bit float dtype the GPU has had since Maxwell: the four operations, the comparisons, the reductions, `Cumo::NMath`, `sort`, `rand`, `dot` through cuBLAS with a float compute type, and the cuDNN methods, which now reach the tensor cores. A 1024x1024 `dot` fell to the mulsum fallback and was 285x slower before it, and `cumsum` stays undefined for half until the scan carries a wider accumulator (issue #107, PR #371, PR #370, PR #368, PR #367, PR #366, PR #365)
+* Copy a transposed 2-d view, and run an elementwise op with a transposed operand, through a 32x32 shared memory tile; on 16M SFloat the copy goes from 177 to 380 GB/s and the op from 247 to 396 GB/s (PR #416, PR #415)
+* Point a derived view at the index arrays it would have copied, one device allocation and one copy per dimension, and let `reverse` and `diagonal` borrow the dimensions they only pass through; on a 2**20 index `transpose` goes from 394us to 0.3us, `reverse(1)` from 327us to 0.2us and `diagonal` from 329us to 0.1us (PR #412, PR #410)
+* Let any device synchronize settle a view's index rather than only the one the index code makes, and record the epoch in every constructor that fills an index; storing sub-narrays into an index-backed view goes from 11.7-18.0us to 5.9-8.9us, and a reversed, expanded or diagonal index view stops waiting once per read (PR #414, PR #388)
+* Wait for a sub-narray index only when the row binds one, rather than whenever any argument uses one; a store into an index-backed destination goes from 6.3us to 1.7us a row (PR #387)
+* Reverse a view's index array with a kernel instead of a host loop that stopped the whole device first; a 2**20 index turns around in 480us rather than 1600us (PR #406)
+* Wait for the device once per row of a host sort instead of once per recursion, a 20000 element row having cost 4003 waits; over 200000 elements `sort(nan: true)` goes from 18.95 to 13.03 ms, `median` from 18.15 to 12.41 and `sort_index` from 20.99 to 14.39 (PR #394)
+* Warn when a view is built holding an index array that nothing owns, which leaks it when that view is the last reference and dangles when it is not (PR #413)
+* Say once what the duplicated templates and kernels had each grown a copy of: the transpose tile walk, the `with_index` block argument and its 0-dimensional clamp, the Bit pointer rebase rule, the min and max rules, the minmax pair, and the byte size a packed bit array takes (PR #422, PR #418, PR #409, PR #408, PR #381, PR #377, PR #374, PR #373)
+* Carry the sub-narray reach per dimension and through the row's own shape rather than through the value channel the store templates read (PR #386, PR #385)
+* Separate the cuDNN scalar, compute and parameter types from the tensor's dtype, so a dtype whose compute type differs is named in `gen/def` rather than in eight templates (PR #369)
+* Delete the `dot` and cuBLAS option code no build compiles, the branch that called a view with no strides contiguous, the second spelling of whether a view's offset reaches a char pointer, and the wait in front of a blocking copy that already queues behind the fill it waited for (PR #417, PR #411, PR #404, PR #403)
+* Count the waits a call makes in the sync probe and probe the block walks, run `Cumo::HFloat` through the structural, math and alt-coverage suites, skip the boundary tests on a machine too small for them, and name the message the gemm row count tests are there for (PR #420, PR #419, PR #405, PR #402, PR #400)
+
 # 0.6.0 (2026/09/06)
 
 Breaking changes:
@@ -5,7 +47,7 @@ Breaking changes:
 * `divmod` and `%` floor the quotient and give the remainder the divisor's sign, as Ruby does, on every signed type: `Cumo::Int32[-7].divmod(3)` answers `[-3, 2]` where it answered `[-2, -1]`, and `-7 % 3` answers `2` rather than `-1`. Dividing a float by zero answers a quotient of `NaN` where it answered `Infinity`. Unsigned types are unchanged (PR #364)
 * A view whose base has since been given a smaller shape raises `RuntimeError` instead of reading past what the base holds (PR #363)
 * `marshal_load` and `initialize` leave an array unallocated when they take a new shape, so it reads as unallocated until something is stored in it, and both refuse a view. `initialize` also refuses a frozen array, as `marshal_load` already did (PR #361, PR #360, PR #355)
-* A range subscript longer than 2**31 answers the length it names instead of an empty view, and one longer than 2**32 no longer answers with the low 32 bits of it (PR #354)
+* A range subscript longer than `2**31` answers the length it names instead of an empty view, and one longer than `2**32` no longer answers with the low 32 bits of it (PR #354)
 * `cov` rejects a 3-dimensional receiver or `y`, 2-dimensional or non-integer `fweights`, and any `ddof` but 0 or 1, where it answered from them; it warns and answers `nan` when the degrees of freedom come out at zero or below (PR #343)
 * `Cumo::Bit#extract` answers with a zero-dimensional `Cumo::Bit`, as every other dtype does, rather than a Ruby Integer read back through a synchronize (PR #341)
 * `batch_norm` and its backward reject an axis that names nothing in `x`, repeats a dimension, or is out of order, where they answered as though a different axis had been given (PR #336)
@@ -265,7 +307,7 @@ Fixes:
 * Backport: improve discrimination between int32/64
 * Backport: missing declaration of id_ge
 * Backport: support casting any object that responds to 'to_a'
-* Backport: fix the range of int32 as -2**31 .. 2**31-1
+* Backport: fix the range of int32 as `-2**31 .. 2**31-1`
 * Backport: if int32_max is Bignum, it should be protected from GC
 * fix: handle compiler flags included in RbConfig::CONFIG properly (#157)
 * Backport: Added more percentile tests

data/README.md

@@ -222,6 +222,126 @@ Cumo::SFloat::Math.atan2(a, 2.0)   #=> Cumo::SFloat
 
 The 0-dimensional form has no effect under Numo, where `[]` returns a Ruby Float.
 
+### Half Precision
+
+`Cumo::HFloat`, also reachable as `Cumo::Float16`, holds IEEE binary16: one sign bit, five of exponent and ten of mantissa.
+It exists for the two things half is good at, moving half as many bytes and reaching the tensor cores, and it promotes exactly as `Cumo::SFloat` does, so an integer array or a Ruby Float mixed into an expression stays half while anything wider takes over.
+
+What it cannot hold is the thing to plan around.
+Integers are exact only to 2048, and the largest finite value is 65504:
+
+```ruby
+Cumo::HFloat[2049.0]        #=> 2048.0
+Cumo::HFloat[50257.0]       #=> 50272.0
+Cumo::HFloat[100000]        #=> Infinity
+```
+
+There is no exception on the way past the top; the value saturates, as it does in every other float type.
+Anything that carries an index rather than a measurement has to be built in a wider type and cast afterwards.
+
+Reductions do not inherit that limit, because they accumulate in single precision and round once at the end.
+A sum of forty thousand ones is forty thousand, not the 2048 a half accumulator would stop at:
+
+```ruby
+Cumo::HFloat.new(40_000).fill(1.0).sum   #=> 40000.0
+```
+
+`sum`, `mean`, `var`, `stddev`, `rms`, `mulsum`, `dot`, `gemm` and `cumsum` all widen this way.
+`prod` and `cumprod` do not: a product leaves half's range long before it loses precision, so widening the accumulator would only hide the overflow.
+
+The widening protects the accumulation, not the answer, which is still stored as half.
+A variance above 65504 therefore saturates even though nothing overflowed while it was being computed, and the standard deviation of the same array is fine because the square root brings it back into range:
+
+```ruby
+a = Cumo::HFloat[2000.0, -2000.0, 1000.0, -1000.0]
+a.var      #=> Infinity
+a.stddev   #=> 1826.0
+```
+
+That distinction is what makes half usable in a transformer's layer normalization, and getting it wrong is the first thing to go wrong there.
+Squaring overflows at 256, since 256 squared is already past the top, and a residual stream with one outlier feature reaches thousands.
+Writing the normalization out as `((x - mean) ** 2).mean` builds those squares as a half array and the answer is `Infinity`, while `x.var(axis: 1)` keeps them in its single-precision accumulator and answers 10208.0 against a single-precision 10209.
+`var` saturates only if the variance itself is out of range, which is a far higher bar than any one deviation being over 256, but it is still a bar.
+The quantity to check against 65504 is the largest variance a layer produces, not the largest activation in it.
+A single outlier of size d among n values contributes only d squared over n to the variance, so a row 768 wide divides it by 768: an activation of 3000 squares to nine million but raises the variance of its row by about twelve thousand.
+An activation that looks safe therefore says nothing about whether `var` overflows, in either direction.
+
+A `dot` whose answer does not fit still saturates, since the result is stored back as half.
+The accumulator itself is single precision, so scaling it down on the way out recovers the value:
+
+```ruby
+a = Cumo::HFloat.ones(1, 1024)
+b = Cumo::HFloat.new(1024, 1).fill(100.0)
+a.dot(b)                   #=> Infinity
+a.gemm(b, alpha: 0.001)    #=> 102.375, the true 102400 scaled down
+```
+
+#### What half is faster at
+
+`gemm` reaches the tensor cores. Square matrices on an RTX 5070 Ti Laptop, median of three runs each:
+
+```
+              HFloat              SFloat              DFloat
+1024x1024     0.043 ms  49.7 TF   0.161 ms  13.3 TF   5.370 ms  0.40 TF
+2048x2048     0.362 ms  47.4 TF   1.360 ms  12.6 TF   43.46 ms  0.40 TF
+4096x4096     3.335 ms  41.2 TF   10.40 ms  13.2 TF   328.1 ms  0.42 TF
+```
+
+An odd number of columns costs half far more than it costs the others, because a row then starts on a two-byte boundary and the vectorized path is gone.
+It is the column count of either operand that matters, not the row count.
+1024x1024 times 1024x1024, with one dimension made odd at a time:
+
+```
+             all even   M odd      K odd      N odd
+HFloat        51.8 TF   49.2 TF    23.3 TF    23.3 TF
+SFloat        11.9 TF      -       11.6 TF    10.8 TF
+```
+
+The run-to-run spread on these is a few per cent and reaches fifteen at the top end, so the M column says the row count does not matter rather than that it costs 5 per cent.
+
+The penalty is on the arithmetic, so it does not reach a matrix-vector product, which is bound by how fast the matrix can be read whatever its shape.
+A 1x768 by 768x50257 gemv takes 0.211 ms with that odd 50257 and 0.205 ms with 50256, a difference inside the noise; the same 768x50257 matrix against 256 rows takes 0.833 ms and 0.481 ms, which is not.
+Pad the inner dimensions of a real matrix product; leave a gemv alone.
+
+`conv` is a different story, and worth reading before reaching for half in a network.
+cuDNN chooses its algorithm from the ones that fit in a scratch buffer, and the half algorithms that use the tensor cores ask for more than the default 8MB ceiling allows.
+Left at the default, a half convolution is no faster than a single-precision one.
+N=32, C=K=64, 56x56, 3x3:
+
+```
+                              HFloat     SFloat
+CUMO_CUDNN_MAX_WORKSPACE_SIZE unset      1.13 ms    1.03 ms
+CUMO_CUDNN_MAX_WORKSPACE_SIZE=268435456  0.56 ms    0.68 ms
+```
+
+Tensor cores also want the channel counts to be multiples of eight, which the first layer of a network never satisfies.
+That layer is still faster in half, but for the other reason:
+
+```
+C=3, K=64, 56x56, 3x3     0.085 ms   0.170 ms
+C=K=64, 56x56, 1x1        0.037 ms   0.120 ms
+```
+
+Neither of those reaches a tensor core; they move half the bytes.
+
+#### Batch normalization takes single-precision parameters
+
+cuDNN derives the descriptor for the batch norm parameters from `x`, and widens it to float when `x` is half.
+`gamma`, `beta`, `running_mean`, `running_var`, `mean` and `inv_std` are therefore `Cumo::SFloat` where `x` is `Cumo::HFloat`, and `batch_norm_backward` answers `gx` in half with `ggamma` and `gbeta` in single:
+
+```ruby
+x = Cumo::HFloat.new(2, 4, 3, 3).seq
+gamma = Cumo::SFloat.ones(4)
+beta = Cumo::SFloat.zeros(4)
+x.batch_norm(gamma, beta, axis: [0, 2, 3])   #=> Cumo::HFloat
+```
+
+Passing half parameters raises `TypeError: gamma must be Cumo::SFloat, not Cumo::HFloat`.
+Keeping the running statistics in single precision is what the arithmetic wants in any case: a momentum update is a long chain of small corrections, and eleven bits of mantissa lose them.
+
+Gradients are the other place half runs out of room.
+Values below `Cumo::HFloat::MIN` of 6.1e-05 fall into the subnormals and then to zero, which is what loss scaling in a training loop exists to prevent.
+
 ### Select a GPU device ID
 
 Set the `CUDA_VISIBLE_DEVICES=id` environment variable, or
@@ -363,6 +483,19 @@ export CUMO_SHOW_WARNING=ON
 export CUMO_SHOW_WARNING_ONCE=OFF
 ```
 
+### Raise the cuDNN workspace ceiling
+
+cuDNN picks a convolution algorithm by benchmarking the ones that fit in a scratch buffer, and the ceiling on that buffer is 8MB.
+The fastest half precision algorithms, the ones that reach the tensor cores, ask for more than that and are left out of the search.
+To let them in:
+
+```
+export CUMO_CUDNN_MAX_WORKSPACE_SIZE=67108864
+```
+
+The value is in bytes and only bounds the search; each convolution reserves what its chosen algorithm actually needs.
+`Cumo::CUDA::CUDNN.max_workspace_size` reads back the value in force.
+
 ## Contributing
 
 Bug reports and pull requests are welcome on GitHub at https://github.com/sonots/cumo.

data/bench/cumo_sync_probe.rb

@@ -90,8 +90,8 @@ end
 
 PROBES = []
 
-def probe(group, name, &setup)
-  PROBES << { group: group, name: name, setup: setup }
+def probe(group, name, counts_only: false, &setup)
+  PROBES << { group: group, name: name, setup: setup, counts_only: counts_only }
 end
 
 # Reading data back is meant to synchronize. The question is whether it says so.
@@ -283,6 +283,47 @@ probe('compute', 'index aref a[idx, true]') do
 end
 
 # Host to device
+# A block walks the array on the host and can queue work on it, so the walk has
+# to wait between elements. It only has to wait when the block queued something,
+# which is what the idle count says: nothing for a block that stays off the
+# device, once per touched element for one that does not.
+BLOCK_N = 32
+
+probe('block', 'each (block off the device)', counts_only: true) do
+  a = Cumo::SFloat.new(BLOCK_N).seq
+  -> { a.each { |x| x } }
+end
+
+probe('block', 'each (block writes)', counts_only: true) do
+  a = Cumo::SFloat.new(BLOCK_N).seq
+  -> { i = 0; a.each { |_x| a[i] = 1; i += 1 } }
+end
+
+probe('block', 'each_with_index (block writes)', counts_only: true) do
+  a = Cumo::SFloat.new(BLOCK_N).seq
+  -> { a.each_with_index { |_x, i| a[i] = 1 } }
+end
+
+probe('block', 'map (block writes)', counts_only: true) do
+  a = Cumo::SFloat.new(BLOCK_N).seq
+  -> { i = 0; a.map { |x| a[i] = 1; i += 1; x } }
+end
+
+probe('block', 'map_with_index (block writes)', counts_only: true) do
+  a = Cumo::SFloat.new(BLOCK_N).seq
+  -> { a.map_with_index { |x, i| a[i] = 1; x } }
+end
+
+probe('block', 'Bit each (block writes)', counts_only: true) do
+  a = Cumo::Bit.new(BLOCK_N).fill(1)
+  -> { i = 0; a.each { |_x| a[i] = 0; i += 1 } }
+end
+
+probe('block', 'Bit each_with_index (block writes)', counts_only: true) do
+  a = Cumo::Bit.new(BLOCK_N).fill(1)
+  -> { a.each_with_index { |_x, i| a[i] = 0 } }
+end
+
 probe('h2d', 'from_binary') do
   bin = Array.new(BIG_N / 4, 1.0).pack('f*')
   -> { Cumo::SFloat.from_binary(bin, [BIG_N / 4]) }
@@ -347,8 +388,8 @@ puts
 
 results = []
 
-puts format('  %-26s %10s %12s  %-9s %s', 'probe', 'idle', 'behind queue', 'verdict', 'warning')
-puts "  #{'-' * 84}"
+puts format('  %-34s %10s %6s %12s  %-9s %s', 'probe', 'idle', 'waits', 'behind queue', 'verdict', 'warning')
+puts "  #{'-' * 99}"
 
 current_group = nil
 PROBES.each do |p|
@@ -363,7 +404,22 @@ PROBES.each do |p|
   begin
     call = p[:setup].call
     if call.nil?
-      puts format('  %-26s %10s %12s  %-9s %s', p[:name], '-', '-', 'skip', '(no such API)')
+      puts format('  %-34s %10s %6s %12s  %-9s %s', p[:name], '-', '-', '-', 'skip', '(no such API)')
+      next
+    end
+
+    # (1b) how many times it waits with nothing queued. With the warning left
+    # on every occurrence this is the wait count, not a yes or no.
+    sync
+    idle_text = capture_stdio { call.call }
+    idle_warns = idle_text.to_s.lines.grep(/warn|sync|Sync|WARN/i).size
+
+    if p[:counts_only]
+      # Timing these would measure the warning writes, not the call.
+      puts format('  %-34s %10s %6d %12s  %-9s %s',
+                  p[:name], '-', idle_warns, '-', 'waits', '-')
+      results << { name: p[:name], group: p[:group], sync: idle_warns.positive?,
+                   warned: idle_warns.positive? }
       next
     end
 
@@ -412,12 +468,12 @@ PROBES.each do |p|
         format('x%d %s', warns.size, warns.first.to_s.strip[0, 40])
       end
 
-    puts format('  %-26s %8.1f us %10.2f ms  %-9s %s',
-                p[:name], idle * 1e6, blocked * 1e3,
+    puts format('  %-34s %8.1f us %6d %10.2f ms  %-9s %s',
+                p[:name], idle * 1e6, idle_warns, blocked * 1e3,
                 is_sync ? 'SYNC' : 'async', note)
   rescue StandardError, NotImplementedError => e
     msg = e.message.to_s.split("\n").first.to_s[0, 40]
-    puts format('  %-26s %10s %12s  %-9s %s', p[:name], '-', '-', 'ERROR', "#{e.class}: #{msg}")
+    puts format('  %-34s %10s %6s %12s  %-9s %s', p[:name], '-', '-', '-', 'ERROR', "#{e.class}: #{msg}")
   end
 end
 
@@ -440,3 +496,10 @@ puts '  The first line is what this is for: those calls stall the pipeline and'
 puts '  CUMO_SHOW_WARNING will not tell you. Either add the warning or, where'
 puts '  the synchronization is avoidable, remove it.'
 puts '  A SYNC in the meta group is waste with nothing to show for it.'
+puts
+puts '  The block group counts waits rather than timing them: a walk that hands'
+puts '  every element to a Ruby block has to wait between elements, since the'
+puts '  block can queue work on the array being walked. A block that stays off'
+puts '  the device should cost none, and one that writes should cost about one'
+puts '  per element. A count of one there means the walk waits once for the whole'
+puts '  row, which is what made Bit#each and map answer differently from numo.'

data/ext/cumo/cuda/cublas.c

@@ -76,46 +76,6 @@ cumo_cuda_cublas_option_value(VALUE value, VALUE default_value)
     return value;
 }
 
-#if 0
-cublasOperation_t
-cumo_cuda_cublas_option_trans(VALUE trans)
-{
-    int opt;
-    char *ptr;
-
-    switch(TYPE(trans)) {
-    case T_NIL:
-    case T_UNDEF:
-    case T_FALSE:
-        return CUBLAS_OP_N;
-    case T_TRUE:
-        return CUBLAS_OP_T;
-    case T_FIXNUM:
-        opt = FIX2INT(trans);
-        if (opt >= CUBLAS_OP_N && opt <= CUBLAS_OP_C) {
-            return opt;
-        }
-        break;
-    case T_SYMBOL:
-        trans = rb_sym2str(trans);
-    case T_STRING:
-        ptr = RSTRING_PTR(trans);
-        if (RSTRING_LEN(trans) > 0) {
-            switch(ptr[0]){
-            case 'N': case 'n':
-                return CUBLAS_OP_N;
-            case 'T': case 't':
-                return CUBLAS_OP_T;
-            case 'C': case 'c':
-                return CUBLAS_OP_C;
-            }
-        }
-        break;
-    }
-    rb_raise(rb_eArgError, "invalid value for cublasOperation_t");
-    return 0;
-}
-#endif
 
 void
 Init_cumo_cuda_cublas(void)

data/ext/cumo/cuda/cudnn.c

@@ -1,6 +1,9 @@
 #include "cumo/cuda/cudnn.h"
 
 #include <assert.h>
+#include <stdlib.h>
+#include <stdint.h>
+#include <errno.h>
 #include <ruby.h>
 #include "cumo/narray.h"
 #include "cumo/template.h"
@@ -21,6 +24,35 @@ cumo_cuda_cudnn_check_status(cudnnStatus_t status)
     }
 }
 
+static size_t cudnn_max_workspace_size = CUMO_CUDA_CUDNN_DEFAULT_MAX_WORKSPACE_SIZE;
+
+size_t
+cumo_cuda_cudnn_max_workspace_size()
+{
+    return cudnn_max_workspace_size;
+}
+
+static void
+init_max_workspace_size(void)
+{
+    const char* env = getenv("CUMO_CUDNN_MAX_WORKSPACE_SIZE");
+    char* end = NULL;
+    unsigned long long v;
+
+    if (env == NULL || *env == '\0') return;
+    // strtoull takes a leading minus and wraps it, so reject the sign itself
+    if (*env == '-' || *env == '+') goto BAD;
+    errno = 0;
+    v = strtoull(env, &end, 10);
+    if (errno != 0 || end == env || *end != '\0' || v == 0 || v > (unsigned long long)SIZE_MAX) goto BAD;
+    cudnn_max_workspace_size = (size_t)v;
+    return;
+
+BAD:
+    rb_warn("CUMO_CUDNN_MAX_WORKSPACE_SIZE=%s is not a positive byte count, using %"PRIuSIZE,
+            env, cudnn_max_workspace_size);
+}
+
 // Lazily initialize cudnn handle, and cache it
 cudnnHandle_t
 cumo_cuda_cudnn_handle()
@@ -51,6 +83,20 @@ cumo_cuda_cudnn_handle()
 
   @return [Boolean] Returns true if cuDNN is available
  */
+#ifdef CUDNN_FOUND
+/*
+  Returns the ceiling cuDNN may use when it searches for a convolution
+  algorithm, set by CUMO_CUDNN_MAX_WORKSPACE_SIZE.
+
+  @return [Integer] bytes
+ */
+static VALUE
+rb_cudnn_max_workspace_size(VALUE self)
+{
+    return SIZET2NUM(cumo_cuda_cudnn_max_workspace_size());
+}
+#endif // CUDNN_FOUND
+
 static VALUE
 rb_cudnn_available_p(VALUE self)
 {
@@ -76,6 +122,8 @@ Init_cumo_cuda_cudnn(void)
 
     rb_define_singleton_method(mCUDNN, "available?", rb_cudnn_available_p, 0);
 #ifdef CUDNN_FOUND
+    init_max_workspace_size();
+    rb_define_singleton_method(mCUDNN, "max_workspace_size", rb_cudnn_max_workspace_size, 0);
     rb_define_const(mCUDNN, "CUDNN_POOLING_MAX", INT2NUM(CUDNN_POOLING_MAX));
     rb_define_const(mCUDNN, "CUDNN_POOLING_MAX_DETERMINISTIC", INT2NUM(CUDNN_POOLING_MAX_DETERMINISTIC));
     rb_define_const(mCUDNN, "CUDNN_POOLING_AVERAGE_COUNT_INCLUDE_PADDING", INT2NUM(CUDNN_POOLING_AVERAGE_COUNT_INCLUDE_PADDING));