32-bit Integer Max/Min on Tenstorrent

Computing max/min of signed (two’s complement) or unsigned 32-bit integers is not directly supported on Tenstorrent’s AI accelerators. The available options are:

• Use SFPSWAP, but this operates on sign-magnitude integers (or floating-point values), and so will not give the correct result on its own.
• Use SFPIADD to compute a - b, treating inputs as two’s complement, and optionally set lane flags depending on the sign of the result. However, this suffers from overflow.

Binary Case

First, we’ll look at calculating max(a, b) and min(a, b), where a and b are tensors and processed elementwise.

32-bit Signed Integers

Looking more carefully at the SFPSWAP option, and using the following definition:

swap_minmax(x, y):
    # returns (min_sm(x, y), max_sm(x, y))
    # sign-magnitude ordering, with -0 < +0

Note that for max:

• swap_minmax(a, b) with the result in b gives the correct result if one or both values are non-negative;
• otherwise, the result will be inverted, i.e. it will be in a.

Similarly, for min:

• swap_minmax(a, b) with the result in a gives the correct result if one or both values are non-negative;
• otherwise, the result will be inverted, i.e. it will be in b.

In both cases, we simply have to check for the case where both values are negative, and invert the result.

Putting this together, we have:

a, b = swap_minmax(a, b)
if msb(a) == 1 and msb(b) == 1:
    a, b = b, a

result_min = a
result_max = b

32-bit Unsigned Integers

An elegant way to handle unsigned integers is to observe that if we treat an unsigned integer as a sign-magnitude integer, then values with MSB=1 become negative. If both values have MSB=1, then when treated as sign-magnitude, their order is reversed, i.e. in order to compute the maximum of two unsigned integers with MSB=1, we need to compute their minimum when treated as sign-magnitude (and vice versa).

If exactly one operand has MSB=1, this continues to work, as sign-magnitude ordering places the unsigned maximum in the min_sm position (as it’s negative in sign-magnitude).

If both values have MSB=0, then we have to invert the result.

a, b = swap_minmax(a, b)
if msb(a) == 0 and msb(b) == 0:
    a, b = b, a

result_max = a
result_min = b

Unary Case

Now consider max(a, k) and min(a, k), where a is a tensor, and k is a fixed scalar, potentially allowing us to precompute a transformed constant like ~k before processing a elementwise.

32-bit Signed Integers

• If msb(k) == 0 this is trivial: we can use SFPSWAP directly:

  sm_min, sm_max = swap_minmax(a, k)
  
  result_min = sm_min
  result_max = sm_max

• If msb(k) == 1, SFPSWAP will give the inverted answer if msb(a) == 1.

  One approach could be to compare-and-swap ~a and ~k instead. ~k flips its MSB to 0, so sign-magnitude treats it as non-negative. Applying the same transformation to a lets us use SFPSWAP on the transformed values.

  Note that we can’t use two’s complement negation, -k, since this would fail to map INT_MIN to the non-negative domain due to overflow (-INT_MIN = INT_MIN).

  not_k = ~k # can be precomputed and stored in constant register
  
  not_a = ~a
  sm_min, sm_max = swap_minmax(not_a, not_k)
  
  result_max = ~sm_min
  result_min = ~sm_max

32-bit Unsigned Integers

• If msb(k) == 1 this is trivial: we can use SFPSWAP directly. Since msb(k) == 1, k is negative in sign-magnitude, so the unsigned maximum ends up in min_sm.

  sm_min, sm_max = swap_minmax(a, k)
  
  result_max = sm_min
  result_min = sm_max

• If msb(k) == 0, SFPSWAP will give the inverted answer if msb(a) == 0.

  We use the same trick as for signed integers: compare-and-swap ~a and ~k.

  not_k = ~k # can be precomputed and stored in constant register
  
  not_a = ~a
  sm_min, sm_max = swap_minmax(not_a, not_k)
  
  result_min = ~sm_min
  result_max = ~sm_max

Future Work

Of course, the above can be optimised via SFPLOADMACRO. Watch this space!

Acknowledgements

Thanks to Tenstorrent for sponsoring this work.