unlock unlimited code runs for $79, limited time →
NOSSLETTER.techLIVE
back to OSSpedia
llvm/llvm-project
senior7/10
algorithmic-optimizationdata-structures

[SelectionDAG] Preserve CTTZ_ELTS lane count during expansion

Project: llvm/llvm-project | Language: C++ | Concepts: algorithmic-optimization, data-structures | Difficulty: senior
Link: https://github.com/llvm/llvm-project/pull/217982

The problem

When lowering the llvm.experimental.cttz.elts intrinsic in LLVM's SelectionDAG, the code derived the vector length (VL) from a legalized step vector. On targets like X86 with AVX512F, this helper vector can be widened for legality, causing its lane count to differ from the logical lane count of the original operand. This mismatch led to incorrect results: for example, an all-zero <4 x i1> mask would return 16 instead of 4.

Why it's non-trivial

A naive fix might just use the lane count of the legalized step vector, but this doesn't match the semantic lane count of the original operation. Legalization can widen vectors for codegen reasons, but the logical "number of elements" in the operation must stay consistent with the input, regardless of how the helper vectors are widened internally.

The approach

The patch extracts the logical element count from the original operand and uses it when materializing the VL value, instead of relying on the (possibly widened) step vector's lane count.

Key change: preserve the original operand's element count:

+  ElementCount EC =
+      Node->getOperand(0).getValueType().getVectorElementCount();

This extracts the element count (EC) from the input operand, ensuring the logical lane count is preserved.

Use the preserved element count when materializing VL:

-  SDValue VL =
-      DAG.getElementCount(DL, StepVT, StepVecVT.getVectorElementCount());
+  SDValue VL = DAG.getElementCount(DL, StepVT, EC);

Previously, VL was derived from the step vector's (possibly widened) lane count. The fix uses the logical lane count (EC) from the original operand.

Test coverage for the affected cases:

+define i32 @ctz_zero_v4i1() {
+; AVX512-LABEL: ctz_zero_v4i1:
+; AVX512:         addb $5, %al
+  %res = call i32 @llvm.experimental.cttz.elts.i32.v4i1(<4 x i1> zeroinitializer, i1 false)
+  ret i32 %res
+}
+
+define i32 @ctz_zero_v8i1() {
+; AVX512-LABEL: ctz_zero_v8i1:
+; AVX512:         addb $9, %al
+  %res = call i32 @llvm.experimental.cttz.elts.i32.v8i1(<8 x i1> zeroinitializer, i1 false)
+  ret i32 %res
+}

New tests ensure that the result matches the logical lane count (4 or 8), not the widened helper vector's lane count.

From the review thread

No notable review discussion on this PR.

Takeaways

  • Always preserve and propagate logical data structure invariants (like element counts) even when helper representations (like legalized vectors) are widened or transformed for backend legality.
  • When writing code that manipulates data structures for target-specific lowering, distinguish between "physical" representation (how the data is laid out for codegen) and "semantic" meaning (what the operation is logically doing).
  • Add targeted regression tests for subtle backend bugs, especially when hardware-specific vectorization or legalization can change the shape of helper data.