x86_64 efficiently set bit N in register DST if N < width(DST), otherwise DST = 0

Question

I am trying to find an efficient way to do the following in x86_64 assembly:

if(N < word_size) {
    dst[N] = 1; // as in Nth bit of dst = 1 
}
else {
    dst[word_size - 1:0] = 0
} 

Alternative I could get the desired result if the "else" case did not unset the other bits, or if the "if" case did unset the other bits. The important thing is that if N > word_size it will not set any bits

I am unable to find any instruction that might do this as bt[s/c], shlx, sal, rol, shld all appear to the take the module of src by width.

The use case is basically I will be iterating over a bit vector with a known length and want to either A) find the first set bit and return its position, or B) test all of the bits and if no set bit is found return length of the vector.

// rsi has length
L(keep_searching):
movq %(rdi), %rax
testq %rax, %rax
jnz L(found)
subq $64, rsi
jbe L(done) // this done will return origional value of rsi
addq $8, %rdi
jmp L(keep_searching)

I figure this could be vastly sped up if I could quickly set a bit in rax if rsi < 64 so I could drop the second branch. But for this to work it needs to have the behavior above i.e it can't set the bit of rsi % 64, it needs to set iff rsi < 64.

Does anyone know of an instruction that can do this? Every instruction I can think of to check uses module on src. Any help would be greatly appreciated.

Thanks!

A few versions that are working well for me for 32 bit. If I use MMX @PeterCordes pointed out pllsq is exactly what I want.

uint64_t __attribute__((noinline, noclone)) shift(uint64_t cnt) {
    uint64_t ret = 0;
    asm volatile(
        "cmpq $32, %[cnt]\n\t"
        "setbe %b[ret]\n\t"
        "shlxq %[cnt], %[ret], %[ret]\n\t"
        : [ ret ] "+r"(ret)
        : [ cnt ] "r"(cnt)
        : "cc");
    return ret;
}


uint64_t __attribute__((noinline, noclone)) shift2(uint64_t cnt) {
    uint64_t ret = 0, tmp = 0;
    asm volatile(
        "leaq -33(%[cnt]), %[tmp]\n\t"
        "movl $1, %k[ret]\n\t"
        "shlxq %[cnt], %[ret], %[ret]\n\t"
        "sarq $63, %[tmp]\n\t"
        "andq %[tmp], %[ret]\n\t"
        : [ ret ] "+r"(ret), [ tmp ] "+r"(tmp), [ cnt ] "+r"(cnt)
        :
        : "cc");
    return ret;
}


uint64_t __attribute__((noinline, noclone)) shift3(uint64_t cnt) {
    uint64_t ret, tmp;
    asm volatile(
        "leaq -33(%[cnt]), %[tmp]\n\t"
        "btsq %[cnt], %[ret]\n\t"
        "sarq $63, %[tmp]\n\t"
        "andq %[tmp], %[ret]\n\t"
        : [ ret ] "+r"(ret), [ tmp ] "+r"(tmp), [ cnt ] "+r"(cnt)
        :
        : "cc");
    return ret;
}

Answer 1

Haven't verified, but

mov rax, 1  // common
mov rdx, 0  // common

cmp rcx, 64
shlxq rbx, rax, rcx
cmova rbx, rdx

could be slightly more performant than the suggested alternative as the comparison and the shift are now independent and can be executed in parallel.

EDIT

From the use case it may seem this is an XY-problem -- an efficient way to iterate over bits in bitset is to use the n & (n-1) trick or variants; popcount(n ^ (n-1)) should give the index of the least bit set. n&=n-1 will clear the LSB.

Answer 2

SIMD shifts (like SSE2 psllq xmm1, xmm2) saturate the count, but that's unlikely to be useful here because you I think want to OR this into the data from memory as an end condition for a scalar version of this loop?

I'd be more inclined to use cmov from a zeroed register using FLAGS still set from sub. You can create 1<<(rsi&63) using BTS into a zeroed register before or after SUB; before SUB is good because BTS modifies CF. Note that rsi&63 is not affected by rsi -= 64.

This is probably not a good choice for a loop condition: just use a single-uop sub/ja, with a separate the test/jnz being normally not taken. One of these goes at the bottom instead of an unconditional jmp: that's the most obvious and basic optimization here: Why are loops always compiled into "do...while" style (tail jump)?

