builtins: Add insertBits

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

Bug: tint:1371
Change-Id: Idd55c0bc9dad1dffb558d0bc57d744f65e9041b5
Reviewed-on: https://dawn-review.googlesource.com/c/tint/+/81701
Kokoro: Kokoro <noreply+kokoro@google.com>
Reviewed-by: David Neto <dneto@google.com>

src/tint/transform/builtin_polyfill.cc

@@ -55,7 +55,7 @@ struct BuiltinPolyfill::State {
       if (width == 1) {
         return b.ty.u32();
       }
-      return b.ty.vec<ProgramBuilder::u32>(width);
+      return b.ty.vec<u32>(width);
     };
     auto V = [&](uint32_t value) -> const ast::Expression* {
       return ScalarOrVector(width, value);
@@ -117,7 +117,7 @@ struct BuiltinPolyfill::State {
       if (width == 1) {
         return b.ty.u32();
       }
-      return b.ty.vec<ProgramBuilder::u32>(width);
+      return b.ty.vec<u32>(width);
     };
     auto V = [&](uint32_t value) -> const ast::Expression* {
       return ScalarOrVector(width, value);
@@ -187,7 +187,7 @@ struct BuiltinPolyfill::State {
       if (width == 1) {
         return value;
       }
-      return b.Construct(b.ty.vec<ProgramBuilder::u32>(width), value);
+      return b.Construct(b.ty.vec<u32>(width), value);
     };
 
     ast::StatementList body = {
@@ -236,7 +236,7 @@ struct BuiltinPolyfill::State {
       if (width == 1) {
         return b.ty.u32();
       }
-      return b.ty.vec<ProgramBuilder::u32>(width);
+      return b.ty.vec<u32>(width);
     };
     auto V = [&](uint32_t value) -> const ast::Expression* {
       return ScalarOrVector(width, value);
@@ -317,7 +317,7 @@ struct BuiltinPolyfill::State {
       if (width == 1) {
         return b.ty.u32();
       }
-      return b.ty.vec<ProgramBuilder::u32>(width);
+      return b.ty.vec<u32>(width);
     };
     auto V = [&](uint32_t value) -> const ast::Expression* {
       return ScalarOrVector(width, value);
@@ -373,16 +373,90 @@ struct BuiltinPolyfill::State {
     return name;
   }
 
+  /// Builds the polyfill function for the `insertBits` builtin
+  /// @param ty the parameter and return type for the function
+  /// @return the polyfill function name
+  Symbol insertBits(const sem::Type* ty) {
+    auto name = b.Symbols().New("tint_insert_bits");
+    uint32_t width = WidthOf(ty);
+
+    constexpr uint32_t W = 32u;  // 32-bit
+
+    auto V = [&](auto value) -> const ast::Expression* {
+      const ast::Expression* expr = b.Expr(value);
+      if (!ty->is_unsigned_scalar_or_vector()) {
+        expr = b.Construct<i32>(expr);
+      }
+      if (ty->Is<sem::Vector>()) {
+        expr = b.Construct(T(ty), expr);
+      }
+      return expr;
+    };
+    auto U = [&](auto value) -> const ast::Expression* {
+      if (width == 1) {
+        return b.Expr(value);
+      }
+      return b.vec(b.ty.u32(), width, value);
+    };
+
+    ast::StatementList body = {
+        b.Decl(b.Const("s", nullptr, b.Call("min", "offset", W))),
+        b.Decl(b.Const("e", nullptr, b.Call("min", W, b.Add("s", "count")))),
+    };
+
+    switch (polyfill.insert_bits) {
+      case Level::kFull:
+        // let mask = ((1 << s) - 1) ^ ((1 << e) - 1)
+        body.emplace_back(b.Decl(b.Const(
+            "mask", nullptr,
+            b.Xor(b.Sub(b.Shl(1u, "s"), 1u), b.Sub(b.Shl(1u, "e"), 1u)))));
+        // return ((n << s) & mask) | (v & ~mask)
+        body.emplace_back(b.Return(b.Or(b.And(b.Shl("n", U("s")), V("mask")),
+                                        b.And("v", V(b.Complement("mask"))))));
+        break;
+      case Level::kClampParameters:
+        body.emplace_back(
+            b.Return(b.Call("insertBits", "v", "n", "s", b.Sub("e", "s"))));
+        break;
+      default:
+        TINT_ICE(Transform, b.Diagnostics())
+            << "unhandled polyfill level: "
+            << static_cast<int>(polyfill.insert_bits);
+        return {};
+    }
+
+    b.Func(name,
+           {
+               b.Param("v", T(ty)),
+               b.Param("n", T(ty)),
+               b.Param("offset", b.ty.u32()),
+               b.Param("count", b.ty.u32()),
+           },
+           T(ty), body);
+
+    return name;
+  }
+
  private:
+  /// Aliases
+  using u32 = ProgramBuilder::u32;
+  using i32 = ProgramBuilder::i32;
+
+  /// @returns the AST type for the given sem type
   const ast::Type* T(const sem::Type* ty) const {
     return CreateASTTypeFor(ctx, ty);
   }
+
+  /// @returns 1 if `ty` is not a vector, otherwise the vector width
   uint32_t WidthOf(const sem::Type* ty) const {
     if (auto* v = ty->As<sem::Vector>()) {
       return v->Width();
     }
     return 1;
   }
+
+  /// @returns a scalar or vector with the given width, with each element with
+  /// the given value.
   template <typename T>
   const ast::Expression* ScalarOrVector(uint32_t width, T value) const {
     if (width == 1) {
@@ -430,6 +504,11 @@ bool BuiltinPolyfill::ShouldRun(const Program* program,
                 return true;
               }
               break;
+            case sem::BuiltinType::kInsertBits:
+              if (builtins.insert_bits != Level::kNone) {
+                return true;
+              }
+              break;
             default:
               break;
           }
@@ -494,6 +573,13 @@ void BuiltinPolyfill::Run(CloneContext& ctx,
                   });
                 }
                 break;
+              case sem::BuiltinType::kInsertBits:
+                if (builtins.insert_bits != Level::kNone) {
+                  polyfill = utils::GetOrCreate(polyfills, builtin, [&] {
+                    return s.insertBits(builtin->ReturnType());
+                  });
+                }
+                break;
               default:
                 break;
             }

src/tint/transform/builtin_polyfill.h

@@ -50,6 +50,8 @@ class BuiltinPolyfill : public Castable<BuiltinPolyfill, Transform> {
     bool first_leading_bit = false;
     /// Should `firstTrailingBit()` be polyfilled?
     bool first_trailing_bit = false;
+    /// Should `insertBits()` be polyfilled?
+    Level insert_bits = Level::kNone;
   };
 
   /// Config is consumed by the BuiltinPolyfill transform.

src/tint/transform/builtin_polyfill_test.cc

@@ -889,6 +889,233 @@ fn f() {
   EXPECT_EQ(expect, str(got));
 }
 
+////////////////////////////////////////////////////////////////////////////////
+// insertBits
+////////////////////////////////////////////////////////////////////////////////
+DataMap polyfillInsertBits(Level level) {
+  BuiltinPolyfill::Builtins builtins;
+  builtins.insert_bits = level;
+  DataMap data;
+  data.Add<BuiltinPolyfill::Config>(builtins);
+  return data;
+}
+
+TEST_F(BuiltinPolyfillTest, ShouldRunInsertBits) {
+  auto* src = R"(
+fn f() {
+  insertBits(1234, 5678, 5u, 6u);
+}
+)";
+
+  EXPECT_FALSE(ShouldRun<BuiltinPolyfill>(src));
+  EXPECT_FALSE(
+      ShouldRun<BuiltinPolyfill>(src, polyfillInsertBits(Level::kNone)));
+  EXPECT_TRUE(ShouldRun<BuiltinPolyfill>(
+      src, polyfillInsertBits(Level::kClampParameters)));
+  EXPECT_TRUE(
+      ShouldRun<BuiltinPolyfill>(src, polyfillInsertBits(Level::kFull)));
+}
+
+TEST_F(BuiltinPolyfillTest, InsertBits_Full_i32) {
+  auto* src = R"(
+fn f() {
+  let r : i32 = insertBits(1234, 5678, 5u, 6u);
+}
+)";
+
+  auto* expect = R"(
+fn tint_insert_bits(v : i32, n : i32, offset : u32, count : u32) -> i32 {
+  let s = min(offset, 32u);
+  let e = min(32u, (s + count));
+  let mask = (((1u << s) - 1u) ^ ((1u << e) - 1u));
+  return (((n << s) & i32(mask)) | (v & i32(~(mask))));
+}
+
+fn f() {
+  let r : i32 = tint_insert_bits(1234, 5678, 5u, 6u);
+}
+)";
+
+  auto got = Run<BuiltinPolyfill>(src, polyfillInsertBits(Level::kFull));
+
+  EXPECT_EQ(expect, str(got));
+}
+
+TEST_F(BuiltinPolyfillTest, InsertBits_Full_u32) {
+  auto* src = R"(
+fn f() {
+  let r : u32 = insertBits(1234u, 5678u, 5u, 6u);
+}
+)";
+
+  auto* expect = R"(
+fn tint_insert_bits(v : u32, n : u32, offset : u32, count : u32) -> u32 {
+  let s = min(offset, 32u);
+  let e = min(32u, (s + count));
+  let mask = (((1u << s) - 1u) ^ ((1u << e) - 1u));
+  return (((n << s) & mask) | (v & ~(mask)));
+}
+
+fn f() {
+  let r : u32 = tint_insert_bits(1234u, 5678u, 5u, 6u);
+}
+)";
+
+  auto got = Run<BuiltinPolyfill>(src, polyfillInsertBits(Level::kFull));
+
+  EXPECT_EQ(expect, str(got));
+}
+
+TEST_F(BuiltinPolyfillTest, InsertBits_Full_vec3_i32) {
+  auto* src = R"(
+fn f() {
+  let r : vec3<i32> = insertBits(vec3<i32>(1234), vec3<i32>(5678), 5u, 6u);
+}
+)";
+
+  auto* expect = R"(
+fn tint_insert_bits(v : vec3<i32>, n : vec3<i32>, offset : u32, count : u32) -> vec3<i32> {
+  let s = min(offset, 32u);
+  let e = min(32u, (s + count));
+  let mask = (((1u << s) - 1u) ^ ((1u << e) - 1u));
+  return (((n << vec3<u32>(s)) & vec3<i32>(i32(mask))) | (v & vec3<i32>(i32(~(mask)))));
+}
+
+fn f() {
+  let r : vec3<i32> = tint_insert_bits(vec3<i32>(1234), vec3<i32>(5678), 5u, 6u);
+}
+)";
+
+  auto got = Run<BuiltinPolyfill>(src, polyfillInsertBits(Level::kFull));
+
+  EXPECT_EQ(expect, str(got));
+}
+
+TEST_F(BuiltinPolyfillTest, InsertBits_Full_vec3_u32) {
+  auto* src = R"(
+fn f() {
+  let r : vec3<u32> = insertBits(vec3<u32>(1234u), vec3<u32>(5678u), 5u, 6u);
+}
+)";
+
+  auto* expect = R"(
+fn tint_insert_bits(v : vec3<u32>, n : vec3<u32>, offset : u32, count : u32) -> vec3<u32> {
+  let s = min(offset, 32u);
+  let e = min(32u, (s + count));
+  let mask = (((1u << s) - 1u) ^ ((1u << e) - 1u));
+  return (((n << vec3<u32>(s)) & vec3<u32>(mask)) | (v & vec3<u32>(~(mask))));
+}
+
+fn f() {
+  let r : vec3<u32> = tint_insert_bits(vec3<u32>(1234u), vec3<u32>(5678u), 5u, 6u);
+}
+)";
+
+  auto got = Run<BuiltinPolyfill>(src, polyfillInsertBits(Level::kFull));
+
+  EXPECT_EQ(expect, str(got));
+}
+
+TEST_F(BuiltinPolyfillTest, InsertBits_Clamp_i32) {
+  auto* src = R"(
+fn f() {
+  let r : i32 = insertBits(1234, 5678, 5u, 6u);
+}
+)";
+
+  auto* expect = R"(
+fn tint_insert_bits(v : i32, n : i32, offset : u32, count : u32) -> i32 {
+  let s = min(offset, 32u);
+  let e = min(32u, (s + count));
+  return insertBits(v, n, s, (e - s));
+}
+
+fn f() {
+  let r : i32 = tint_insert_bits(1234, 5678, 5u, 6u);
+}
+)";
+
+  auto got =
+      Run<BuiltinPolyfill>(src, polyfillInsertBits(Level::kClampParameters));
+
+  EXPECT_EQ(expect, str(got));
+}
+
+TEST_F(BuiltinPolyfillTest, InsertBits_Clamp_u32) {
+  auto* src = R"(
+fn f() {
+  let r : u32 = insertBits(1234u, 5678u, 5u, 6u);
+}
+)";
+
+  auto* expect = R"(
+fn tint_insert_bits(v : u32, n : u32, offset : u32, count : u32) -> u32 {
+  let s = min(offset, 32u);
+  let e = min(32u, (s + count));
+  return insertBits(v, n, s, (e - s));
+}
+
+fn f() {
+  let r : u32 = tint_insert_bits(1234u, 5678u, 5u, 6u);
+}
+)";
+
+  auto got =
+      Run<BuiltinPolyfill>(src, polyfillInsertBits(Level::kClampParameters));
+
+  EXPECT_EQ(expect, str(got));
+}
+
+TEST_F(BuiltinPolyfillTest, InsertBits_Clamp_vec3_i32) {
+  auto* src = R"(
+fn f() {
+  let r : vec3<i32> = insertBits(vec3<i32>(1234), vec3<i32>(5678), 5u, 6u);
+}
+)";
+
+  auto* expect = R"(
+fn tint_insert_bits(v : vec3<i32>, n : vec3<i32>, offset : u32, count : u32) -> vec3<i32> {
+  let s = min(offset, 32u);
+  let e = min(32u, (s + count));
+  return insertBits(v, n, s, (e - s));
+}
+
+fn f() {
+  let r : vec3<i32> = tint_insert_bits(vec3<i32>(1234), vec3<i32>(5678), 5u, 6u);
+}
+)";
+
+  auto got =
+      Run<BuiltinPolyfill>(src, polyfillInsertBits(Level::kClampParameters));
+
+  EXPECT_EQ(expect, str(got));
+}
+
+TEST_F(BuiltinPolyfillTest, InsertBits_Clamp_vec3_u32) {
+  auto* src = R"(
+fn f() {
+  let r : vec3<u32> = insertBits(vec3<u32>(1234u), vec3<u32>(5678u), 5u, 6u);
+}
+)";
+
+  auto* expect = R"(
+fn tint_insert_bits(v : vec3<u32>, n : vec3<u32>, offset : u32, count : u32) -> vec3<u32> {
+  let s = min(offset, 32u);
+  let e = min(32u, (s + count));
+  return insertBits(v, n, s, (e - s));
+}
+
+fn f() {
+  let r : vec3<u32> = tint_insert_bits(vec3<u32>(1234u), vec3<u32>(5678u), 5u, 6u);
+}
+)";
+
+  auto got =
+      Run<BuiltinPolyfill>(src, polyfillInsertBits(Level::kClampParameters));
+
+  EXPECT_EQ(expect, str(got));
+}
+
 }  // namespace
 }  // namespace transform
 }  // namespace tint

src/tint/transform/glsl.cc

@@ -59,6 +59,7 @@ Output Glsl::Run(const Program* in, const DataMap& inputs) const {
     polyfills.extract_bits = BuiltinPolyfill::Level::kClampParameters;
     polyfills.first_leading_bit = true;
     polyfills.first_trailing_bit = true;
+    polyfills.insert_bits = BuiltinPolyfill::Level::kClampParameters;
     data.Add<BuiltinPolyfill::Config>(polyfills);
     manager.Add<BuiltinPolyfill>();
   }