builtins: Add countTrailingZeros

Requires polyfilling for all but the MSL backend.

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

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

src/tint/builtins.def

@@ -290,6 +290,8 @@ fn countLeadingZeros<T: iu32>(T) -> T
 fn countLeadingZeros<N: num, T: iu32>(vec<N, T>) -> vec<N, T>
 fn countOneBits<T: iu32>(T) -> T
 fn countOneBits<N: num, T: iu32>(vec<N, T>) -> vec<N, T>
+fn countTrailingZeros<T: iu32>(T) -> T
+fn countTrailingZeros<N: num, T: iu32>(vec<N, T>) -> vec<N, T>
 fn cross(vec3<f32>, vec3<f32>) -> vec3<f32>
 fn degrees(f32) -> f32
 fn degrees<N: num>(vec<N, f32>) -> vec<N, f32>

src/tint/sem/builtin_type.cc

@@ -72,6 +72,9 @@ BuiltinType ParseBuiltinType(const std::string& name) {
   if (name == "countOneBits") {
     return BuiltinType::kCountOneBits;
   }
+  if (name == "countTrailingZeros") {
+    return BuiltinType::kCountTrailingZeros;
+  }
   if (name == "cross") {
     return BuiltinType::kCross;
   }
@@ -374,6 +377,8 @@ const char* str(BuiltinType i) {
       return "countLeadingZeros";
     case BuiltinType::kCountOneBits:
       return "countOneBits";
+    case BuiltinType::kCountTrailingZeros:
+      return "countTrailingZeros";
     case BuiltinType::kCross:
       return "cross";
     case BuiltinType::kDegrees:

src/tint/sem/builtin_type.h

@@ -48,6 +48,7 @@ enum class BuiltinType {
   kCosh,
   kCountLeadingZeros,
   kCountOneBits,
+  kCountTrailingZeros,
   kCross,
   kDegrees,
   kDeterminant,

src/tint/transform/builtin_polyfill.cc

@@ -108,6 +108,80 @@ struct BuiltinPolyfill::State {
     return name;
   }
 
+  /// Builds the polyfill function for the `countTrailingZeros` builtin
+  /// @param ty the parameter and return type for the function
+  /// @return the polyfill function name
+  Symbol countTrailingZeros(const sem::Type* ty) {
+    auto name = b.Symbols().New("tint_count_trailing_zeros");
+    uint32_t width = 1;
+    if (auto* v = ty->As<sem::Vector>()) {
+      width = v->Width();
+    }
+
+    // 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* {
+      if (width == 1) {
+        return b.Expr(value);
+      }
+      return b.Construct(b.ty.vec<ProgramBuilder::u32>(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);
+    };
+    b.Func(
+        name, {b.Param("v", T(ty))}, T(ty),
+        {
+            // var x = U(v);
+            b.Decl(b.Var("x", nullptr, b.Construct(U(), b.Expr("v")))),
+            // let b16 = select(16, 0, bool(x & 0x0000ffff));
+            b.Decl(b.Const(
+                "b16", nullptr,
+                b.Call("select", V(16), V(0), B(b.And("x", V(0x0000ffff)))))),
+            // x = x >> b16;
+            b.Assign("x", b.Shr("x", "b16")),
+            // let b8  = select(8,  0, bool(x & 0x000000ff));
+            b.Decl(b.Const(
+                "b8", nullptr,
+                b.Call("select", V(8), V(0), B(b.And("x", V(0x000000ff)))))),
+            // x = x >> b8;
+            b.Assign("x", b.Shr("x", "b8")),
+            // let b4  = select(4,  0, bool(x & 0x0000000f));
+            b.Decl(b.Const(
+                "b4", nullptr,
+                b.Call("select", V(4), V(0), B(b.And("x", V(0x0000000f)))))),
+            // x = x >> b4;
+            b.Assign("x", b.Shr("x", "b4")),
+            // let b2  = select(2,  0, bool(x & 0x00000003));
+            b.Decl(b.Const(
+                "b2", nullptr,
+                b.Call("select", V(2), V(0), B(b.And("x", V(0x00000003)))))),
+            // x = x >> b2;
+            b.Assign("x", b.Shr("x", "b2")),
+            // let b1  = select(1,  0, bool(x & 0x00000001));
+            b.Decl(b.Const(
+                "b1", nullptr,
+                b.Call("select", V(1), V(0), B(b.And("x", V(0x00000001)))))),
+            // let is_zero  = select(0, 1, x == 0);
+            b.Decl(b.Const("is_zero", nullptr,
+                           b.Call("select", V(0), V(1), b.Equal("x", V(0))))),
+            // return R((b16 | b8 | b4 | b2 | b1) + zero);
+            b.Return(b.Construct(
+                T(ty),
+                b.Add(b.Or(b.Or(b.Or(b.Or("b16", "b8"), "b4"), "b2"), "b1"),
+                      "is_zero"))),
+        });
+    return name;
+  }
+
  private:
   const ast::Type* T(const sem::Type* ty) { return CreateASTTypeFor(ctx, ty); }
 };
@@ -130,6 +204,11 @@ bool BuiltinPolyfill::ShouldRun(const Program* program,
                 return true;
               }
               break;
+            case sem::BuiltinType::kCountTrailingZeros:
+              if (builtins.count_trailing_zeros) {
+                return true;
+              }
+              break;
             default:
               break;
           }
@@ -166,6 +245,13 @@ void BuiltinPolyfill::Run(CloneContext& ctx,
                   });
                 }
                 break;
