Cube option

[maxtremblay]

This make Option<T> where T: LaunchArg launchable in a cubecl kernel.
Combined with

match comptime!(option) { \* ... *\ }

it allows to match at comptime data that are runtime known. For example, for the first test named kernel_option_scalar, we can generate two different kernels based on an argument
of type Option<i32>.

// When option is None
extern "C" __global__ void kernel_option_scalar(int output_0[], uint info[]) {

  int threadIdxGlobal = threadIdx.x + threadIdx.y * blockDim.x +
                        threadIdx.z * (blockDim.x * blockDim.y);
  const bool l_0 = threadIdxGlobal == uint(0);
  if (l_0) {
  }
}

// When option is Some(value: i32)
extern "C" __global__ void kernel_option_scalar(int output_0[], uint info[],
                                                int scalars_i32[]) {

  int threadIdxGlobal = threadIdx.x + threadIdx.y * blockDim.x +
                        threadIdx.z * (blockDim.x * blockDim.y);
  const bool l_0 = threadIdxGlobal == uint(0);
  if (l_0) {
    const uint l_1 = info[uint(0)];
    const bool l_2 = uint(0) < l_1;
    if (l_2) {
      output_0[uint(0)] = scalars_i32[0];
    }
  }
}

Limitations

First, we can't use

let x = Some(y);

directly in a kernel. Thus, I introduce a new function named some that does exactly. However, it has some issue with type inference. A proper fix of the way we expand enum variant should solve the issue.

Second, we can't use stuff like option.map(|x| 2 * x) to generate code at comptime. Currently, we have to use a match statement to achieve this behavior.

Design choices

I hesitated between introducing a new enum

pub enum CubeOption<T> {
    Some(T),
    None
}

for which we could implement more utilities similar to std::Option. However, I decided to go with std::Option directly for now as a first step.

Justification

I believe this will quite useful to implement quantization in the matmul without doubling all implementations.

[louisfd]

LGTM