Or even better, use SSE2 (baseline for x86-64) to check 16 bytes (2 qwords or 4 dwords) at a time for all-non-zero. Or even POR together a couple vectors to check all at once if you expect not to find the first set bit soon, i.e. tune for large trip counts at the expense of slower handling of eventually finding it. (The last last iteration can be scalar).

(Have a look at glibc's strlen or especially memchr for more ideas about optimizing the large-array not-found-early case with SIMD. In that case they're using pminub to get a zero if any vector had a zero at that position, but you want the opposite: por to get a non-zero if any had a non-zero.)

ORing together two values from memory works for scalar, too, as a way to unroll.

    mov  (%rdi), %rax
    or  8(%rdi), %rax
    jnz  found
    ...
    add  $16, %rdi

But note that or/jnz is 2 uops while test/jnz is 1.
OTOH, getting cmpq $0, (%rdi) / jne to micro- and macro-fuse on Intel may not be possible; IIRC maybe with a register source. So memory-source or may be costing 2 more uops to do twice as much work, instead of just 1 more, if you tune really aggressively. You'd need to compare against a loop where you unroll and do two separate load/test/jcc or cmp-mem/jcc, to keep it fair for the loop overhead of pointer-increment logic. (And also unrolling logic to handle a possible odd number of qwords.)

---

But just as an exercise, let's see what we can do with your idea: in this case the non-zero shift result can be computed once ahead of time (because rsi-=64 doesn't change rsi%64), and hoisted out of the loop.

   xor  %edx, %edx
   bts  %rsi, %rdx        # rdx = 1 << (rsi&63)

// rsi has length
L(keep_searching):
   add  $8, %rdi
   xor  %eax, %eax        # need to re-create a zero every time
   sub  $64, %rsi
   cmovbe %rdx, %rax      # 0  or  1<<(rsi&63) to put a bit there for us to find

   or  -8(%rdi), %rax
   jnz  L(keep_searching)

found_or_done:
   tzcnt %rax, %rax
   add   orig_rsi?, %rax
   ...

Unfortunately OR can't macro-fuse with JCC the way TEST can. (Or SUB on Intel SnB-family). But memory-source OR is a single uop for the front-end.

Unfortunately cmovbe and cmova cost 2 uops because they need CF and ZF. (See What is a Partial Flag Stall? - recent Intel don't have partial-flag stalls or even merging, just CF vs. the rest (SPAZO), with uops reading both inputs separately if they need them.) But for no apparent reason, setbe and seta are also 2 uops (https://uops.info) - maybe Intel never updated the setcc uop format to work as a 3-input uop (including the full register that they merge into the low byte of). Fun fact: this leads to setcc in general only being able to decode in the "complex" decoder, if your code isn't in the uop cache: Can the simple decoders in recent Intel microarchitectures handle all 1-µop instructions?

The loop body is 7 uops total on Intel thanks to cmov. 6 on AMD.

Compare vs. 4 uops for a simple scalar search loop without this trick. This can run as fast as 1 cycle per iteration on Intel (Haswell and later can run 2 branches per clock as long as at most 1 is taken). Also on AMD Zen I think. So we're searching about 8 bytes per cycle, about half what we could do with SIMD. But the startup and end overhead is lower.

L(loop):                 # do {
    mov   (%rdi), %rax       # 1 uop
    test  %rax,%rax
    jnz   L(found)           # 1 uop (macro-fused)
    add   $8, %rdi           # 1 uop
    cmp   %rdi, %rdx
    jbe   L(loop)         # }while(p < endp)   # 1 uops (macro-fused)
L(done):

---

If you counted a negative array index up towards zero, you could avoid having both add and sub in the loop: use FLAGS set by add. (Or for the simple version, avoiding the CMP, letting

You'd need an or (%r10, %r11, 8), %rax or something like that, but Haswell and later can keep that indexed addressing mode micro-fused as part of a 2-operand instruction with a RW destination: Micro fusion and addressing modes)

---

setcc r/m8 is inconvenient because it only works on an 8-bit register. But if you had a zeroed register like RDX you could setbe dl / shlx %rsi, %rdx, %rcx. ORing that into RAX before a test / jz keep_searching seems unlikely to be worth it even if you don't want to use SIMD.

A correctly predicted not-taken test/jnz is a single uop, cheaper than all the work you're doing to create an RAX value.