triton.language.extra.cann.extension.conv1d

triton.language.extra.cann.extension.conv1d(input: tensor, weight: tensor, bias: tensor = None, stride=None, padding_size=None, dilation=None, groups=None, _semantic=None) → tensor

Applies a 1D convolution over an input signal.

参数:

• input (tensor) -- Input tensor of shape (N, C_in, L_in) or (C_in, L_in). N is a batch size, C denotes a number of channels, L is a length of signal sequence.

• weight (tensor) -- Weight tensor of shape (C_out, C_in // groups, kernel_size).

• bias (tensor or None) -- Bias tensor of shape (C_out) or None. Default: None.

• stride (int or Tuple[int]) -- The stride of the convolution kernel. Can be an int or a 1-element tuple.

• padding_size (int, Tuple[int], or str) -- Padding added to both sides of the input. Can be an int, a 1-element tuple, or a string. Can be a string {'valid', 'same'}, single number or a one-element tuple. padding_size='valid' is the same as no padding. padding_size='same' pads the input so the output has the same shape as the input. However, this mode doesn't support any stride values other than 1.

• dilation (int or Tuple[int]) -- The spacing between kernel elements. Can be an int or a 1-element tuple.

• groups (int) -- Number of blocked connections from input to output channels.

Example:

@triton.jit
def conv_kernel(input_ptr, weight_ptr, bias_ptr, output_ptr, N, C, L, K, BLOCK_SIZE: tl.constexpr):
    # Load a tile of input and weight
    input_block = tl.load(input_ptr + ...)
    weight_block = tl.load(weight_ptr + ...)

    # Perform 1D convolution
    # Using default stride=1, padding_size=0, dilation=1, groups=1
    conv_output = al.conv1d(
        input_block,
        weight_block,
        bias=None,
        stride=1,
        padding_size=0,
        dilation=1,
        groups=1,
    )

    # Store the result
    tl.store(output_ptr + ..., conv_output)

返回:

The output tensor of shape (N, C_out, L_out).

返回类型:

tensor