Block Boundary Version

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 0 0]  // Second block. data(7) + padding to fill grid(1) + (CONV_2-1)(3) 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 = stack_allocation[dtype=dtype, address_space=AddressSpace.SHARED](row_major[TPB + CONV_2 - 1]())
   shared_b = stack_allocation[dtype=dtype, address_space=AddressSpace.SHARED](row_major[CONV_2]())
   

   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]
   else:
       shared_a[local_i] = 0
   
   # 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: Note how starting from thread 4, global_i + j < SIZE_2 evaluates to False and hence iterations are skipped.

   Thread 0: [8  9 10 11] × [0 1 2 3]
   Thread 1: [9 10 11 12] × [0 1 2 3]
   ...
   Thread 4: [12 13 14] × [0 1 2]       // Zero padding at end
   Thread 5: [13 14]    × [0 1]
   Thread 6: [14]       × [0]
   Thread 7: skipped                    // global_i + j < SIZE_2 evaluates to false for all j, no computation
   

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 pattern 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 TileTensor abstractions
• Minimal synchronization overhead
• Mathematically sound boundary handling