Skip to content

Commit 1482ea6

Browse files
committed
* Specialize singleton strides for vectorization
1 parent 86dca3b commit 1482ea6

1 file changed

Lines changed: 7 additions & 2 deletions

File tree

cext/tile_kernel.cpp

Lines changed: 7 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -771,10 +771,15 @@ static ArraySpecializationBits compute_array_specialization_bits(
771771
bool is_shape_byte_aligned = shape_bitwidth % BYTE_BITWIDTH == 0;
772772
bool is_shape_divisible_by_16 = shape % DIVISOR_16 == 0;
773773

774-
if (is_stride_byte_aligned && is_stride_16_byte_divisible)
774+
// A size-one axis has no address contribution for dense indexing.
775+
// Keep its physical stride dynamic, but preserve the 16-byte
776+
// alignment specialization used by layout/vectorization analysis.
777+
bool is_singleton_stride_one = shape == 1 && stride == 1;
778+
if ((is_stride_byte_aligned && is_stride_16_byte_divisible) ||
779+
is_singleton_stride_one)
775780
ret.stride_16byte_divisible |= 1u << i;
776781

777-
if (stride == 1)
782+
if (stride == 1 && !is_singleton_stride_one)
778783
ret.stride_one |= 1u << i;
779784

780785
if (is_shape_byte_aligned && is_shape_divisible_by_16)

0 commit comments

Comments
 (0)