pub struct Lppool3dForward<D: Float> {
pub name: &'static str,
pub id: String,
pub kd: i32,
pub kh: i32,
pub kw: i32,
pub stride_d: i32,
pub stride_h: i32,
pub stride_w: i32,
pub block_ow: i32,
pub kernel_source: String,
pub entry_point_source: String,
pub source: String,
/* private fields */
}Expand description
3-D Lp-norm pooling forward pass.
Grid: pid = (((b * C + c) * OD + od) * OH + oh) * num_ow_tiles + ow_tile
Constraints: no padding; OD = (D-KD)/STRIDE_D+1, etc.
Fields§
§name: &'static strThe kernel function’s name (e.g. "flash_attention2_forward").
id: StringUnique kernel identifier: fn_name + dtype(s) + const values joined by “__”.
kd: i32Compile-time kernel constant, from the annotated fn’s const generics.
kh: i32Compile-time kernel constant, from the annotated fn’s const generics.
kw: i32Compile-time kernel constant, from the annotated fn’s const generics.
stride_d: i32Compile-time kernel constant, from the annotated fn’s const generics.
stride_h: i32Compile-time kernel constant, from the annotated fn’s const generics.
stride_w: i32Compile-time kernel constant, from the annotated fn’s const generics.
block_ow: i32Compile-time kernel constant, from the annotated fn’s const generics.
kernel_source: StringThe original kernel function source.
entry_point_source: StringThe Rust source of the generated C-ABI entry-point wrapper function.
source: StringCombined source (kernel_source + "\n\n" + entry_point_source); used by the Kernel trait.
Implementations§
Trait Implementations§
Source§impl<D: Float> Kernel for Lppool3dForward<D>
impl<D: Float> Kernel for Lppool3dForward<D>
Source§type Args<'__a> = (*mut D, *mut D, i32, i32, i32, i32, i32, i32, i32, i32, f32)
type Args<'__a> = (*mut D, *mut D, i32, i32, i32, i32, i32, i32, i32, i32, f32)
Source§fn kernel_source(&self) -> &str
fn kernel_source(&self) -> &str
Source§fn entry_point_source(&self) -> &str
fn entry_point_source(&self) -> &str
Source§fn entry_point_name(&self) -> String
fn entry_point_name(&self) -> String
"{name}_entry_point".Source§impl<D: Float + Send + Sync + 'static> RuntimeOp for Lppool3dForward<D>
impl<D: Float + Send + Sync + 'static> RuntimeOp for Lppool3dForward<D>
Source§fn n_activation_inputs(&self) -> usize
fn n_activation_inputs(&self) -> usize
Source§fn param_shapes(&self, _: &[&[usize]], _: &[usize]) -> Vec<Vec<usize>>
fn param_shapes(&self, _: &[&[usize]], _: &[usize]) -> Vec<Vec<usize>>
LoadedModel::load() time to pre-allocate device buffers.
input_shapes / output_shape are concrete (batch dim resolved).Source§fn pack_args(
&self,
_: &[(RawPtr, &[usize])],
_: &[RawPtr],
_: RawPtr,
_: &[usize],
_: i32,
_: &mut dyn ArgVisitor,
)
fn pack_args( &self, _: &[(RawPtr, &[usize])], _: &[RawPtr], _: RawPtr, _: &[usize], _: i32, _: &mut dyn ArgVisitor, )
visitor in the correct order. Read moreSource§fn grid(&self, _: &[usize]) -> [u32; 3]
fn grid(&self, _: &[usize]) -> [u32; 3]
Source§fn param_names(&self) -> &'static [&'static str]
fn param_names(&self) -> &'static [&'static str]
param_shapes], in the same order.
Used as the suffix in the dotted key {node_name}.{slot_name}.
Return an empty slice for ops that have no named parameters.Source§fn forward_output_row_stride(&self, output_shape: &[usize]) -> usize
fn forward_output_row_stride(&self, output_shape: &[usize]) -> usize
output_shape[-1]). Kernels using TMA must round up to satisfy the
16-byte alignment constraint (e.g. 4 elements for f32).Source§fn param_init_data(&self, _param_idx: usize) -> Option<Vec<u8>>
fn param_init_data(&self, _param_idx: usize) -> Option<Vec<u8>>
param_idx
immediately after device buffer allocation. Return None to leave the
slot zero-initialised (the default for trained parameters).
Byte count must equal param_shapes()[param_idx].iter().product() * dtype_bytes.Source§fn compute_concrete_output_shape(
&self,
_input_shapes: &[&[usize]],
resolved: &[usize],
) -> Vec<usize>
fn compute_concrete_output_shape( &self, _input_shapes: &[&[usize]], resolved: &[usize], ) -> Vec<usize>
Source§fn n_launches(&self) -> usize
fn n_launches(&self) -> usize
Source§fn pack_args_for_launch(
&self,
launch_idx: usize,
inputs: &[(*mut c_void, &[usize])],
params: &[*mut c_void],
output: *mut c_void,
output_shape: &[usize],
output_row_stride: i32,
visitor: &mut dyn ArgVisitor,
)
fn pack_args_for_launch( &self, launch_idx: usize, inputs: &[(*mut c_void, &[usize])], params: &[*mut c_void], output: *mut c_void, output_shape: &[usize], output_row_stride: i32, visitor: &mut dyn ArgVisitor, )
i (0-indexed). Read moreSource§fn grid_for_launch(
&self,
launch_idx: usize,
input_shapes: &[&[usize]],
output_shape: &[usize],
) -> [u32; 3]
fn grid_for_launch( &self, launch_idx: usize, input_shapes: &[&[usize]], output_shape: &[usize], ) -> [u32; 3]
i. Receives concrete input shapes so that per-chunk
grids can be computed without storing them in the op. Read moreSource§fn has_backward(&self) -> bool
fn has_backward(&self) -> bool
Source§fn backward_grad_output_row_stride(&self, output_shape: &[usize]) -> usize
fn backward_grad_output_row_stride(&self, output_shape: &[usize]) -> usize
pack_backward_args. The default is the natural row-major
stride (output_shape[-1]). Kernels using TMA must round up to satisfy
the 16-byte alignment constraint (e.g. 4 elements for f32).Source§fn pack_backward_args(
&self,
inputs: &[(*mut c_void, &[usize])],
params: &[*mut c_void],
output: *mut c_void,
output_shape: &[usize],
grad_output: *mut c_void,
grad_output_row_stride: i32,
grad_inputs: &[*mut c_void],
grad_params: &[*mut c_void],
visitor: &mut dyn ArgVisitor,
)
fn pack_backward_args( &self, inputs: &[(*mut c_void, &[usize])], params: &[*mut c_void], output: *mut c_void, output_shape: &[usize], grad_output: *mut c_void, grad_output_row_stride: i32, grad_inputs: &[*mut c_void], grad_params: &[*mut c_void], visitor: &mut dyn ArgVisitor, )
Source§fn backward_block(&self) -> [u32; 3]
fn backward_block(&self) -> [u32; 3]
Source§fn backward_grid(
&self,
input_shapes: &[&[usize]],
output_shape: &[usize],
) -> [u32; 3]
fn backward_grid( &self, input_shapes: &[&[usize]], output_shape: &[usize], ) -> [u32; 3]
Source§fn n_backward_launches(&self) -> usize
fn n_backward_launches(&self) -> usize
Source§fn pack_backward_args_for_launch(
&self,
launch_idx: usize,
inputs: &[(*mut c_void, &[usize])],
params: &[*mut c_void],
output: *mut c_void,
output_shape: &[usize],
grad_output: *mut c_void,
grad_output_row_stride: i32,
grad_inputs: &[*mut c_void],
grad_params: &[*mut c_void],
visitor: &mut dyn ArgVisitor,
)
fn pack_backward_args_for_launch( &self, launch_idx: usize, inputs: &[(*mut c_void, &[usize])], params: &[*mut c_void], output: *mut c_void, output_shape: &[usize], grad_output: *mut c_void, grad_output_row_stride: i32, grad_inputs: &[*mut c_void], grad_params: &[*mut c_void], visitor: &mut dyn ArgVisitor, )
i (0-indexed). Read more