Fix codegen of NotMask on x64
Project: dotnet/runtime | Language: C# | Concepts: compiler-internals, algorithmic-optimization, testing-strategy | Difficulty: senior
Link: https://github.com/dotnet/runtime/pull/132437
The problem
On x64, the JIT's code generation for the NotMask AVX512 intrinsic was incorrect when the number of SIMD lanes was less than 8 (e.g., 2 or 4). The knotb instruction inverts all 8 bits of the mask register, but for vectors with fewer lanes, only the lower N bits are valid. The upper unused bits were left set, leading to incorrect results when those bits were later counted (e.g., by popcnt). This caused test failures such as in System.Numerics.Tensors.Tests.UInt32GenericTensorPrimitives.HammingDistance_AllLengths.
Why it's non-trivial
A naive fix would just apply the NotMask operation, assuming the upper bits are ignored. However, downstream consumers (like ExtractMostSignificantBits or PopCount) treat all bits as significant, so stale bits in the upper part of the mask register can corrupt results. The fix must clear these upper bits reliably, and do so in a way that doesn't cause code bloat or regressions elsewhere in the JIT's output.
The approach
1. Add a helper to clear unused mask bits
A new helper method was added to the code generator to zero out the unused upper bits of the mask register after a NotMask operation for small lane counts:
+void CodeGen::ClearUnusedMaskBits(regNumber maskReg, uint32_t count)
+{
+ assert((count == 2) || (count == 4));
+ assert(emitter::isMaskReg(maskReg));
+
+ emitter* emit = GetEmitter();
+
+ emit->emitIns_R_R_I(INS_kshiftlb, EA_8BYTE, maskReg, maskReg, (int8_t)(8 - count));
+ emit->emitIns_R_R_I(INS_kshiftrb, EA_8BYTE, maskReg, maskReg, (int8_t)(8 - count));
+}
Annotation:
This helper shifts the mask register left by the number of unused bits, then right by the same amount, effectively zeroing the upper bits above the live lane count. This is needed because knotb always inverts all 8 bits.
2. Use the helper after NotMask codegen for small lane counts
Wherever NotMask is emitted for a mask with fewer than 8 lanes, the code now calls the helper to clear the upper bits:
+ if (count < 8)
+ {
+ // Emit shifts to clear bits N to 7 for 2-bit or 4-bit NotMask.
+ // There is no 2 or 4-bit knot*. Normally not an issue, but would cause wrong codegen
+ // if k is used in a kmovb+POPCNT for example.
+
+ ClearUnusedMaskBits(targetReg, count);
+ }
Annotation:
This ensures that after a NotMask operation, only the valid bits for the actual SIMD width remain set. This block appears in multiple places where NotMask is used.
3. Remove redundant clearing from lowering phase
Previously, the lowering phase inserted explicit shift nodes to clear unused mask bits. With the new helper in codegen, this is no longer necessary:
- cnsNode = m_compiler->gtNewIconNode(8 - count);
- BlockRange().InsertAfter(maskNode, cnsNode);
-
- maskNode =
- m_compiler->gtNewSimdHWIntrinsicNode(TYP_MASK, maskNode, cnsNode, NI_AVX512_ShiftLeftMask,
- maskBaseType, simdSize);
- BlockRange().InsertAfter(cnsNode, maskNode);
- LowerNode(maskNode);
-
- cnsNode = m_compiler->gtNewIconNode(8 - count);
- BlockRange().InsertAfter(maskNode, cnsNode);
-
- maskNode =
- m_compiler->gtNewSimdHWIntrinsicNode(TYP_MASK, maskNode, cnsNode, NI_AVX512_ShiftRightMask,
- maskBaseType, simdSize);
- BlockRange().InsertAfter(cnsNode, maskNode);
- LowerNode(maskNode);
Annotation:
This code was responsible for inserting the left/right shifts at a higher IR level. It's now removed, relying on the new codegen-time fixup instead.
4. Optimize mask logic transformations
To recover codegen diffs and improve generated code, a new transformation was added:
~(a | ~b) → ~a & b
This enables the use of AndNotMask and avoids unnecessary fixups for NotMask:
+ if (mskIntrin->GetHWIntrinsicId() == NI_AVX512_OrMask)
+ {
+ // Transform ~(a | ~b) into ~a & b, to enable AndNotMask.
+ // Also has the nice effect of not having to fixup knotb during emit.
+ ...
+ if (transform)
+ {
+ mskIntrin->ChangeHWIntrinsicId(NI_AVX512_AndNotMask, op1, op2);
+ return true;
+ }
+ }
Annotation:
This pattern-matching optimization rewrites certain mask expressions to more efficient forms, reducing the need for bit-clearing fixups and producing better code.
5. Add targeted regression tests
A new test was added to verify correct codegen for 2- and 4-lane masks:
+ [Fact]
+ public static void TestEntryPoint()
+ {
+ // Vector128<uint> is 4 lanes, so only bits 0-3 may ever be set.
+ Vector128<uint> zero32 = Vector128<uint>.Zero;
+ Assert.Equal(0xFu, NotAnd128UInt32(zero32, zero32, zero32, zero32));
+ ...
+ // Vector128<ulong> is 2 lanes, the worst case at 6 stale bits.
+ Vector128<ulong> zero64 = Vector128<ulong>.Zero;
+ Assert.Equal(0x3u, NotAnd128UInt64(zero64, zero64, zero64, zero64));
+ ...
+ }
Annotation:
These tests confirm that only the intended bits are set in the result, catching the bug if stale upper bits are not cleared.
From the review thread
tannergooding: Is it worth extracting this to some
ClearUnusedMaskBitshelper since it's reusable and already in two places?
dhartglassMSFT: I think we might want both, I'll try a prototype locally but probably end up opening a follow-up to see if we can fix this up also in
fgOptimizeHWIntrinsic
Takeaways
- Always clear unused bits in mask registers when generating code for SIMD masks with fewer than the register width's worth of lanes. Downstream consumers may treat all bits as significant, so stale bits can cause subtle bugs.
- Centralize repeated logic in reusable helpers. Extracting bit-clearing logic into a helper reduces code duplication and ensures consistency.
- Pattern-matching and optimizing boolean expressions at the IR level can both improve code quality and reduce the need for downstream fixups. Recognizing and transforming expressions like
~(a | ~b)into~a & benables more efficient instruction selection and avoids unnecessary work. - Regression tests for edge cases (like minimal lane counts) are critical in compiler development. Subtle bugs often only manifest in these scenarios.