Block Boundary Version

fn conv_1d_block_boundary[
    in_layout: Layout, out_layout: Layout, conv_layout: Layout, dtype: DType
](
    output: LayoutTensor[mut=False, dtype, out_layout],
    a: LayoutTensor[mut=False, dtype, in_layout],
    b: LayoutTensor[mut=False, dtype, conv_layout],
):
    global_i = block_dim.x * block_idx.x + thread_idx.x
    local_i = thread_idx.x
    # first: need to account for padding
    shared_a = tb[dtype]().row_major[TPB + CONV_2 - 1]().shared().alloc()
    shared_b = tb[dtype]().row_major[CONV_2]().shared().alloc()
    if global_i < SIZE_2:
        shared_a[local_i] = a[global_i]

    # second: load elements needed for convolution at block boundary
    if local_i < CONV_2 - 1:
        # indices from next block
        next_idx = global_i + TPB
        if next_idx < SIZE_2:
            shared_a[TPB + local_i] = a[next_idx]
        else:
            # Initialize out-of-bounds elements to 0 to avoid reading from uninitialized memory
            # which is an undefined behavior
            shared_a[TPB + local_i] = 0

    if local_i < CONV_2:
        shared_b[local_i] = b[local_i]

    barrier()

    if global_i < SIZE_2:
        var local_sum: output.element_type = 0

        @parameter
        for j in range(CONV_2):
            if global_i + j < SIZE_2:
                local_sum += shared_a[local_i + j] * shared_b[j]

        output[global_i] = local_sum

The solution handles block boundary cases in 1D convolution using extended shared memory. Here’s a detailed analysis:

Memory layout and sizing

Test Configuration:
- Full array size: SIZE_2 = 15 elements
- Grid: 2 blocks × 8 threads
- Convolution kernel: CONV_2 = 4 elements

Block 0 shared memory:  [0 1 2 3 4 5 6 7|8 9 10]  // TPB(8) + (CONV_2-1)(3) padding
Block 1 shared memory:  [8 9 10 11 12 13 14|0 0]  // Second block with padding

Size calculation:
- Main data: TPB elements (8)
- Overlap: CONV_2 - 1 elements (4 - 1 = 3)
- Total: TPB + CONV_2 - 1 = 8 + 4 - 1 = 11 elements

Implementation details

1. Shared Memory Allocation:

   # First: account for padding needed for convolution window
   shared_a = tb[dtype]().row_major[TPB + CONV_2 - 1]().shared().alloc()
   shared_b = tb[dtype]().row_major[CONV_2]().shared().alloc()
   

   This allocation pattern ensures we have enough space for both the block’s data and the overlap region.

2. Data Loading Strategy:

   # Main block data
   if global_i < SIZE_2:
       shared_a[local_i] = a[global_i]
   
   # Boundary data from next block
   if local_i < CONV_2 - 1:
       next_idx = global_i + TPB
       if next_idx < SIZE_2:
           shared_a[TPB + local_i] = a[next_idx]
       else:
           # Initialize out-of-bounds elements to 0 to avoid reading from uninitialized memory
           # which is an undefined behavior
           shared_a[TPB + local_i] = 0
   

   • Only threads with local_i < CONV_2 - 1 load boundary data
   • Prevents unnecessary thread divergence
   • Maintains memory coalescing for main data load
   • Explicitly zeroes out-of-bounds elements to avoid undefined behavior

3. Kernel Loading:

   if local_i < b_size:
       shared_b[local_i] = b[local_i]
   

   • Single load per thread
   • Bounded by kernel size

4. Convolution Computation:

   if global_i < SIZE_2:
       var local_sum: output.element_type = 0
       @parameter
       for j in range(CONV_2):
           if global_i + j < SIZE_2:
               local_sum += shared_a[local_i + j] * shared_b[j]
   

   • Uses @parameter for compile-time loop unrolling
   • Proper type inference with output.element_type
   • Semantically correct bounds check: only compute convolution for valid input positions

Memory access pattern analysis

1. Block 0 Access Pattern:

   Thread 0: [0 1 2 3] × [0 1 2 3]
   Thread 1: [1 2 3 4] × [0 1 2 3]
   Thread 2: [2 3 4 5] × [0 1 2 3]
   ...
   Thread 7: [7 8 9 10] × [0 1 2 3]  // Uses overlap data
   

2. Block 1 Access Pattern:

   Thread 0: [8 9 10 11] × [0 1 2 3]
   Thread 1: [9 10 11 12] × [0 1 2 3]
   ...
   Thread 7: [14 0 0 0] × [0 1 2 3]  // Zero padding at end
   

Performance optimizations

1. Memory Coalescing:

   • Main data load: Consecutive threads access consecutive memory
   • Boundary data: Only necessary threads participate
   • Single barrier synchronization point

2. Thread Divergence Minimization:

   • Clean separation of main and boundary loading
   • Uniform computation pattern within warps
   • Efficient bounds checking

3. Shared Memory Usage:

   • Optimal sizing to handle block boundaries
   • No bank conflicts in access pattern
   • Efficient reuse of loaded data

4. Boundary Handling:

   • Explicit zero initialization for out-of-bounds elements which prevents reading from uninitialized shared memory
   • Semantically correct boundary checking using global_i + j < SIZE_2 instead of shared memory bounds
   • Proper handling of edge cases without over-computation

Boundary Condition Improvement

The solution uses if global_i + j < SIZE_2: rather than checking shared memory bounds. This approach is:

• Mathematically correct: Only computes convolution where input data actually exists
• More efficient: Avoids unnecessary computations for positions beyond the input array
• Safer: Prevents reliance on zero-padding behavior in shared memory

This implementation achieves efficient cross-block convolution while maintaining:

• Memory safety through proper bounds checking
• High performance through optimized memory access
• Clean code structure using LayoutTensor abstractions
• Minimal synchronization overhead
• Mathematically sound boundary handling