Add knowledge of cooperative matrices

Some optimizations are not aware of cooperative matrices, and either do nothing or assert. This commits fixes that up.
s-perron · Jun 21, 2024 · 86d31ab · 86d31ab
1 parent 0cfe9e7
commit 86d31ab
Show file tree

Hide file tree

Showing 6 changed files with 49 additions and 5 deletions.
diff --git a/source/opt/aggressive_dead_code_elim_pass.cpp b/source/opt/aggressive_dead_code_elim_pass.cpp
@@ -1004,7 +1004,8 @@ void AggressiveDCEPass::InitExtensions() {
       "SPV_NV_bindless_texture",
       "SPV_EXT_shader_atomic_float_add",
       "SPV_EXT_fragment_shader_interlock",
-      "SPV_NV_compute_shader_derivatives"
+      "SPV_NV_compute_shader_derivatives",
+      "SPV_KHR_cooperative_matrix"
   });
   // clang-format on
 }

diff --git a/source/opt/folding_rules.cpp b/source/opt/folding_rules.cpp
@@ -394,6 +394,11 @@ FoldingRule MergeNegateMulDivArithmetic() {
     analysis::ConstantManager* const_mgr = context->get_constant_mgr();
     const analysis::Type* type =
         context->get_type_mgr()->GetType(inst->type_id());
+
+    if (type->kind() == analysis::Type::kCooperativeMatrixKHR) {
+      return false;
+    }
+
     if (HasFloatingPoint(type) && !inst->IsFloatingPointFoldingAllowed())
       return false;
 
@@ -455,6 +460,11 @@ FoldingRule MergeNegateAddSubArithmetic() {
     analysis::ConstantManager* const_mgr = context->get_constant_mgr();
     const analysis::Type* type =
         context->get_type_mgr()->GetType(inst->type_id());
+
+    if (type->kind() == analysis::Type::kCooperativeMatrixKHR) {
+      return false;
+    }
+
     if (HasFloatingPoint(type) && !inst->IsFloatingPointFoldingAllowed())
       return false;
 
@@ -1068,6 +1078,11 @@ FoldingRule MergeSubNegateArithmetic() {
     analysis::ConstantManager* const_mgr = context->get_constant_mgr();
     const analysis::Type* type =
         context->get_type_mgr()->GetType(inst->type_id());
+
+    if (type->kind() == analysis::Type::kCooperativeMatrixKHR) {
+      return false;
+    }
+
     bool uses_float = HasFloatingPoint(type);
     if (uses_float && !inst->IsFloatingPointFoldingAllowed()) return false;
 
@@ -1116,6 +1131,11 @@ FoldingRule MergeAddAddArithmetic() {
            inst->opcode() == spv::Op::OpIAdd);
     const analysis::Type* type =
         context->get_type_mgr()->GetType(inst->type_id());
+
+    if (type->kind() == analysis::Type::Kind::kCooperativeMatrixKHR) {
+      return false;
+    }
+
     analysis::ConstantManager* const_mgr = context->get_constant_mgr();
     bool uses_float = HasFloatingPoint(type);
     if (uses_float && !inst->IsFloatingPointFoldingAllowed()) return false;
@@ -1164,6 +1184,11 @@ FoldingRule MergeAddSubArithmetic() {
            inst->opcode() == spv::Op::OpIAdd);
     const analysis::Type* type =
         context->get_type_mgr()->GetType(inst->type_id());
+
+    if (type->kind() == analysis::Type::Kind::kCooperativeMatrixKHR) {
+      return false;
+    }
+
     analysis::ConstantManager* const_mgr = context->get_constant_mgr();
     bool uses_float = HasFloatingPoint(type);
     if (uses_float && !inst->IsFloatingPointFoldingAllowed()) return false;
@@ -1224,6 +1249,11 @@ FoldingRule MergeSubAddArithmetic() {
            inst->opcode() == spv::Op::OpISub);
     const analysis::Type* type =
         context->get_type_mgr()->GetType(inst->type_id());
+
+    if (type->kind() == analysis::Type::kCooperativeMatrixKHR) {
+      return false;
+    }
+
     analysis::ConstantManager* const_mgr = context->get_constant_mgr();
     bool uses_float = HasFloatingPoint(type);
     if (uses_float && !inst->IsFloatingPointFoldingAllowed()) return false;
@@ -1290,6 +1320,11 @@ FoldingRule MergeSubSubArithmetic() {
            inst->opcode() == spv::Op::OpISub);
     const analysis::Type* type =
         context->get_type_mgr()->GetType(inst->type_id());
+
+    if (type->kind() == analysis::Type::kCooperativeMatrixKHR) {
+      return false;
+    }
+
     analysis::ConstantManager* const_mgr = context->get_constant_mgr();
     bool uses_float = HasFloatingPoint(type);
     if (uses_float && !inst->IsFloatingPointFoldingAllowed()) return false;
@@ -1383,6 +1418,11 @@ FoldingRule MergeGenericAddSubArithmetic() {
            inst->opcode() == spv::Op::OpIAdd);
     const analysis::Type* type =
         context->get_type_mgr()->GetType(inst->type_id());
+
+    if (type->kind() == analysis::Type::kCooperativeMatrixKHR) {
+      return false;
+    }
+
     bool uses_float = HasFloatingPoint(type);
     if (uses_float && !inst->IsFloatingPointFoldingAllowed()) return false;
 

diff --git a/source/opt/local_access_chain_convert_pass.cpp b/source/opt/local_access_chain_convert_pass.cpp
@@ -428,8 +428,8 @@ void LocalAccessChainConvertPass::InitExtensions() {
        "SPV_KHR_uniform_group_instructions",
        "SPV_KHR_fragment_shader_barycentric", "SPV_KHR_vulkan_memory_model",
        "SPV_NV_bindless_texture", "SPV_EXT_shader_atomic_float_add",
-       "SPV_EXT_fragment_shader_interlock",
-       "SPV_NV_compute_shader_derivatives"});
+       "SPV_EXT_fragment_shader_interlock", "SPV_NV_compute_shader_derivatives",
+       "SPV_KHR_cooperative_matrix"});
 }
 
 bool LocalAccessChainConvertPass::AnyIndexIsOutOfBounds(

diff --git a/source/opt/local_single_block_elim_pass.cpp b/source/opt/local_single_block_elim_pass.cpp
@@ -291,7 +291,8 @@ void LocalSingleBlockLoadStoreElimPass::InitExtensions() {
                                 "SPV_NV_bindless_texture",
                                 "SPV_EXT_shader_atomic_float_add",
                                 "SPV_EXT_fragment_shader_interlock",
-                                "SPV_NV_compute_shader_derivatives"});
+                                "SPV_NV_compute_shader_derivatives",
+                                "SPV_KHR_cooperative_matrix"});
 }
 
 }  // namespace opt

diff --git a/source/opt/local_single_store_elim_pass.cpp b/source/opt/local_single_store_elim_pass.cpp
@@ -141,7 +141,8 @@ void LocalSingleStoreElimPass::InitExtensionAllowList() {
                                 "SPV_NV_bindless_texture",
                                 "SPV_EXT_shader_atomic_float_add",
                                 "SPV_EXT_fragment_shader_interlock",
-                                "SPV_NV_compute_shader_derivatives"});
+                                "SPV_NV_compute_shader_derivatives",
+                                "SPV_KHR_cooperative_matrix"});
 }
 bool LocalSingleStoreElimPass::ProcessVariable(Instruction* var_inst) {
   std::vector<Instruction*> users;

diff --git a/source/opt/mem_pass.cpp b/source/opt/mem_pass.cpp
@@ -43,6 +43,7 @@ bool MemPass::IsBaseTargetType(const Instruction* typeInst) const {
     case spv::Op::OpTypeSampler:
     case spv::Op::OpTypeSampledImage:
     case spv::Op::OpTypePointer:
+    case spv::Op::OpTypeCooperativeMatrixKHR:
       return true;
     default:
       break;