+              case sem::BuiltinType::kCountTrailingZeros:
+                if (builtins.count_trailing_zeros) {
+                  polyfill = utils::GetOrCreate(polyfills, builtin, [&] {
+                    return s.countTrailingZeros(builtin->ReturnType());
+                  });
+                }
+                break;
               default:
                 break;
             }

src/tint/transform/builtin_polyfill.h

@@ -32,6 +32,8 @@ class BuiltinPolyfill : public Castable<BuiltinPolyfill, Transform> {
   struct Builtins {
     /// Should `countLeadingZeros()` be polyfilled?
     bool count_leading_zeros = false;
+    /// Should `countTrailingZeros()` be polyfilled?
+    bool count_trailing_zeros = false;
   };
 
   /// Config is consumed by the BuiltinPolyfill transform.

src/tint/transform/builtin_polyfill_test.cc

@@ -194,6 +194,160 @@ fn f() {
   EXPECT_EQ(expect, str(got));
 }
 
+////////////////////////////////////////////////////////////////////////////////
+// countTrailingZeros
+////////////////////////////////////////////////////////////////////////////////
+DataMap polyfillCountTrailingZeros() {
+  BuiltinPolyfill::Builtins builtins;
+  builtins.count_trailing_zeros = true;
+  DataMap data;
+  data.Add<BuiltinPolyfill::Config>(builtins);
+  return data;
+}
+
+TEST_F(BuiltinPolyfillTest, ShouldRunCountTrailingZeros) {
+  auto* src = R"(
+fn f() {
+  countTrailingZeros(0xf);
+}
+)";
+
+  EXPECT_FALSE(ShouldRun<BuiltinPolyfill>(src));
+  EXPECT_TRUE(ShouldRun<BuiltinPolyfill>(src, polyfillCountTrailingZeros()));
+}
+
+TEST_F(BuiltinPolyfillTest, CountTrailingZeros_i32) {
+  auto* src = R"(
+fn f() {
+  let r : i32 = countTrailingZeros(15);
+}
+)";
+
+  auto* expect = R"(
+fn tint_count_trailing_zeros(v : i32) -> i32 {
+  var x = u32(v);
+  let b16 = select(16u, 0u, bool((x & 65535u)));
+  x = (x >> b16);
+  let b8 = select(8u, 0u, bool((x & 255u)));
+  x = (x >> b8);
+  let b4 = select(4u, 0u, bool((x & 15u)));
+  x = (x >> b4);
+  let b2 = select(2u, 0u, bool((x & 3u)));
+  x = (x >> b2);
+  let b1 = select(1u, 0u, bool((x & 1u)));
+  let is_zero = select(0u, 1u, (x == 0u));
+  return i32((((((b16 | b8) | b4) | b2) | b1) + is_zero));
+}
+
+fn f() {
+  let r : i32 = tint_count_trailing_zeros(15);
+}
+)";
+
+  auto got = Run<BuiltinPolyfill>(src, polyfillCountTrailingZeros());
+
+  EXPECT_EQ(expect, str(got));
+}
+
+TEST_F(BuiltinPolyfillTest, CountTrailingZeros_u32) {
+  auto* src = R"(
+fn f() {
+  let r : u32 = countTrailingZeros(15u);
+}
+)";
+
+  auto* expect = R"(
+fn tint_count_trailing_zeros(v : u32) -> u32 {
+  var x = u32(v);
+  let b16 = select(16u, 0u, bool((x & 65535u)));
+  x = (x >> b16);
+  let b8 = select(8u, 0u, bool((x & 255u)));
+  x = (x >> b8);
+  let b4 = select(4u, 0u, bool((x & 15u)));
+  x = (x >> b4);
+  let b2 = select(2u, 0u, bool((x & 3u)));
+  x = (x >> b2);
+  let b1 = select(1u, 0u, bool((x & 1u)));
+  let is_zero = select(0u, 1u, (x == 0u));
+  return u32((((((b16 | b8) | b4) | b2) | b1) + is_zero));
+}
+
+fn f() {
+  let r : u32 = tint_count_trailing_zeros(15u);
+}
+)";
+
+  auto got = Run<BuiltinPolyfill>(src, polyfillCountTrailingZeros());
+
+  EXPECT_EQ(expect, str(got));
+}
+
+TEST_F(BuiltinPolyfillTest, CountTrailingZeros_vec3_i32) {
+  auto* src = R"(
+fn f() {
+  let r : vec3<i32> = countTrailingZeros(vec3<i32>(15));
+}
+)";
+
+  auto* expect = R"(
+fn tint_count_trailing_zeros(v : vec3<i32>) -> vec3<i32> {
+  var x = vec3<u32>(v);
+  let b16 = select(vec3<u32>(16u), vec3<u32>(0u), vec3<bool>((x & vec3<u32>(65535u))));
+  x = (x >> b16);
+  let b8 = select(vec3<u32>(8u), vec3<u32>(0u), vec3<bool>((x & vec3<u32>(255u))));
+  x = (x >> b8);
+  let b4 = select(vec3<u32>(4u), vec3<u32>(0u), vec3<bool>((x & vec3<u32>(15u))));
+  x = (x >> b4);
+  let b2 = select(vec3<u32>(2u), vec3<u32>(0u), vec3<bool>((x & vec3<u32>(3u))));
+  x = (x >> b2);
+  let b1 = select(vec3<u32>(1u), vec3<u32>(0u), vec3<bool>((x & vec3<u32>(1u))));
+  let is_zero = select(vec3<u32>(0u), vec3<u32>(1u), (x == vec3<u32>(0u)));
+  return vec3<i32>((((((b16 | b8) | b4) | b2) | b1) + is_zero));
+}
+
+fn f() {
+  let r : vec3<i32> = tint_count_trailing_zeros(vec3<i32>(15));
+}
+)";
+
+  auto got = Run<BuiltinPolyfill>(src, polyfillCountTrailingZeros());
+
+  EXPECT_EQ(expect, str(got));
+}
+
+TEST_F(BuiltinPolyfillTest, CountTrailingZeros_vec3_u32) {
+  auto* src = R"(
+fn f() {
+  let r : vec3<u32> = countTrailingZeros(vec3<u32>(15u));
+}
+)";
+
+  auto* expect = R"(
+fn tint_count_trailing_zeros(v : vec3<u32>) -> vec3<u32> {
+  var x = vec3<u32>(v);
+  let b16 = select(vec3<u32>(16u), vec3<u32>(0u), vec3<bool>((x & vec3<u32>(65535u))));
+  x = (x >> b16);
+  let b8 = select(vec3<u32>(8u), vec3<u32>(0u), vec3<bool>((x & vec3<u32>(255u))));
+  x = (x >> b8);
+  let b4 = select(vec3<u32>(4u), vec3<u32>(0u), vec3<bool>((x & vec3<u32>(15u))));
+  x = (x >> b4);
+  let b2 = select(vec3<u32>(2u), vec3<u32>(0u), vec3<bool>((x & vec3<u32>(3u))));
+  x = (x >> b2);
+  let b1 = select(vec3<u32>(1u), vec3<u32>(0u), vec3<bool>((x & vec3<u32>(1u))));
+  let is_zero = select(vec3<u32>(0u), vec3<u32>(1u), (x == vec3<u32>(0u)));
+  return vec3<u32>((((((b16 | b8) | b4) | b2) | b1) + is_zero));
+}
+
+fn f() {
+  let r : vec3<u32> = tint_count_trailing_zeros(vec3<u32>(15u));
+}
+)";
+
+  auto got = Run<BuiltinPolyfill>(src, polyfillCountTrailingZeros());
+
+  EXPECT_EQ(expect, str(got));
+}
+
 }  // namespace
 }  // namespace transform
 }  // namespace tint

