cubecl_runtime/
server.rs

1use crate::{
2    DeviceProperties,
3    compiler::CompilationError,
4    kernel::KernelMetadata,
5    logging::ServerLogger,
6    memory_management::{
7        MemoryAllocationMode, MemoryHandle, MemoryUsage,
8        memory_pool::{SliceBinding, SliceHandle},
9    },
10    storage::{BindingResource, ComputeStorage},
11    tma::{OobFill, TensorMapFormat, TensorMapInterleave, TensorMapPrefetch, TensorMapSwizzle},
12};
13use alloc::collections::BTreeMap;
14use alloc::string::String;
15use alloc::sync::Arc;
16use alloc::vec;
17use alloc::vec::Vec;
18use core::fmt::Debug;
19use cubecl_common::{
20    ExecutionMode, bytes::Bytes, device, future::DynFut, profile::ProfileDuration,
21    stream_id::StreamId,
22};
23use cubecl_ir::StorageType;
24use thiserror::Error;
25
26#[derive(Debug, Clone)]
27/// An error during profiling.
28pub enum ProfileError {
29    /// Unknown error.
30    Unknown(String),
31    /// When no profiling has been registered.
32    NotRegistered,
33    /// An error happened when launching a kernel.
34    Launch(LaunchError),
35    /// An error happened when executing runtime operations.
36    Execution(ExecutionError),
37}
38
39impl From<LaunchError> for ProfileError {
40    fn from(val: LaunchError) -> Self {
41        ProfileError::Launch(val)
42    }
43}
44
45impl From<ExecutionError> for ProfileError {
46    fn from(val: ExecutionError) -> Self {
47        Self::Execution(val)
48    }
49}
50
51#[derive(Debug)]
52/// Contains many different types that are useful for server implementations and compute clients.
53pub struct ServerUtilities<Server: ComputeServer> {
54    /// The time when `profile-tracy` is activated.
55    #[cfg(feature = "profile-tracy")]
56    pub epoch_time: web_time::Instant,
57    /// The GPU client when `profile-tracy` is activated.
58    #[cfg(feature = "profile-tracy")]
59    pub gpu_client: tracy_client::GpuContext,
60    /// Information shared between all servers.
61    pub properties: DeviceProperties,
62    /// Information specific to the current server.
63    pub info: Server::Info,
64    /// The logger based on global cubecl configs.
65    pub logger: Arc<ServerLogger>,
66}
67
68impl<S: ComputeServer> ServerUtilities<S> {
69    /// Creates a new server utilities.
70    pub fn new(properties: DeviceProperties, logger: Arc<ServerLogger>, info: S::Info) -> Self {
71        // Start a tracy client if needed.
72        #[cfg(feature = "profile-tracy")]
73        let client = tracy_client::Client::start();
74
75        Self {
76            properties,
77            logger,
78            // Create the GPU client if needed.
79            #[cfg(feature = "profile-tracy")]
80            gpu_client: client
81                .clone()
82                .new_gpu_context(
83                    Some(&format!("{info:?}")),
84                    // In the future should ask the server what makes sense here. 'Invalid' atm is a generic stand-in (Tracy doesn't have CUDA/RocM atm anyway).
85                    tracy_client::GpuContextType::Invalid,
86                    0,   // Timestamps are manually aligned to this epoch so start at 0.
87                    1.0, // Timestamps are manually converted to be nanoseconds so period is 1.
88                )
89                .unwrap(),
90            #[cfg(feature = "profile-tracy")]
91            epoch_time: web_time::Instant::now(),
92            info,
93        }
94    }
95}
96
97/// Error that can happen when calling [ComputeServer::execute];
98///
99/// # Notes
100///
101/// Not all errors are going to be catched when calling [ComputeServer::execute] only the one that
102/// won't block the compute queue.
103#[derive(Debug, PartialEq, Eq, Clone, Hash)]
104#[cfg_attr(std_io, derive(serde::Serialize, serde::Deserialize))]
105pub enum LaunchError {
106    /// The given kernel can't be compiled.
107    CompilationError(CompilationError),
108    /// The server is out of memory.
109    OutOfMemory {
110        /// The details of the memory error.
111        context: String,
112    },
113    /// Unknown launch error.
114    Unknown {
115        /// The details of the unknown error.
116        context: String,
117    },
118    /// Can't launch because of an IO Error.
119    IoError(IoError),
120}
121
122/// Error that can happen asynchronously while executing registered kernels.
123#[derive(Debug, PartialEq, Eq, Clone, Hash)]
124#[cfg_attr(std_io, derive(serde::Serialize, serde::Deserialize))]
125pub enum ExecutionError {
126    /// A generic runtime error.
127    Generic {
128        /// The details of the generic error.
129        context: String,
130    },
131    /// When multiple errors happened during runtime.
132    Composed {
133        /// The details of the error.
134        context: String,
135        /// All the underlying errors.
136        errors: Vec<Self>,
137    },
138}
139
140impl From<CompilationError> for LaunchError {
141    fn from(value: CompilationError) -> Self {
142        Self::CompilationError(value)
143    }
144}
145
146impl From<IoError> for LaunchError {
147    fn from(value: IoError) -> Self {
148        Self::IoError(value)
149    }
150}
151
152impl core::fmt::Display for LaunchError {
153    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
154        match self {
155            LaunchError::CompilationError(err) => f.write_fmt(format_args!(
156                "A compilation error happened during launch: {err}"
157            )),
158            LaunchError::OutOfMemory { context } => f.write_fmt(format_args!(
159                "Out of memory error happened during launch: {context}"
160            )),
161            LaunchError::Unknown { context } => f.write_fmt(format_args!(
162                "An unknown error happened during launch: {context}"
163            )),
164            LaunchError::IoError(err) => {
165                f.write_fmt(format_args!("Can't launch because of an IO error: {err}"))
166            }
167        }
168    }
169}
170
171/// The compute server is responsible for handling resources and computations over resources.
172///
173/// Everything in the server is mutable, therefore it should be solely accessed through the
174/// [compute channel](crate::channel::ComputeChannel) for thread safety.
175pub trait ComputeServer:
176    Send + core::fmt::Debug + ServerCommunication + device::DeviceState + 'static
177where
178    Self: Sized,
179{
180    /// The kernel type defines the computation algorithms.
181    type Kernel: KernelMetadata;
182    /// Information that can be retrieved for the runtime.
183    type Info: Debug + Send + Sync;
184    /// The [storage](ComputeStorage) type defines how data is stored and accessed.
185    type Storage: ComputeStorage;
186
187    /// Reserves `size` bytes in the storage, and returns a handle over them.
188    fn create(
189        &mut self,
190        descriptors: Vec<AllocationDescriptor<'_>>,
191        stream_id: StreamId,
192    ) -> Result<Vec<Allocation>, IoError>;
193
194    /// Reserves N [Bytes] of the provided sizes to be used as staging to load data.
195    fn staging(&mut self, _sizes: &[usize], _stream_id: StreamId) -> Result<Vec<Bytes>, IoError> {
196        Err(IoError::UnsupportedIoOperation)
197    }
198
199    /// Retrieve the server logger.
200    fn logger(&self) -> Arc<ServerLogger>;
201
202    /// Retrieve the server utilities.
203    fn utilities(&self) -> Arc<ServerUtilities<Self>>;
204
205    /// Utility to create a new buffer and immediately copy contiguous data into it
206    fn create_with_data(&mut self, data: &[u8], stream_id: StreamId) -> Result<Handle, IoError> {
207        let alloc = self
208            .create(
209                vec![AllocationDescriptor::new(
210                    AllocationKind::Contiguous,
211                    &[data.len()],
212                    1,
213                )],
214                stream_id,
215            )?
216            .remove(0);
217        self.write(
218            vec![(
219                CopyDescriptor::new(
220                    alloc.handle.clone().binding(),
221                    &[data.len()],
222                    &alloc.strides,
223                    1,
224                ),
225                Bytes::from_bytes_vec(data.to_vec()),
226            )],
227            stream_id,
228        )?;
229        Ok(alloc.handle)
230    }
231
232    /// Utility to create a new buffer and immediately copy contiguous data into it
233    fn create_with_bytes(&mut self, data: Bytes, stream_id: StreamId) -> Result<Handle, IoError> {
234        let alloc = self
235            .create(
236                vec![AllocationDescriptor::new(
237                    AllocationKind::Contiguous,
238                    &[data.len()],
239                    1,
240                )],
241                stream_id,
242            )?
243            .remove(0);
244        self.write(
245            vec![(
246                CopyDescriptor::new(
247                    alloc.handle.clone().binding(),
248                    &[data.len()],
249                    &alloc.strides,
250                    1,
251                ),
252                data,
253            )],
254            stream_id,
255        )?;
256        Ok(alloc.handle)
257    }
258
259    /// Given bindings, returns the owned resources as bytes.
260    fn read<'a>(
261        &mut self,
262        descriptors: Vec<CopyDescriptor<'a>>,
263        stream_id: StreamId,
264    ) -> DynFut<Result<Vec<Bytes>, IoError>>;
265
266    /// Writes the specified bytes into the buffers given
267    fn write(
268        &mut self,
269        descriptors: Vec<(CopyDescriptor<'_>, Bytes)>,
270        stream_id: StreamId,
271    ) -> Result<(), IoError>;
272
273    /// Wait for the completion of every task in the server.
274    fn sync(&mut self, stream_id: StreamId) -> DynFut<Result<(), ExecutionError>>;
275
276    /// Given a resource handle, returns the storage resource.
277    fn get_resource(
278        &mut self,
279        binding: Binding,
280        stream_id: StreamId,
281    ) -> BindingResource<<Self::Storage as ComputeStorage>::Resource>;
282
283    /// Executes the `kernel` over the given memory `handles`.
284    ///
285    /// Kernels have mutable access to every resource they are given
286    /// and are responsible of determining which should be read or written.
287    ///
288    /// # Safety
289    ///
290    /// When executing with mode [ExecutionMode::Unchecked], out-of-bound reads and writes can happen.
291    unsafe fn launch(
292        &mut self,
293        kernel: Self::Kernel,
294        count: CubeCount,
295        bindings: Bindings,
296        kind: ExecutionMode,
297        stream_id: StreamId,
298    ) -> Result<(), LaunchError>;
299
300    /// Flush all outstanding tasks in the server.
301    fn flush(&mut self, stream_id: StreamId);
302
303    /// The current memory usage of the server.
304    fn memory_usage(&mut self, stream_id: StreamId) -> MemoryUsage;
305
306    /// Ask the server to release memory that it can release.
307    fn memory_cleanup(&mut self, stream_id: StreamId);
308
309    /// Enable collecting timestamps.
310    fn start_profile(&mut self, stream_id: StreamId) -> ProfilingToken;
311
312    /// Disable collecting timestamps.
313    fn end_profile(
314        &mut self,
315        stream_id: StreamId,
316        token: ProfilingToken,
317    ) -> Result<ProfileDuration, ProfileError>;
318
319    /// Update the memory mode of allocation in the server.
320    fn allocation_mode(&mut self, mode: MemoryAllocationMode, stream_id: StreamId);
321}
322
323/// Defines functions for optimized data transfer between servers, supporting custom communication
324/// mechanisms such as peer-to-peer communication or specialized implementations.
325pub trait ServerCommunication {
326    /// Indicates whether server-to-server communication is enabled for this implementation.
327    const SERVER_COMM_ENABLED: bool;
328
329    /// Copies data from a source server to a destination server.
330    ///
331    /// # Arguments
332    ///
333    /// * `server_src` - A mutable reference to the source server from which data is copied.
334    /// * `server_dst` - A mutable reference to the destination server receiving the data.
335    /// * `src` - A descriptor specifying the data to be copied, including shape, strides, and binding.
336    /// * `stream_id_src` - The stream ID associated with the source server's operation.
337    /// * `stream_id_dst` - The stream ID associated with the destination server's operation.
338    ///
339    /// # Returns
340    ///
341    /// Returns a `Result` containing an `Allocation` on success, or an `IoError` if the operation fails.
342    ///
343    /// # Panics
344    ///
345    /// Panics if server communication is not enabled (`SERVER_COMM_ENABLED` is `false`) or if the
346    /// trait is incorrectly implemented by the server.
347    #[allow(unused_variables)]
348    fn copy(
349        server_src: &mut Self,
350        server_dst: &mut Self,
351        src: CopyDescriptor<'_>,
352        stream_id_src: StreamId,
353        stream_id_dst: StreamId,
354    ) -> Result<Allocation, IoError> {
355        if !Self::SERVER_COMM_ENABLED {
356            panic!("Server-to-server communication is not supported by this server.");
357        } else {
358            panic!(
359                "[Internal Error] The `ServerCommunication` trait is incorrectly implemented by the server."
360            );
361        }
362    }
363}
364
365#[derive(Debug, Clone, Copy, Hash, PartialEq, Eq)]
366/// Profiling identification so that the server can support recursive and overlapping profilings.
367pub struct ProfilingToken {
368    /// The token value.
369    pub id: u64,
370}
371
372/// Server handle containing the [memory handle](crate::server::Handle).
373#[derive(new, Debug, PartialEq, Eq)]
374pub struct Handle {
375    /// Memory handle.
376    pub memory: SliceHandle,
377    /// Memory offset in bytes.
378    pub offset_start: Option<u64>,
379    /// Memory offset in bytes.
380    pub offset_end: Option<u64>,
381    /// The stream where the data was created.
382    pub stream: cubecl_common::stream_id::StreamId,
383    /// The stream position when the tensor became available.
384    pub cursor: u64,
385    /// Length of the underlying buffer ignoring offsets
386    size: u64,
387}
388
389/// Type of allocation, either contiguous or optimized (row-aligned when possible)
390#[derive(Debug, Clone, Copy, Hash, PartialEq, Eq)]
391pub enum AllocationKind {
392    /// Contiguous layout, with no padding
393    Contiguous,
394    /// Optimized for access speed. In practice this means row-aligned with padding for runtimes
395    /// that support it.
396    Optimized,
397}
398
399/// Descriptor for a new tensor allocation
400#[derive(new, Debug, Clone, Copy)]
401pub struct AllocationDescriptor<'a> {
402    /// Layout for the tensor
403    pub kind: AllocationKind,
404    /// Shape of the tensor
405    pub shape: &'a [usize],
406    /// Size of each element in the tensor (used for conversion of shape to bytes)
407    pub elem_size: usize,
408}
409
410impl<'a> AllocationDescriptor<'a> {
411    /// Create an optimized allocation descriptor
412    pub fn optimized(shape: &'a [usize], elem_size: usize) -> Self {
413        AllocationDescriptor::new(AllocationKind::Optimized, shape, elem_size)
414    }
415
416    /// Create a contiguous allocation descriptor
417    pub fn contiguous(shape: &'a [usize], elem_size: usize) -> Self {
418        AllocationDescriptor::new(AllocationKind::Contiguous, shape, elem_size)
419    }
420}
421
422/// An allocation with associated strides. Strides depend on tensor layout.
423#[derive(new, Debug)]
424pub struct Allocation {
425    /// The handle for the memory resource
426    pub handle: Handle,
427    /// The strides of the tensor
428    pub strides: Vec<usize>,
429}
430
431/// Error returned from `create`/`read`/`write` functions. Due to async execution not all errors
432/// are able to be caught, so some IO errors will still panic.
433#[derive(Debug, Error, PartialEq, Eq, Clone, Hash)]
434#[cfg_attr(std_io, derive(serde::Serialize, serde::Deserialize))]
435pub enum IoError {
436    /// Buffer size exceeds the max available
437    #[error("can't allocate buffer of size")]
438    BufferTooBig(usize),
439    /// Strides aren't supported for this copy operation on this runtime
440    #[error("the provided strides are not supported for this operation")]
441    UnsupportedStrides,
442    /// Handle wasn't found in the memory pool
443    #[error("couldn't find resource for that handle")]
444    InvalidHandle,
445    /// Unknown error happened during execution
446    #[error("Unknown error happened during execution")]
447    Unknown(String),
448    /// The current IO operation is not supported
449    #[error("The current IO operation is not supported")]
450    UnsupportedIoOperation,
451    /// Can't perform the IO operation because of a runtime error.
452    #[error("Can't perform the IO operation because of a runtime error")]
453    Execution(ExecutionError),
454}
455
456impl From<ExecutionError> for IoError {
457    fn from(value: ExecutionError) -> Self {
458        Self::Execution(value)
459    }
460}
461
462impl Handle {
463    /// Add to the current offset in bytes.
464    pub fn offset_start(mut self, offset: u64) -> Self {
465        if let Some(val) = &mut self.offset_start {
466            *val += offset;
467        } else {
468            self.offset_start = Some(offset);
469        }
470
471        self
472    }
473    /// Add to the current offset in bytes.
474    pub fn offset_end(mut self, offset: u64) -> Self {
475        if let Some(val) = &mut self.offset_end {
476            *val += offset;
477        } else {
478            self.offset_end = Some(offset);
479        }
480
481        self
482    }
483
484    /// Get the size of the handle, in bytes, accounting for offsets
485    pub fn size(&self) -> u64 {
486        self.size - self.offset_start.unwrap_or(0) - self.offset_end.unwrap_or(0)
487    }
488}
489
490/// Bindings to execute a kernel.
491#[derive(Debug, Default)]
492pub struct Bindings {
493    /// Buffer bindings
494    pub buffers: Vec<Binding>,
495    /// Packed metadata for tensor bindings (len, shape, stride, etc).
496    /// Ordered by inputs, then outputs, then tensormaps
497    pub metadata: MetadataBinding,
498    /// Scalar bindings
499    pub scalars: BTreeMap<StorageType, ScalarBinding>,
500    /// Tensor map bindings
501    pub tensor_maps: Vec<TensorMapBinding>,
502}
503
504impl Bindings {
505    /// Create a new bindings struct
506    pub fn new() -> Self {
507        Self::default()
508    }
509
510    /// Add a buffer binding
511    pub fn with_buffer(mut self, binding: Binding) -> Self {
512        self.buffers.push(binding);
513        self
514    }
515
516    /// Extend the buffers with `bindings`
517    pub fn with_buffers(mut self, bindings: Vec<Binding>) -> Self {
518        self.buffers.extend(bindings);
519        self
520    }
521
522    /// Add a scalar parameter
523    pub fn with_scalar(mut self, ty: StorageType, length: usize, data: Vec<u64>) -> Self {
524        self.scalars
525            .insert(ty, ScalarBinding::new(ty, length, data));
526        self
527    }
528
529    /// Extend the scalars with `bindings`
530    pub fn with_scalars(mut self, bindings: Vec<ScalarBinding>) -> Self {
531        self.scalars
532            .extend(bindings.into_iter().map(|binding| (binding.ty, binding)));
533        self
534    }
535
536    /// Set the metadata to `meta`
537    pub fn with_metadata(mut self, meta: MetadataBinding) -> Self {
538        self.metadata = meta;
539        self
540    }
541
542    /// Extend the tensor maps with `bindings`
543    pub fn with_tensor_maps(mut self, bindings: Vec<TensorMapBinding>) -> Self {
544        self.tensor_maps.extend(bindings);
545        self
546    }
547}
548
549/// Binding of a set of scalars of the same type to execute a kernel.
550#[derive(new, Debug, Default)]
551pub struct MetadataBinding {
552    /// Metadata values
553    pub data: Vec<u32>,
554    /// Length of the static portion (rank, len, buffer_len, shape_offsets, stride_offsets).
555    pub static_len: usize,
556}
557
558/// Binding of a set of scalars of the same type to execute a kernel.
559#[derive(new, Debug, Clone)]
560pub struct ScalarBinding {
561    /// Type of the scalars
562    pub ty: StorageType,
563    /// Unpadded length of the underlying data
564    pub length: usize,
565    /// Type-erased data of the scalars. Padded and represented by u64 to prevent misalignment.
566    pub data: Vec<u64>,
567}
568
569impl ScalarBinding {
570    /// Get data as byte slice
571    pub fn data(&self) -> &[u8] {
572        bytemuck::cast_slice(&self.data)
573    }
574}
575
576/// Binding of a [tensor handle](Handle) to execute a kernel.
577#[derive(new, Debug)]
578pub struct Binding {
579    /// Memory binding.
580    pub memory: SliceBinding,
581    /// Memory offset in bytes.
582    pub offset_start: Option<u64>,
583    /// Memory offset in bytes.
584    pub offset_end: Option<u64>,
585    /// The stream where the data was created.
586    pub stream: cubecl_common::stream_id::StreamId,
587    /// The stream position when the tensor became available.
588    pub cursor: u64,
589    /// Size in bytes
590    size: u64,
591}
592
593impl Binding {
594    /// Get the size of the handle, in bytes, accounting for offsets
595    pub fn size(&self) -> u64 {
596        self.size - self.offset_start.unwrap_or(0) - self.offset_end.unwrap_or(0)
597    }
598}
599
600/// A binding with shape and stride info for non-contiguous reading
601#[derive(new, Debug, Clone)]
602pub struct CopyDescriptor<'a> {
603    /// Binding for the memory resource
604    pub binding: Binding,
605    /// Shape of the resource
606    pub shape: &'a [usize],
607    /// Strides of the resource
608    pub strides: &'a [usize],
609    /// Size of each element in the resource
610    pub elem_size: usize,
611}
612
613/// A tensor map used with TMA ops
614#[derive(new, Debug, Clone)]
615pub struct TensorMapBinding {
616    /// The binding for the backing tensor
617    pub binding: Binding,
618    /// The tensormap metadata
619    pub map: TensorMapMeta,
620}
621
622/// TensorMap metadata for the opaque proxy used in TMA copies
623#[derive(Debug, Clone)]
624pub struct TensorMapMeta {
625    /// Tensormap format (tiled or im2col)
626    pub format: TensorMapFormat,
627    /// Rank of the backing tensor
628    pub rank: usize,
629    /// Shape of the backing tensor
630    pub shape: Vec<usize>,
631    /// Strides of the backing tensor
632    pub strides: Vec<usize>,
633    /// Element stride, usually 1 but may be 2 for complex tensors
634    /// For im2col, this is equivalent to the kernel stride
635    pub elem_stride: Vec<usize>,
636    /// Interleave mode
637    pub interleave: TensorMapInterleave,
638    /// Swizzle mode
639    pub swizzle: TensorMapSwizzle,
640    /// Prefetch settings
641    pub prefetch: TensorMapPrefetch,
642    /// OOB fill value
643    pub oob_fill: OobFill,
644    /// Storage type
645    pub storage_ty: StorageType,
646}
647
648impl Handle {
649    /// If the tensor handle can be reused inplace.
650    pub fn can_mut(&self) -> bool {
651        self.memory.can_mut() && self.stream == StreamId::current()
652    }
653}
654
655impl Handle {
656    /// Convert the [handle](Handle) into a [binding](Binding).
657    pub fn binding(self) -> Binding {
658        Binding {
659            memory: MemoryHandle::binding(self.memory),
660            offset_start: self.offset_start,
661            offset_end: self.offset_end,
662            size: self.size,
663            stream: self.stream,
664            cursor: self.cursor,
665        }
666    }
667
668    /// Convert the [handle](Handle) into a [binding](Binding) with shape and stride metadata.
669    pub fn copy_descriptor<'a>(
670        &'a self,
671        shape: &'a [usize],
672        strides: &'a [usize],
673        elem_size: usize,
674    ) -> CopyDescriptor<'a> {
675        CopyDescriptor {
676            shape,
677            strides,
678            elem_size,
679            binding: self.clone().binding(),
680        }
681    }
682}
683
684impl Clone for Handle {
685    fn clone(&self) -> Self {
686        Self {
687            memory: self.memory.clone(),
688            offset_start: self.offset_start,
689            offset_end: self.offset_end,
690            size: self.size,
691            stream: self.stream,
692            cursor: self.cursor,
693        }
694    }
695}
696
697impl Clone for Binding {
698    fn clone(&self) -> Self {
699        Self {
700            memory: self.memory.clone(),
701            offset_start: self.offset_start,
702            offset_end: self.offset_end,
703            size: self.size,
704            stream: self.stream,
705            cursor: self.cursor,
706        }
707    }
708}
709
710/// Specifieds the number of cubes to be dispatched for a kernel.
711///
712/// This translates to eg. a grid for CUDA, or to num_workgroups for wgsl.
713#[allow(clippy::large_enum_variant)]
714pub enum CubeCount {
715    /// Dispatch a known count of x, y, z cubes.
716    Static(u32, u32, u32),
717    /// Dispatch an amount based on the values in this buffer. The buffer should contain a u32 array [x, y, z].
718    Dynamic(Binding),
719}
720
721impl CubeCount {
722    /// Create a new static cube count with the given x = y = z = 1.
723    pub fn new_single() -> Self {
724        CubeCount::Static(1, 1, 1)
725    }
726
727    /// Create a new static cube count with the given x, and y = z = 1.
728    pub fn new_1d(x: u32) -> Self {
729        CubeCount::Static(x, 1, 1)
730    }
731
732    /// Create a new static cube count with the given x and y, and z = 1.
733    pub fn new_2d(x: u32, y: u32) -> Self {
734        CubeCount::Static(x, y, 1)
735    }
736
737    /// Create a new static cube count with the given x, y and z.
738    pub fn new_3d(x: u32, y: u32, z: u32) -> Self {
739        CubeCount::Static(x, y, z)
740    }
741}
742
743impl Debug for CubeCount {
744    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
745        match self {
746            CubeCount::Static(x, y, z) => f.write_fmt(format_args!("({x}, {y}, {z})")),
747            CubeCount::Dynamic(_) => f.write_str("binding"),
748        }
749    }
750}
751
752impl Clone for CubeCount {
753    fn clone(&self) -> Self {
754        match self {
755            Self::Static(x, y, z) => Self::Static(*x, *y, *z),
756            Self::Dynamic(handle) => Self::Dynamic(handle.clone()),
757        }
758    }
759}
cubecl_runtime/server.rs

cubecl_runtime/
server.rs
