Add geometry flags support (#1044)

Extend acceleration structure geometries (both AABBs and traingles) with support of flags.
google · Jun 25, 2024 · 4fa3746 · 4fa3746
1 parent 449b821
commit 4fa3746
Show file tree

Hide file tree

Showing 5 changed files with 115 additions and 1 deletion.
diff --git a/src/acceleration_structure.h b/src/acceleration_structure.h
@@ -49,6 +49,9 @@ class Geometry {
   void SetData(std::vector<float>& data) { data_.swap(data); }
   std::vector<float>& GetData() { return data_; }
 
+  void SetFlags(uint32_t flags) { flags_ = flags; }
+  uint32_t GetFlags() { return flags_; }
+
   size_t getVertexCount() const {
     return data_.size() / 3;  // Three floats to define vertex
   }
@@ -65,6 +68,7 @@ class Geometry {
  private:
   GeometryType type_ = GeometryType::kUnknown;
   std::vector<float> data_;
+  uint32_t flags_ = 0u;
 };
 
 class BLAS {

diff --git a/src/amberscript/parser.cc b/src/amberscript/parser.cc
@@ -4027,6 +4027,7 @@ Result Parser::ParseBLAS() {
 Result Parser::ParseBLASTriangle(BLAS* blas) {
   std::unique_ptr<Geometry> geometry = MakeUnique<Geometry>();
   std::vector<float> g;
+  uint32_t flags = 0;
   geometry->SetType(GeometryType::kTriangle);
 
   while (true) {
@@ -4041,6 +4042,10 @@ Result Parser::ParseBLASTriangle(BLAS* blas) {
       std::string tok = token->AsString();
       if (tok == "END") {
         break;
+      } else if (tok == "FLAGS") {
+        Result r = ParseGeometryFlags(&flags);
+        if (!r.IsSuccess())
+          return r;
       } else {
         return Result("END or float value is expected");
       }
@@ -4061,6 +4066,7 @@ Result Parser::ParseBLASTriangle(BLAS* blas) {
     return Result("Each triangle should include three vertices.");
 
   geometry->SetData(g);
+  geometry->SetFlags(flags);
 
   blas->AddGeometry(&geometry);
 
@@ -4070,6 +4076,7 @@ Result Parser::ParseBLASTriangle(BLAS* blas) {
 Result Parser::ParseBLASAABB(BLAS* blas) {
   std::unique_ptr<Geometry> geometry = MakeUnique<Geometry>();
   std::vector<float> g;
+  uint32_t flags = 0;
   geometry->SetType(GeometryType::kAABB);
 
   while (true) {
@@ -4084,6 +4091,10 @@ Result Parser::ParseBLASAABB(BLAS* blas) {
       std::string tok = token->AsString();
       if (tok == "END") {
         break;
+      } else if (tok == "FLAGS") {
+        Result r = ParseGeometryFlags(&flags);
+        if (!r.IsSuccess())
+          return r;
       } else {
         return Result("END or float value is expected");
       }
@@ -4105,12 +4116,54 @@ Result Parser::ParseBLASAABB(BLAS* blas) {
         "include two vertices.");
 
   geometry->SetData(g);
+  geometry->SetFlags(flags);
 
   blas->AddGeometry(&geometry);
 
   return {};
 }
 
+Result Parser::ParseGeometryFlags(uint32_t* flags) {
+  std::unique_ptr<Token> token;
+  bool first_eol = true;
+  bool singleline = true;
+  Result r;
+
+  while (true) {
+    token = tokenizer_->NextToken();
+    if (token->IsEOL()) {
+      if (first_eol) {
+        first_eol = false;
+        singleline = (*flags != 0);
+      }
+      if (singleline)
+        break;
+      else
+        continue;
+    }
+    if (token->IsEOS())
+      return Result("END command missing");
+
+    if (token->IsIdentifier()) {
+      if (token->AsString() == "END")
+        break;
+      else if (token->AsString() == "OPAQUE")
+        *flags |= VK_GEOMETRY_OPAQUE_BIT_KHR;
+      else if (token->AsString() == "NO_DUPLICATE_ANY_HIT")
+        *flags |= VK_GEOMETRY_NO_DUPLICATE_ANY_HIT_INVOCATION_BIT_KHR;
+      else
+        return Result("Unknown flag: " + token->AsString());
+    } else {
+      r = Result("Identifier expected");
+    }
+
+    if (!r.IsSuccess())
+      return r;
+  }
+
+  return {};
+}
+
 Result Parser::ParseTLAS() {
   auto token = tokenizer_->NextToken();
   if (!token->IsIdentifier())

diff --git a/src/amberscript/parser.h b/src/amberscript/parser.h
@@ -102,6 +102,7 @@ class Parser : public amber::Parser {
   Result ParseBLAS();
   Result ParseBLASTriangle(BLAS* blas);
   Result ParseBLASAABB(BLAS* blas);
+  Result ParseGeometryFlags(uint32_t* flags);
   Result ParseTLAS();
   Result ParseBLASInstance(TLAS* tlas);
   Result ParseBLASInstanceTransform(BLASInstance* instance);

diff --git a/src/amberscript/parser_raytracing_test.cc b/src/amberscript/parser_raytracing_test.cc
@@ -195,6 +195,33 @@ ACCELERATION_STRUCTURE BOTTOM_LEVEL blas_name
   EXPECT_EQ("3: Unexpected data type", r.Error());
 }
 
+TEST_F(AmberScriptParserTest, RayTracingBlasTriangleGeometryFlags) {
+  {
+    std::string in = R"(
+ACCELERATION_STRUCTURE BOTTOM_LEVEL blas_name
+  GEOMETRY TRIANGLES
+    FLAGS OPAQUE NO_DUPLICATE_ANY_HIT NO_SUCH_FLAG
+)";
+
+    Parser parser;
+    Result r = parser.Parse(in);
+    ASSERT_FALSE(r.IsSuccess());
+    EXPECT_EQ("4: Unknown flag: NO_SUCH_FLAG", r.Error());
+  }
+  {
+    std::string in = R"(
+ACCELERATION_STRUCTURE BOTTOM_LEVEL blas_name
+  GEOMETRY TRIANGLES
+    FLAGS 1
+)";
+
+    Parser parser;
+    Result r = parser.Parse(in);
+    ASSERT_FALSE(r.IsSuccess());
+    EXPECT_EQ("4: Identifier expected", r.Error());
+  }
+}
+
 TEST_F(AmberScriptParserTest, RayTracingBlasAABBEmpty) {
   std::string in = R"(
 ACCELERATION_STRUCTURE BOTTOM_LEVEL blas_name
@@ -249,6 +276,33 @@ ACCELERATION_STRUCTURE BOTTOM_LEVEL blas_name
   EXPECT_EQ("3: Unexpected data type", r.Error());
 }
 
+TEST_F(AmberScriptParserTest, RayTracingBlasAABBGeometryFlags) {
+  {
+    std::string in = R"(
+ACCELERATION_STRUCTURE BOTTOM_LEVEL blas_name
+  GEOMETRY AABBS
+    FLAGS OPAQUE NO_DUPLICATE_ANY_HIT NO_SUCH_FLAG
+)";
+
+    Parser parser;
+    Result r = parser.Parse(in);
+    ASSERT_FALSE(r.IsSuccess());
+    EXPECT_EQ("4: Unknown flag: NO_SUCH_FLAG", r.Error());
+  }
+  {
+    std::string in = R"(
+ACCELERATION_STRUCTURE BOTTOM_LEVEL blas_name
+  GEOMETRY AABBS
+    FLAGS 1
+)";
+
+    Parser parser;
+    Result r = parser.Parse(in);
+    ASSERT_FALSE(r.IsSuccess());
+    EXPECT_EQ("4: Identifier expected", r.Error());
+  }
+}
+
 TEST_F(AmberScriptParserTest, RayTracingTlasName) {
   std::string in = R"(
 ACCELERATION_STRUCTURE TOP_LEVEL

diff --git a/src/vulkan/blas.cc b/src/vulkan/blas.cc
@@ -91,7 +91,7 @@ Result BLAS::CreateBLAS(amber::BLAS* blas) {
             nullptr,
             geometryType,
             geometry,
-            0u,
+            VkGeometryFlagsKHR(geometryData->GetFlags())
         };
     const VkAccelerationStructureBuildRangeInfoKHR
         accelerationStructureBuildRangeInfosKHR = {
@@ -209,6 +209,8 @@ Result BLAS::CreateBLAS(amber::BLAS* blas) {
       } else {
         assert(false && "unknown geometry type");
       }
+      accelerationStructureGeometriesKHR_[geometryNdx].flags =
+          VkGeometryFlagsKHR(geometries[geometryNdx]->GetFlags());
     }
   }