src/tint/transform/glsl.cc

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

src/tint/writer/hlsl/generator_impl.cc

@@ -144,6 +144,7 @@ SanitizedResult Sanitize(
   {  // Builtin polyfills
     transform::BuiltinPolyfill::Builtins polyfills;
     polyfills.count_leading_zeros = true;
+    polyfills.count_trailing_zeros = true;
     data.Add<transform::BuiltinPolyfill::Config>(polyfills);
     manager.Add<transform::BuiltinPolyfill>();
   }

src/tint/writer/msl/generator_impl.cc

@@ -1351,6 +1351,9 @@ std::string GeneratorImpl::generate_builtin_name(const sem::Builtin* builtin) {
     case sem::BuiltinType::kCountOneBits:
       out += "popcount";
       break;
+    case sem::BuiltinType::kCountTrailingZeros:
+      out += "ctz";
+      break;
     case sem::BuiltinType::kDpdx:
     case sem::BuiltinType::kDpdxCoarse:
     case sem::BuiltinType::kDpdxFine:

src/tint/writer/msl/generator_impl_builtin_test.cc

@@ -126,6 +126,7 @@ const ast::CallExpression* GenerateCall(BuiltinType builtin,
       }
     case BuiltinType::kCountLeadingZeros:
     case BuiltinType::kCountOneBits:
+    case BuiltinType::kCountTrailingZeros:
     case BuiltinType::kReverseBits:
       return builder->Call(str.str(), "u2");
     case BuiltinType::kMax:
@@ -215,6 +216,7 @@ INSTANTIATE_TEST_SUITE_P(
         BuiltinData{BuiltinType::kCosh, ParamType::kF32, "cosh"},
         BuiltinData{BuiltinType::kCountLeadingZeros, ParamType::kU32, "clz"},
         BuiltinData{BuiltinType::kCountOneBits, ParamType::kU32, "popcount"},
+        BuiltinData{BuiltinType::kCountTrailingZeros, ParamType::kU32, "ctz"},
         BuiltinData{BuiltinType::kCross, ParamType::kF32, "cross"},
         BuiltinData{BuiltinType::kDeterminant, ParamType::kF32, "determinant"},
         BuiltinData{BuiltinType::kDistance, ParamType::kF32, "distance"},

src/tint/writer/spirv/builder.cc

@@ -262,6 +262,7 @@ SanitizedResult Sanitize(const Program* in,
   {  // Builtin polyfills
     transform::BuiltinPolyfill::Builtins polyfills;
     polyfills.count_leading_zeros = true;
+    polyfills.count_trailing_zeros = true;
     data.Add<transform::BuiltinPolyfill::Config>(polyfills);
     manager.Add<transform::BuiltinPolyfill>();
   }