builtins: Add firstLeadingBit

Currently polyfilled for all backends.
HLSL should be able to map this to 'firstbithigh', but there might need
to be some special case handling for 0 (undocumented behavior). For now
just polyfill.

CTS tests: https://github.com/gpuweb/cts/pull/1004

Bug: tint:1367
Bug: tint:1449
Change-Id: I9c9a08ea93d1c4a602e0ab763e95e2eea336fb0d
Reviewed-on: https://dawn-review.googlesource.com/c/tint/+/81503
Reviewed-by: David Neto <dneto@google.com>
Kokoro: Kokoro <noreply+kokoro@google.com>
Commit-Queue: Ben Clayton <bclayton@google.com>

src/tint/transform/builtin_polyfill.cc

@@ -55,10 +55,7 @@ struct BuiltinPolyfill::State {
       return b.ty.vec<ProgramBuilder::u32>(width);
     };
     auto V = [&](uint32_t value) -> const ast::Expression* {
-      if (width == 1) {
-        return b.Expr(value);
-      }
-      return b.Construct(b.ty.vec<ProgramBuilder::u32>(width), value);
+      return ScalarOrVector(width, value);
     };
     b.Func(
         name, {b.Param("v", T(ty))}, T(ty),
@@ -120,10 +117,7 @@ struct BuiltinPolyfill::State {
       return b.ty.vec<ProgramBuilder::u32>(width);
     };
     auto V = [&](uint32_t value) -> const ast::Expression* {
-      if (width == 1) {
-        return b.Expr(value);
-      }
-      return b.Construct(b.ty.vec<ProgramBuilder::u32>(width), value);
+      return ScalarOrVector(width, value);
     };
     auto B = [&](const ast::Expression* value) -> const ast::Expression* {
       if (width == 1) {
@@ -176,6 +170,87 @@ struct BuiltinPolyfill::State {
     return name;
   }
 
+  /// Builds the polyfill function for the `firstLeadingBit` builtin
+  /// @param ty the parameter and return type for the function
+  /// @return the polyfill function name
+  Symbol firstLeadingBit(const sem::Type* ty) {
+    auto name = b.Symbols().New("tint_first_leading_bit");
+    uint32_t width = WidthOf(ty);
+
+    // Returns either u32 or vecN<u32>
+    auto U = [&]() -> const ast::Type* {
+      if (width == 1) {
+        return b.ty.u32();
+      }
+      return b.ty.vec<ProgramBuilder::u32>(width);
+    };
+    auto V = [&](uint32_t value) -> const ast::Expression* {
+      return ScalarOrVector(width, value);
+    };
+    auto B = [&](const ast::Expression* value) -> const ast::Expression* {
+      if (width == 1) {
+        return b.Construct<bool>(value);
+      }
+      return b.Construct(b.ty.vec<bool>(width), value);
+    };
+
+    const ast::Expression* x = nullptr;
+    if (ty->is_unsigned_scalar_or_vector()) {
+      x = b.Expr("v");
+    } else {
+      // If ty is signed, then the value is inverted if the sign is negative
+      x = b.Call("select",                             //
+                 b.Construct(U(), "v"),                //
+                 b.Construct(U(), b.Complement("v")),  //
+                 b.LessThan("v", ScalarOrVector(width, 0)));
+    }
+
+    b.Func(name, {b.Param("v", T(ty))}, T(ty),
+           {
+               // var x = v;                          (unsigned)
+               // var x = select(U(v), ~U(v), v < 0); (signed)
+               b.Decl(b.Var("x", nullptr, x)),
+               // let b16 = select(0, 16, bool(x & 0xffff0000));
+               b.Decl(b.Const("b16", nullptr,
+                              b.Call("select", V(0), V(16),
+                                     B(b.And("x", V(0xffff0000)))))),
+               // x = x >> b16;
+               b.Assign("x", b.Shr("x", "b16")),
+               // let b8  = select(0, 8,  bool(x & 0x0000ff00));
+               b.Decl(b.Const(
+                   "b8", nullptr,
+                   b.Call("select", V(0), V(8), B(b.And("x", V(0x0000ff00)))))),
+               // x = x >> b8;
+               b.Assign("x", b.Shr("x", "b8")),
+               // let b4  = select(0, 4,  bool(x & 0x000000f0));
+               b.Decl(b.Const(
+                   "b4", nullptr,
+                   b.Call("select", V(0), V(4), B(b.And("x", V(0x000000f0)))))),
+               // x = x >> b4;
+               b.Assign("x", b.Shr("x", "b4")),
+               // let b2  = select(0, 2,  bool(x & 0x0000000c));
+               b.Decl(b.Const(
+                   "b2", nullptr,
+                   b.Call("select", V(0), V(2), B(b.And("x", V(0x0000000c)))))),
+               // x = x >> b2;
+               b.Assign("x", b.Shr("x", "b2")),
+               // let b1  = select(0, 1,  bool(x & 0x00000002));
+               b.Decl(b.Const(
+                   "b1", nullptr,
+                   b.Call("select", V(0), V(1), B(b.And("x", V(0x00000002)))))),
+               // let is_zero  = select(0, 0xffffffff, x == 0);
+               b.Decl(b.Const("is_zero", nullptr,
+                              b.Call("select", V(0), V(0xffffffff),
+                                     b.Equal("x", V(0))))),
+               // return R(b16 | b8 | b4 | b2 | b1 | zero);
+               b.Return(b.Construct(
+                   T(ty),
+                   b.Or(b.Or(b.Or(b.Or(b.Or("b16", "b8"), "b4"), "b2"), "b1"),
+                        "is_zero"))),
+           });
+    return name;
+  }
+
   /// Builds the polyfill function for the `firstTrailingBit` builtin
   /// @param ty the parameter and return type for the function
   /// @return the polyfill function name
@@ -191,10 +266,7 @@ struct BuiltinPolyfill::State {
       return b.ty.vec<ProgramBuilder::u32>(width);
     };
     auto V = [&](uint32_t value) -> const ast::Expression* {
-      if (width == 1) {
-        return b.Expr(value);
-      }
-      return b.Construct(b.ty.vec<ProgramBuilder::u32>(width), value);
+      return ScalarOrVector(width, value);
     };
     auto B = [&](const ast::Expression* value) -> const ast::Expression* {
       if (width == 1) {
@@ -257,6 +329,13 @@ struct BuiltinPolyfill::State {
     }
     return 1;
   }
+  template <typename T>
+  const ast::Expression* ScalarOrVector(uint32_t width, T value) const {
+    if (width == 1) {
+      return b.Expr(value);
+    }
+    return b.Construct(b.ty.vec<T>(width), value);
+  }
 };
 
 BuiltinPolyfill::BuiltinPolyfill() = default;
@@ -282,6 +361,11 @@ bool BuiltinPolyfill::ShouldRun(const Program* program,
                 return true;
               }
               break;
+            case sem::BuiltinType::kFirstLeadingBit:
+              if (builtins.first_leading_bit) {
+                return true;
+              }
+              break;
             case sem::BuiltinType::kFirstTrailingBit:
               if (builtins.first_trailing_bit) {
                 return true;
@@ -330,6 +414,13 @@ void BuiltinPolyfill::Run(CloneContext& ctx,
                   });
                 }
                 break;
+              case sem::BuiltinType::kFirstLeadingBit:
+                if (builtins.first_leading_bit) {
+                  polyfill = utils::GetOrCreate(polyfills, builtin, [&] {
+                    return s.firstLeadingBit(builtin->ReturnType());
+                  });
+                }
+                break;
               case sem::BuiltinType::kFirstTrailingBit:
                 if (builtins.first_trailing_bit) {
                   polyfill = utils::GetOrCreate(polyfills, builtin, [&] {

src/tint/transform/builtin_polyfill.h

@@ -34,6 +34,8 @@ class BuiltinPolyfill : public Castable<BuiltinPolyfill, Transform> {
     bool count_leading_zeros = false;
     /// Should `countTrailingZeros()` be polyfilled?
     bool count_trailing_zeros = false;
+    /// Should `firstLeadingBit()` be polyfilled?
+    bool first_leading_bit = false;
     /// Should `firstTrailingBit()` be polyfilled?
     bool first_trailing_bit = false;
   };

src/tint/transform/builtin_polyfill_test.cc

@@ -348,6 +348,160 @@ fn f() {
   EXPECT_EQ(expect, str(got));
 }
 
+////////////////////////////////////////////////////////////////////////////////
+// firstLeadingBit
+////////////////////////////////////////////////////////////////////////////////
+DataMap polyfillFirstLeadingBit() {
+  BuiltinPolyfill::Builtins builtins;
+  builtins.first_leading_bit = true;
+  DataMap data;
+  data.Add<BuiltinPolyfill::Config>(builtins);
+  return data;
+}
+
+TEST_F(BuiltinPolyfillTest, ShouldRunFirstLeadingBit) {
+  auto* src = R"(
+fn f() {
+  firstLeadingBit(0xf);
+}
+)";
+
+  EXPECT_FALSE(ShouldRun<BuiltinPolyfill>(src));
+  EXPECT_TRUE(ShouldRun<BuiltinPolyfill>(src, polyfillFirstLeadingBit()));
+}
+
+TEST_F(BuiltinPolyfillTest, FirstLeadingBit_i32) {
+  auto* src = R"(
+fn f() {
+  let r : i32 = firstLeadingBit(15);
+}
+)";
+
+  auto* expect = R"(
+fn tint_first_leading_bit(v : i32) -> i32 {
+  var x = select(u32(v), u32(~(v)), (v < 0));
+  let b16 = select(0u, 16u, bool((x & 4294901760u)));
+  x = (x >> b16);
+  let b8 = select(0u, 8u, bool((x & 65280u)));
+  x = (x >> b8);
+  let b4 = select(0u, 4u, bool((x & 240u)));
+  x = (x >> b4);
+  let b2 = select(0u, 2u, bool((x & 12u)));
+  x = (x >> b2);
+  let b1 = select(0u, 1u, bool((x & 2u)));
+  let is_zero = select(0u, 4294967295u, (x == 0u));
+  return i32((((((b16 | b8) | b4) | b2) | b1) | is_zero));
+}
+
+fn f() {
+  let r : i32 = tint_first_leading_bit(15);
+}
+)";
+
+  auto got = Run<BuiltinPolyfill>(src, polyfillFirstLeadingBit());
+
+  EXPECT_EQ(expect, str(got));
+}
+
+TEST_F(BuiltinPolyfillTest, FirstLeadingBit_u32) {
+  auto* src = R"(
+fn f() {
+  let r : u32 = firstLeadingBit(15u);
+}
+)";
+
+  auto* expect = R"(
+fn tint_first_leading_bit(v : u32) -> u32 {
+  var x = v;
+  let b16 = select(0u, 16u, bool((x & 4294901760u)));
+  x = (x >> b16);
+  let b8 = select(0u, 8u, bool((x & 65280u)));
+  x = (x >> b8);
+  let b4 = select(0u, 4u, bool((x & 240u)));
+  x = (x >> b4);
+  let b2 = select(0u, 2u, bool((x & 12u)));
+  x = (x >> b2);
+  let b1 = select(0u, 1u, bool((x & 2u)));
+  let is_zero = select(0u, 4294967295u, (x == 0u));
+  return u32((((((b16 | b8) | b4) | b2) | b1) | is_zero));
+}
+
+fn f() {
+  let r : u32 = tint_first_leading_bit(15u);
+}
+)";
+
+  auto got = Run<BuiltinPolyfill>(src, polyfillFirstLeadingBit());
+
+  EXPECT_EQ(expect, str(got));
+}
+
+TEST_F(BuiltinPolyfillTest, FirstLeadingBit_vec3_i32) {
+  auto* src = R"(
+fn f() {
+  let r : vec3<i32> = firstLeadingBit(vec3<i32>(15));
+}
+)";
+
+  auto* expect = R"(
+fn tint_first_leading_bit(v : vec3<i32>) -> vec3<i32> {
+  var x = select(vec3<u32>(v), vec3<u32>(~(v)), (v < vec3<i32>(0)));
+  let b16 = select(vec3<u32>(0u), vec3<u32>(16u), vec3<bool>((x & vec3<u32>(4294901760u))));
+  x = (x >> b16);
+  let b8 = select(vec3<u32>(0u), vec3<u32>(8u), vec3<bool>((x & vec3<u32>(65280u))));
+  x = (x >> b8);
+  let b4 = select(vec3<u32>(0u), vec3<u32>(4u), vec3<bool>((x & vec3<u32>(240u))));
+  x = (x >> b4);
+  let b2 = select(vec3<u32>(0u), vec3<u32>(2u), vec3<bool>((x & vec3<u32>(12u))));
+  x = (x >> b2);
+  let b1 = select(vec3<u32>(0u), vec3<u32>(1u), vec3<bool>((x & vec3<u32>(2u))));
+  let is_zero = select(vec3<u32>(0u), vec3<u32>(4294967295u), (x == vec3<u32>(0u)));
+  return vec3<i32>((((((b16 | b8) | b4) | b2) | b1) | is_zero));
+}
+
+fn f() {
+  let r : vec3<i32> = tint_first_leading_bit(vec3<i32>(15));
+}
+)";
+
+  auto got = Run<BuiltinPolyfill>(src, polyfillFirstLeadingBit());
+
+  EXPECT_EQ(expect, str(got));
+}
+
+TEST_F(BuiltinPolyfillTest, FirstLeadingBit_vec3_u32) {
+  auto* src = R"(
+fn f() {
+  let r : vec3<u32> = firstLeadingBit(vec3<u32>(15u));
+}
+)";
+
+  auto* expect = R"(
+fn tint_first_leading_bit(v : vec3<u32>) -> vec3<u32> {
+  var x = v;
+  let b16 = select(vec3<u32>(0u), vec3<u32>(16u), vec3<bool>((x & vec3<u32>(4294901760u))));
+  x = (x >> b16);
+  let b8 = select(vec3<u32>(0u), vec3<u32>(8u), vec3<bool>((x & vec3<u32>(65280u))));
+  x = (x >> b8);
+  let b4 = select(vec3<u32>(0u), vec3<u32>(4u), vec3<bool>((x & vec3<u32>(240u))));
+  x = (x >> b4);
+  let b2 = select(vec3<u32>(0u), vec3<u32>(2u), vec3<bool>((x & vec3<u32>(12u))));
+  x = (x >> b2);
+  let b1 = select(vec3<u32>(0u), vec3<u32>(1u), vec3<bool>((x & vec3<u32>(2u))));
+  let is_zero = select(vec3<u32>(0u), vec3<u32>(4294967295u), (x == vec3<u32>(0u)));
+  return vec3<u32>((((((b16 | b8) | b4) | b2) | b1) | is_zero));
+}
+
+fn f() {
+  let r : vec3<u32> = tint_first_leading_bit(vec3<u32>(15u));
+}
+)";
+
+  auto got = Run<BuiltinPolyfill>(src, polyfillFirstLeadingBit());
+
+  EXPECT_EQ(expect, str(got));
+}
+
 ////////////////////////////////////////////////////////////////////////////////
 // firstTrailingBit
 ////////////////////////////////////////////////////////////////////////////////

src/tint/transform/glsl.cc

@@ -56,6 +56,7 @@ Output Glsl::Run(const Program* in, const DataMap& inputs) const {
     BuiltinPolyfill::Builtins polyfills;
     polyfills.count_leading_zeros = true;
     polyfills.count_trailing_zeros = true;
+    polyfills.first_leading_bit = true;
     polyfills.first_trailing_bit = true;
     data.Add<BuiltinPolyfill::Config>(polyfills);
     manager.Add<BuiltinPolyfill>();