From 1121c23198901cf7c9f9b615f14045493a346e20 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Nathan=20Gau=C3=ABr?= <brioche@google.com>
Date: Tue, 5 Sep 2023 15:47:46 +0200
Subject: [PATCH] opt: add Int64 capability to trim pass (#5398)

Adds support for Int64 capability trimming.
---
 source/opt/trim_capabilities_pass.cpp    | 16 ++++++++--
 source/opt/trim_capabilities_pass.h      |  1 +
 test/opt/trim_capabilities_pass_test.cpp | 39 ++++++++++++++++++++++++
 3 files changed, 54 insertions(+), 2 deletions(-)

diff --git a/source/opt/trim_capabilities_pass.cpp b/source/opt/trim_capabilities_pass.cpp
index f8e4d8187..6c6b4a099 100644
--- a/source/opt/trim_capabilities_pass.cpp
+++ b/source/opt/trim_capabilities_pass.cpp
@@ -40,6 +40,7 @@ constexpr uint32_t kOpTypePointerStorageClassIndex = 0;
 constexpr uint32_t kTypeArrayTypeIndex = 0;
 constexpr uint32_t kOpTypeScalarBitWidthIndex = 0;
 constexpr uint32_t kTypePointerTypeIdInIdx = 1;
+constexpr uint32_t kOpTypeIntSizeIndex = 0;
 
 // DFS visit of the type defined by `instruction`.
 // If `condition` is true, children of the current node are visited.
@@ -255,13 +256,24 @@ static std::optional<spv::Capability> Handler_OpTypePointer_StorageUniform16(
                         : std::nullopt;
 }
 
+static std::optional<spv::Capability> Handler_OpTypeInt_Int64(
+    const Instruction* instruction) {
+  assert(instruction->opcode() == spv::Op::OpTypeInt &&
+         "This handler only support OpTypeInt opcodes.");
+
+  const uint32_t size =
+      instruction->GetSingleWordInOperand(kOpTypeIntSizeIndex);
+  return size == 64 ? std::optional(spv::Capability::Int64) : std::nullopt;
+}
+
 // Opcode of interest to determine capabilities requirements.
-constexpr std::array<std::pair<spv::Op, OpcodeHandler>, 4> kOpcodeHandlers{{
+constexpr std::array<std::pair<spv::Op, OpcodeHandler>, 5> kOpcodeHandlers{{
     // clang-format off
     {spv::Op::OpTypePointer, Handler_OpTypePointer_StorageInputOutput16},
     {spv::Op::OpTypePointer, Handler_OpTypePointer_StoragePushConstant16},
     {spv::Op::OpTypePointer, Handler_OpTypePointer_StorageUniformBufferBlock16},
-    {spv::Op::OpTypePointer, Handler_OpTypePointer_StorageUniform16}
+    {spv::Op::OpTypePointer, Handler_OpTypePointer_StorageUniform16},
+    {spv::Op::OpTypeInt, Handler_OpTypeInt_Int64 },
     // clang-format on
 }};
 
diff --git a/source/opt/trim_capabilities_pass.h b/source/opt/trim_capabilities_pass.h
index 16b5781ae..272861799 100644
--- a/source/opt/trim_capabilities_pass.h
+++ b/source/opt/trim_capabilities_pass.h
@@ -78,6 +78,7 @@ class TrimCapabilitiesPass : public Pass {
       spv::Capability::FragmentShaderSampleInterlockEXT,
       spv::Capability::FragmentShaderShadingRateInterlockEXT,
       spv::Capability::Groups,
+      spv::Capability::Int64,
       spv::Capability::Linkage,
       spv::Capability::MinLod,
       spv::Capability::Shader,
diff --git a/test/opt/trim_capabilities_pass_test.cpp b/test/opt/trim_capabilities_pass_test.cpp
index 90954439f..fb64ecd3a 100644
--- a/test/opt/trim_capabilities_pass_test.cpp
+++ b/test/opt/trim_capabilities_pass_test.cpp
@@ -1591,6 +1591,45 @@ TEST_F(TrimCapabilitiesPassTest,
   EXPECT_EQ(std::get<1>(result), Pass::Status::SuccessWithChange);
 }
 
+TEST_F(TrimCapabilitiesPassTest, Int64_RemovedWhenUnused) {
+  const std::string kTest = R"(
+               OpCapability Int64
+; CHECK-NOT:   OpCapability Int64
+               OpCapability Shader
+               OpMemoryModel Logical GLSL450
+               OpEntryPoint GLCompute %1 "main"
+       %void = OpTypeVoid
+          %3 = OpTypeFunction %void
+          %1 = OpFunction %void None %3
+          %6 = OpLabel
+               OpReturn
+               OpFunctionEnd;
+  )";
+  const auto result =
+      SinglePassRunAndMatch<TrimCapabilitiesPass>(kTest, /* skip_nop= */ false);
+  EXPECT_EQ(std::get<1>(result), Pass::Status::SuccessWithChange);
+}
+
+TEST_F(TrimCapabilitiesPassTest, Int64_RemainsWhenUsed) {
+  const std::string kTest = R"(
+               OpCapability Int64
+; CHECK:       OpCapability Int64
+               OpCapability Shader
+               OpMemoryModel Logical GLSL450
+               OpEntryPoint GLCompute %1 "main"
+       %void = OpTypeVoid
+        %int = OpTypeInt 64 0
+          %3 = OpTypeFunction %void
+          %1 = OpFunction %void None %3
+          %6 = OpLabel
+               OpReturn
+               OpFunctionEnd;
+  )";
+  const auto result =
+      SinglePassRunAndMatch<TrimCapabilitiesPass>(kTest, /* skip_nop= */ false);
+  EXPECT_EQ(std::get<1>(result), Pass::Status::SuccessWithoutChange);
+}
+
 }  // namespace
 }  // namespace opt
 }  // namespace spvtools