nir/spirv: initial handling of OpenCL.std extension opcodes

Not complete, mostly just adding things as I encounter them in CTS. But
not getting far enough yet to hit most of the OpenCL.std instructions.

Anyway, this is better than nothing and covers the most common builtins.

v2: add hadd proof from Jason
    move some of the lowering into opt_algebraic and create new nir opcodes
    simplify nextafter lowering
    fix normalize lowering for inf
    rework upsample to use nir_pack_bits
    add missing files to build systems
v3: split lines of iadd/sub_sat expressions

Signed-off-by: Karol Herbst <[email protected]>
Reviewed-by: Jason Ekstrand <[email protected]>

6 files changed, 304 insertions, 2 deletions

diff --git a/src/compiler/nir/meson.build b/src/compiler/nir/meson.build
index d7f88f391f4..af781f9a2d5 100644
--- a/src/compiler/nir/meson.build
+++ b/src/compiler/nir/meson.build
@@ -203,6 +203,7 @@ files_libnir = files(
   '../spirv/GLSL.std.450.h',
   '../spirv/gl_spirv.c',
   '../spirv/nir_spirv.h',
+  '../spirv/OpenCL.std.h',
   '../spirv/spirv.h',
   '../spirv/spirv_info.h',
   '../spirv/spirv_to_nir.c',
@@ -210,6 +211,7 @@ files_libnir = files(
   '../spirv/vtn_amd.c',
   '../spirv/vtn_cfg.c',
   '../spirv/vtn_glsl450.c',
+  '../spirv/vtn_opencl.c',
   '../spirv/vtn_private.h',
   '../spirv/vtn_subgroup.c',
   '../spirv/vtn_variables.c',
diff --git a/src/compiler/nir/nir.h b/src/compiler/nir/nir.h
index 401635712f9..a6690682c83 100644
--- a/src/compiler/nir/nir.h
+++ b/src/compiler/nir/nir.h
@@ -2252,6 +2252,9 @@ typedef struct nir_shader_compiler_options {
    /* Set if nir_lower_wpos_ytransform() should also invert gl_PointCoord. */
    bool lower_wpos_pntc;
 
+   bool lower_hadd;
+   bool lower_add_sat;
+
    /**
     * Should nir_lower_io() create load_interpolated_input intrinsics?
     *
diff --git a/src/compiler/nir/nir_builtin_builder.c b/src/compiler/nir/nir_builtin_builder.c
index 252a7691f36..a050306765e 100644
--- a/src/compiler/nir/nir_builtin_builder.c
+++ b/src/compiler/nir/nir_builtin_builder.c
@@ -21,11 +21,13 @@
  * IN THE SOFTWARE.
  */
 
+#include <math.h>
+
 #include "nir.h"
 #include "nir_builtin_builder.h"
 
 nir_ssa_def*
-nir_cross(nir_builder *b, nir_ssa_def *x, nir_ssa_def *y)
+nir_cross3(nir_builder *b, nir_ssa_def *x, nir_ssa_def *y)
 {
    unsigned yzx[3] = { 1, 2, 0 };
    unsigned zxy[3] = { 2, 0, 1 };
@@ -37,6 +39,33 @@ nir_cross(nir_builder *b, nir_ssa_def *x, nir_ssa_def *y)
 }
 
 nir_ssa_def*
+nir_cross4(nir_builder *b, nir_ssa_def *x, nir_ssa_def *y)
+{
+   nir_ssa_def *cross = nir_cross3(b, x, y);
+
+   return nir_vec4(b,
+      nir_channel(b, cross, 0),
+      nir_channel(b, cross, 1),
+      nir_channel(b, cross, 2),
+      nir_imm_intN_t(b, 0, cross->bit_size));
+}
+
+nir_ssa_def*
+nir_length(nir_builder *b, nir_ssa_def *vec)
+{
+   nir_ssa_def *finf = nir_imm_floatN_t(b, INFINITY, vec->bit_size);
+
+   nir_ssa_def *abs = nir_fabs(b, vec);
+   if (vec->num_components == 1)
+      return abs;
+
+   nir_ssa_def *maxc = nir_fmax_abs_vec_comp(b, abs);
+   abs = nir_fdiv(b, abs, maxc);
+   nir_ssa_def *res = nir_fmul(b, nir_fsqrt(b, nir_fdot(b, abs, abs)), maxc);
+   return nir_bcsel(b, nir_feq(b, maxc, finf), maxc, res);
+}
+
+nir_ssa_def*
 nir_fast_length(nir_builder *b, nir_ssa_def *vec)
 {
    switch (vec->num_components) {
@@ -50,6 +79,72 @@ nir_fast_length(nir_builder *b, nir_ssa_def *vec)
 }
 
 nir_ssa_def*
+nir_nextafter(nir_builder *b, nir_ssa_def *x, nir_ssa_def *y)
+{
+   nir_ssa_def *zero = nir_imm_intN_t(b, 0, x->bit_size);
+   nir_ssa_def *one = nir_imm_intN_t(b, 1, x->bit_size);
+
+   nir_ssa_def *condeq = nir_feq(b, x, y);
+   nir_ssa_def *conddir = nir_flt(b, x, y);
+   nir_ssa_def *condzero = nir_feq(b, x, zero);
+
+   /* beware of: +/-0.0 - 1 == NaN */
+   nir_ssa_def *xn =
+      nir_bcsel(b,
+                condzero,
+                nir_imm_intN_t(b, (1 << (x->bit_size - 1)) + 1, x->bit_size),
+                nir_isub(b, x, one));
+
+   /* beware of -0.0 + 1 == -0x1p-149 */
+   nir_ssa_def *xp = nir_bcsel(b, condzero, one, nir_iadd(b, x, one));
+
+   /* nextafter can be implemented by just +/- 1 on the int value */
+   nir_ssa_def *res =
+      nir_bcsel(b, nir_ixor(b, conddir, nir_flt(b, x, zero)), xp, xn);
+
+   return nir_nan_check2(b, x, y, nir_bcsel(b, condeq, x, res));
+}
+
+nir_ssa_def*
+nir_normalize(nir_builder *b, nir_ssa_def *vec)
+{
+   if (vec->num_components == 1)
+      return nir_fsign(b, vec);
+
+   nir_ssa_def *f0 = nir_imm_floatN_t(b, 0.0, vec->bit_size);
+   nir_ssa_def *f1 = nir_imm_floatN_t(b, 1.0, vec->bit_size);
+   nir_ssa_def *finf = nir_imm_floatN_t(b, INFINITY, vec->bit_size);
+
+   /* scale the input to increase precision */
+   nir_ssa_def *maxc = nir_fmax_abs_vec_comp(b, vec);
+   nir_ssa_def *svec = nir_fdiv(b, vec, maxc);
+   /* for inf */
+   nir_ssa_def *finfvec = nir_copysign(b, nir_bcsel(b, nir_feq(b, vec, finf), f1, f0), f1);
+
+   nir_ssa_def *temp = nir_bcsel(b, nir_feq(b, maxc, finf), finfvec, svec);
+   nir_ssa_def *res = nir_fmul(b, temp, nir_frsq(b, nir_fdot(b, temp, temp)));
+
+   return nir_bcsel(b, nir_feq(b, maxc, f0), vec, res);
+}
+
+nir_ssa_def*
+nir_rotate(nir_builder *b, nir_ssa_def *x, nir_ssa_def *y)
+{
+   nir_ssa_def *shift_mask = nir_imm_int(b, x->bit_size - 1);
+
+   if (y->bit_size != 32)
+      y = nir_u2u32(b, y);
+
+   nir_ssa_def *lshift = nir_iand(b, y, shift_mask);
+   nir_ssa_def *rshift = nir_isub(b, nir_imm_int(b, x->bit_size), lshift);
+
+   nir_ssa_def *hi = nir_ishl(b, x, lshift);
+   nir_ssa_def *lo = nir_ushr(b, x, rshift);
+
+   return nir_ior(b, hi, lo);
+}
+
+nir_ssa_def*
 nir_smoothstep(nir_builder *b, nir_ssa_def *edge0, nir_ssa_def *edge1, nir_ssa_def *x)
 {
    nir_ssa_def *f2 = nir_imm_floatN_t(b, 2.0, x->bit_size);
@@ -63,3 +158,18 @@ nir_smoothstep(nir_builder *b, nir_ssa_def *edge0, nir_ssa_def *edge1, nir_ssa_d
    /* result = t * t * (3 - 2 * t) */
    return nir_fmul(b, t, nir_fmul(b, t, nir_fsub(b, f3, nir_fmul(b, f2, t))));
 }
+
+nir_ssa_def*
+nir_upsample(nir_builder *b, nir_ssa_def *hi, nir_ssa_def *lo)
+{
+   assert(lo->num_components == hi->num_components);
+   assert(lo->bit_size == hi->bit_size);
+
+   nir_ssa_def *res[NIR_MAX_VEC_COMPONENTS];
+   for (unsigned i = 0; i < lo->num_components; ++i) {
+      nir_ssa_def *vec = nir_vec2(b, nir_channel(b, lo, i), nir_channel(b, hi, i));
+      res[i] = nir_pack_bits(b, vec, vec->bit_size * 2);
+   }
+
+   return nir_vec(b, res, lo->num_components);
+}
diff --git a/src/compiler/nir/nir_builtin_builder.h b/src/compiler/nir/nir_builtin_builder.h
index 792fa08e8c5..39a1e0bd3cd 100644
--- a/src/compiler/nir/nir_builtin_builder.h
+++ b/src/compiler/nir/nir_builtin_builder.h
@@ -31,10 +31,55 @@
  * Definitions for functions in the C file come first.
  */
 
-nir_ssa_def* nir_cross(nir_builder *b, nir_ssa_def *x, nir_ssa_def *y);
+nir_ssa_def* nir_cross3(nir_builder *b, nir_ssa_def *x, nir_ssa_def *y);
+nir_ssa_def* nir_cross4(nir_builder *b, nir_ssa_def *x, nir_ssa_def *y);
+nir_ssa_def* nir_length(nir_builder *b, nir_ssa_def *vec);
 nir_ssa_def* nir_fast_length(nir_builder *b, nir_ssa_def *vec);
+nir_ssa_def* nir_nextafter(nir_builder *b, nir_ssa_def *x, nir_ssa_def *y);
+nir_ssa_def* nir_normalize(nir_builder *b, nir_ssa_def *vec);
+nir_ssa_def* nir_rotate(nir_builder *b, nir_ssa_def *x, nir_ssa_def *y);
 nir_ssa_def* nir_smoothstep(nir_builder *b, nir_ssa_def *edge0,
                             nir_ssa_def *edge1, nir_ssa_def *x);
+nir_ssa_def* nir_upsample(nir_builder *b, nir_ssa_def *hi, nir_ssa_def *lo);
+
+static inline nir_ssa_def *
+nir_nan_check2(nir_builder *b, nir_ssa_def *x, nir_ssa_def *y, nir_ssa_def *res)
+{
+   return nir_bcsel(b, nir_fne(b, x, x), x, nir_bcsel(b, nir_fne(b, y, y), y, res));
+}
+
+static inline nir_ssa_def *
+nir_fmax_abs_vec_comp(nir_builder *b, nir_ssa_def *vec)
+{
+   nir_ssa_def *res = nir_channel(b, vec, 0);
+   for (unsigned i = 1; i < vec->num_components; ++i)
+      res = nir_fmax(b, res, nir_fabs(b, nir_channel(b, vec, i)));
+   return res;
+}
+
+static inline nir_ssa_def *
+nir_iabs_diff(nir_builder *b, nir_ssa_def *x, nir_ssa_def *y)
+{
+   nir_ssa_def *cond = nir_ige(b, x, y);
+   nir_ssa_def *res0 = nir_isub(b, x, y);
+   nir_ssa_def *res1 = nir_isub(b, y, x);
+   return nir_bcsel(b, cond, res0, res1);
+}
+
+static inline nir_ssa_def *
+nir_uabs_diff(nir_builder *b, nir_ssa_def *x, nir_ssa_def *y)
+{
+   nir_ssa_def *cond = nir_uge(b, x, y);
+   nir_ssa_def *res0 = nir_isub(b, x, y);
+   nir_ssa_def *res1 = nir_isub(b, y, x);
+   return nir_bcsel(b, cond, res0, res1);
+}
+
+static inline nir_ssa_def *
+nir_bitselect(nir_builder *b, nir_ssa_def *x, nir_ssa_def *y, nir_ssa_def *s)
+{
+   return nir_ior(b, nir_iand(b, nir_inot(b, s), x), nir_iand(b, s, y));
+}
 
 static inline nir_ssa_def *
 nir_fclamp(nir_builder *b,
@@ -58,12 +103,41 @@ nir_uclamp(nir_builder *b,
 }
 
 static inline nir_ssa_def *
+nir_copysign(nir_builder *b, nir_ssa_def *x, nir_ssa_def *y)
+{
+   uint64_t masks = 1ull << (x->bit_size - 1);
+   uint64_t maskv = ~masks;
+
+   nir_ssa_def *s = nir_imm_intN_t(b, masks, x->bit_size);
+   nir_ssa_def *v = nir_imm_intN_t(b, maskv, x->bit_size);
+
+   return nir_ior(b, nir_iand(b, x, v), nir_iand(b, y, s));
+}
+
+static inline nir_ssa_def *
 nir_degrees(nir_builder *b, nir_ssa_def *val)
 {
    return nir_fmul_imm(b, val, 180.0 / M_PI);
 }
 
 static inline nir_ssa_def *
+nir_fdim(nir_builder *b, nir_ssa_def *x, nir_ssa_def *y)
+{
+   nir_ssa_def *cond = nir_flt(b, y, x);
+   nir_ssa_def *res = nir_fsub(b, x, y);
+   nir_ssa_def *zero = nir_imm_floatN_t(b, 0.0, x->bit_size);
+
+   // return NaN if either x or y are NaN, else x-y if x>y, else +0.0
+   return nir_nan_check2(b, x, y, nir_bcsel(b, cond, res, zero));
+}
+
+static inline nir_ssa_def *
+nir_distance(nir_builder *b, nir_ssa_def *x, nir_ssa_def *y)
+{
+   return nir_length(b, nir_fsub(b, x, y));
+}
+
+static inline nir_ssa_def *
 nir_fast_distance(nir_builder *b, nir_ssa_def *x, nir_ssa_def *y)
 {
    return nir_fast_length(b, nir_fsub(b, x, y));
@@ -75,10 +149,64 @@ nir_fast_normalize(nir_builder *b, nir_ssa_def *vec)
    return nir_fdiv(b, vec, nir_fast_length(b, vec));
 }
 
+static inline nir_ssa_def*
+nir_fmad(nir_builder *b, nir_ssa_def *x, nir_ssa_def *y, nir_ssa_def *z)
+{
+   return nir_fadd(b, nir_fmul(b, x, y), z);
+}
+
+static inline nir_ssa_def*
+nir_maxmag(nir_builder *b, nir_ssa_def *x, nir_ssa_def *y)
+{
+   nir_ssa_def *xabs = nir_fabs(b, x);
+   nir_ssa_def *yabs = nir_fabs(b, y);
+
+   nir_ssa_def *condy = nir_flt(b, xabs, yabs);
+   nir_ssa_def *condx = nir_flt(b, yabs, xabs);
+
+   return nir_bcsel(b, condy, y, nir_bcsel(b, condx, x, nir_fmax(b, x, y)));
+}
+
+static inline nir_ssa_def*
+nir_minmag(nir_builder *b, nir_ssa_def *x, nir_ssa_def *y)
+{
+   nir_ssa_def *xabs = nir_fabs(b, x);
+   nir_ssa_def *yabs = nir_fabs(b, y);
+
+   nir_ssa_def *condx = nir_flt(b, xabs, yabs);
+   nir_ssa_def *condy = nir_flt(b, yabs, xabs);
+
+   return nir_bcsel(b, condy, y, nir_bcsel(b, condx, x, nir_fmin(b, x, y)));
+}
+
+static inline nir_ssa_def*
+nir_nan(nir_builder *b, nir_ssa_def *x)
+{
+   nir_ssa_def *nan = nir_imm_floatN_t(b, NAN, x->bit_size);
+   if (x->num_components == 1)
+      return nan;
+
+   nir_ssa_def *nans[NIR_MAX_VEC_COMPONENTS];
+   for (unsigned i = 0; i < x->num_components; ++i)
+      nans[i] = nan;
+
+   return nir_vec(b, nans, x->num_components);
+}
+
 static inline nir_ssa_def *
 nir_radians(nir_builder *b, nir_ssa_def *val)
 {
    return nir_fmul_imm(b, val, M_PI / 180.0);
 }
 
+static inline nir_ssa_def *
+nir_select(nir_builder *b, nir_ssa_def *x, nir_ssa_def *y, nir_ssa_def *s)
+{
+   if (s->num_components != 1) {
+      uint64_t mask = 1ull << (s->bit_size - 1);
+      s = nir_iand(b, s, nir_imm_intN_t(b, mask, s->bit_size));
+   }
+   return nir_bcsel(b, nir_ieq(b, s, nir_imm_intN_t(b, 0, s->bit_size)), x, y);
+}
+
 #endif /* NIR_BUILTIN_BUILDER_H */
diff --git a/src/compiler/nir/nir_opcodes.py b/src/compiler/nir/nir_opcodes.py
index 42f8662352e..5da3cafabb2 100644
--- a/src/compiler/nir/nir_opcodes.py
+++ b/src/compiler/nir/nir_opcodes.py
@@ -466,8 +466,20 @@ def binop_reduce(name, output_size, output_type, src_type, prereduce_expr,
 
 binop("fadd", tfloat, commutative + associative, "src0 + src1")
 binop("iadd", tint, commutative + associative, "src0 + src1")
+binop("iadd_sat", tint, commutative + associative, """
+      src1 > 0 ?
+         (src0 + src1 < src0 ? (1ull << (bit_size - 1)) - 1 : src0 + src1) :
+         (src0 < src0 + src1 ? (1ull << (bit_size - 1))     : src0 + src1)
+""")
 binop("uadd_sat", tuint, commutative,
       "(src0 + src1) < src0 ? UINT64_MAX : (src0 + src1)")
+binop("isub_sat", tint, "", """
+      src1 < 0 ?
+         (src0 - src1 < src0 ? (1ull << (bit_size - 1)) - 1 : src0 - src1) :
+         (src0 < src0 - src1 ? (1ull << (bit_size - 1))     : src0 - src1)
+""")
+binop("usub_sat", tuint, "", "src0 < src1 ? 0 : src0 - src1")
+
 binop("fsub", tfloat, "", "src0 - src1")
 binop("isub", tint, "", "src0 - src1")
 
@@ -536,6 +548,32 @@ binop_convert("uadd_carry", tuint, tuint, commutative, "src0 + src1 < src0")
 
 binop_convert("usub_borrow", tuint, tuint, "", "src0 < src1")
 
+# hadd: (a + b) >> 1 (without overflow)
+# x + y = x - (x & ~y) + (x & ~y) + y - (~x & y) + (~x & y)
+#       =      (x & y) + (x & ~y) +      (x & y) + (~x & y)
+#       = 2 *  (x & y) + (x & ~y) +                (~x & y)
+#       =     ((x & y) << 1) + (x ^ y)
+#
+# Since we know that the bottom bit of (x & y) << 1 is zero,
+#
+# (x + y) >> 1 = (((x & y) << 1) + (x ^ y)) >> 1
+#              =   (x & y) +      ((x ^ y)  >> 1)
+binop("ihadd", tint, commutative, "(src0 & src1) + ((src0 ^ src1) >> 1)")
+binop("uhadd", tuint, commutative, "(src0 & src1) + ((src0 ^ src1) >> 1)")
+
+# rhadd: (a + b + 1) >> 1 (without overflow)
+# x + y + 1 = x + (~x & y) - (~x & y) + y + (x & ~y) - (x & ~y) + 1
+#           =      (x | y) - (~x & y) +      (x | y) - (x & ~y) + 1
+#           = 2 *  (x | y) - ((~x & y) +               (x & ~y)) + 1
+#           =     ((x | y) << 1) - (x ^ y) + 1
+#
+# Since we know that the bottom bit of (x & y) << 1 is zero,
+#
+# (x + y + 1) >> 1 = (x | y) + (-(x ^ y) + 1) >> 1)
+#                  = (x | y) -  ((x ^ y)      >> 1)
+binop("irhadd", tint, commutative, "(src0 | src1) + ((src0 ^ src1) >> 1)")
+binop("urhadd", tuint, commutative, "(src0 | src1) + ((src0 ^ src1) >> 1)")
+
 binop("umod", tuint, "", "src1 == 0 ? 0 : src0 % src1")
 
 # For signed integers, there are several different possible definitions of
diff --git a/src/compiler/nir/nir_opt_algebraic.py b/src/compiler/nir/nir_opt_algebraic.py
index e8191b72271..5b2e7ee2405 100644
--- a/src/compiler/nir/nir_opt_algebraic.py
+++ b/src/compiler/nir/nir_opt_algebraic.py
@@ -678,6 +678,12 @@ optimizations = [
     ('bcsel', ('ilt', 31, 'bits'), 'insert',
               ('bfi', ('bfm', 'bits', 'offset'), 'insert', 'base')),
     'options->lower_bitfield_insert'),
+   (('ihadd', a, b), ('iadd', ('iand', a, b), ('ishr', ('ixor', a, b), 1)), 'options->lower_hadd'),
+   (('uhadd', a, b), ('iadd', ('iand', a, b), ('ushr', ('ixor', a, b), 1)), 'options->lower_hadd'),
+   (('irhadd', a, b), ('isub', ('ior', a, b), ('ishr', ('ixor', a, b), 1)), 'options->lower_hadd'),
+   (('urhadd', a, b), ('isub', ('ior', a, b), ('ushr', ('ixor', a, b), 1)), 'options->lower_hadd'),
+   (('uadd_sat', a, b), ('bcsel', ('ult', ('iadd', a, b), a), -1, ('iadd', a, b)), 'options->lower_add_sat'),
+   (('usub_sat', a, b), ('bcsel', ('ult', a, b), 0, ('isub', a, b)), 'options->lower_add_sat'),
 
    # Alternative lowering that doesn't rely on bfi.
    (('bitfield_insert', 'base', 'insert', 'offset', 'bits'),
@@ -794,6 +800,21 @@ optimizations = [
    (('isign', a), ('imin', ('imax', a, -1), 1), 'options->lower_isign'),
 ]
 
+# bit_size dependent lowerings
+for bit_size in [8, 16, 32, 64]:
+   # convenience constants
+   intmax = (1 << (bit_size - 1)) - 1
+   intmin = 1 << (bit_size - 1)
+
+   optimizations += [
+      (('iadd_sat@' + str(bit_size), a, b),
+       ('bcsel', ('ige', b, 1), ('bcsel', ('ilt', ('iadd', a, b), a), intmax, ('iadd', a, b)),
+                                ('bcsel', ('ilt', a, ('iadd', a, b)), intmin, ('iadd', a, b))), 'options->lower_add_sat'),
+      (('isub_sat@' + str(bit_size), a, b),
+       ('bcsel', ('ilt', b, 0), ('bcsel', ('ilt', ('isub', a, b), a), intmax, ('isub', a, b)),
+                                ('bcsel', ('ilt', a, ('isub', a, b)), intmin, ('isub', a, b))), 'options->lower_add_sat'),
+   ]
+
 invert = OrderedDict([('feq', 'fne'), ('fne', 'feq'), ('fge', 'flt'), ('flt', 'fge')])
 
 for left, right in itertools.combinations_with_replacement(invert.keys(), 2):