Google Git
Sign in
skia/external/github.com/KhronosGroup/SPIRV-Tools/fbcb6cf4c820ec83d0eff650a8323b8f92c137c9^!/.
commit
fbcb6cf4c820ec83d0eff650a8323b8f92c137c9
[log]
author
manas-kulkarni <manas@conffx.com>
Thu Jun 16 23:24:12 2022 +0530
committer
GitHub <noreply@github.com>
Thu Jun 16 13:54:12 2022 -0400
tree
71d7b9416807e0555e4c121901fee5a0cac2855c
parent
bfc611b03e0fe70c229d0f3b021cce465107f6f1 [diff]
Ability to fold Constant Vector times Matrix and Matrix times vector instructions (#4818)


diff --git a/source/opt/const_folding_rules.cpp b/source/opt/const_folding_rules.cpp
index 249e11e..cb36087 100644
--- a/source/opt/const_folding_rules.cpp
+++ b/source/opt/const_folding_rules.cpp

@@ -251,6 +251,193 @@
   };
 }
 
+ConstantFoldingRule FoldVectorTimesMatrix() {
+  return [](IRContext* context, Instruction* inst,
+            const std::vector<const analysis::Constant*>& constants)
+             -> const analysis::Constant* {
+    assert(inst->opcode() == SpvOpVectorTimesMatrix);
+    analysis::ConstantManager* const_mgr = context->get_constant_mgr();
+    analysis::TypeManager* type_mgr = context->get_type_mgr();
+
+    if (!inst->IsFloatingPointFoldingAllowed()) {
+      if (HasFloatingPoint(type_mgr->GetType(inst->type_id()))) {
+        return nullptr;
+      }
+    }
+
+    const analysis::Constant* c1 = constants[0];
+    const analysis::Constant* c2 = constants[1];
+
+    if (c1 == nullptr || c2 == nullptr) {
+      return nullptr;
+    }
+
+    // Check result type.
+    const analysis::Type* result_type = type_mgr->GetType(inst->type_id());
+    const analysis::Vector* vector_type = result_type->AsVector();
+    assert(vector_type != nullptr);
+    const analysis::Type* element_type = vector_type->element_type();
+    assert(element_type != nullptr);
+    const analysis::Float* float_type = element_type->AsFloat();
+    assert(float_type != nullptr);
+
+    // Check types of c1 and c2.
+    assert(c1->type()->AsVector() == vector_type);
+    assert(c1->type()->AsVector()->element_type() == element_type &&
+           c2->type()->AsMatrix()->element_type() == vector_type);
+
+    // Get a float vector that is the result of vector-times-matrix.
+    std::vector<const analysis::Constant*> c1_components =
+        c1->GetVectorComponents(const_mgr);
+    std::vector<const analysis::Constant*> c2_components =
+        c2->AsMatrixConstant()->GetComponents();
+    uint32_t resultVectorSize = result_type->AsVector()->element_count();
+
+    std::vector<uint32_t> ids;
+
+    if ((c1 && c1->IsZero()) || (c2 && c2->IsZero())) {
+      std::vector<uint32_t> words(float_type->width() / 32, 0);
+      for (uint32_t i = 0; i < resultVectorSize; ++i) {
+        const analysis::Constant* new_elem =
+            const_mgr->GetConstant(float_type, words);
+        ids.push_back(const_mgr->GetDefiningInstruction(new_elem)->result_id());
+      }
+      return const_mgr->GetConstant(vector_type, ids);
+    }
+
+    if (float_type->width() == 32) {
+      for (uint32_t i = 0; i < resultVectorSize; ++i) {
+        float result_scalar = 0.0f;
+        const analysis::VectorConstant* c2_vec =
+            c2_components[i]->AsVectorConstant();
+        for (uint32_t j = 0; j < c2_vec->GetComponents().size(); ++j) {
+          float c1_scalar = c1_components[j]->GetFloat();
+          float c2_scalar = c2_vec->GetComponents()[j]->GetFloat();
+          result_scalar += c1_scalar * c2_scalar;
+        }
+        utils::FloatProxy<float> result(result_scalar);
+        std::vector<uint32_t> words = result.GetWords();
+        const analysis::Constant* new_elem =
+            const_mgr->GetConstant(float_type, words);
+        ids.push_back(const_mgr->GetDefiningInstruction(new_elem)->result_id());
+      }
+      return const_mgr->GetConstant(vector_type, ids);
+    } else if (float_type->width() == 64) {
+      for (uint32_t i = 0; i < c2_components.size(); ++i) {
+        double result_scalar = 0.0;
+        const analysis::VectorConstant* c2_vec =
+            c2_components[i]->AsVectorConstant();
+        for (uint32_t j = 0; j < c2_vec->GetComponents().size(); ++j) {
+          double c1_scalar = c1_components[j]->GetDouble();
+          double c2_scalar = c2_vec->GetComponents()[j]->GetDouble();
+          result_scalar += c1_scalar * c2_scalar;
+        }
+        utils::FloatProxy<double> result(result_scalar);
+        std::vector<uint32_t> words = result.GetWords();
+        const analysis::Constant* new_elem =
+            const_mgr->GetConstant(float_type, words);
+        ids.push_back(const_mgr->GetDefiningInstruction(new_elem)->result_id());
+      }
+      return const_mgr->GetConstant(vector_type, ids);
+    }
+    return nullptr;
+  };
+}
+
+ConstantFoldingRule FoldMatrixTimesVector() {
+  return [](IRContext* context, Instruction* inst,
+            const std::vector<const analysis::Constant*>& constants)
+             -> const analysis::Constant* {
+    assert(inst->opcode() == SpvOpMatrixTimesVector);
+    analysis::ConstantManager* const_mgr = context->get_constant_mgr();
+    analysis::TypeManager* type_mgr = context->get_type_mgr();
+
+    if (!inst->IsFloatingPointFoldingAllowed()) {
+      if (HasFloatingPoint(type_mgr->GetType(inst->type_id()))) {
+        return nullptr;
+      }
+    }
+
+    const analysis::Constant* c1 = constants[0];
+    const analysis::Constant* c2 = constants[1];
+
+    if (c1 == nullptr || c2 == nullptr) {
+      return nullptr;
+    }
+
+    // Check result type.
+    const analysis::Type* result_type = type_mgr->GetType(inst->type_id());
+    const analysis::Vector* vector_type = result_type->AsVector();
+    assert(vector_type != nullptr);
+    const analysis::Type* element_type = vector_type->element_type();
+    assert(element_type != nullptr);
+    const analysis::Float* float_type = element_type->AsFloat();
+    assert(float_type != nullptr);
+
+    // Check types of c1 and c2.
+    assert(c1->type()->AsMatrix()->element_type() == vector_type);
+    assert(c2->type()->AsVector()->element_type() == element_type);
+
+    // Get a float vector that is the result of matrix-times-vector.
+    std::vector<const analysis::Constant*> c1_components =
+        c1->AsMatrixConstant()->GetComponents();
+    std::vector<const analysis::Constant*> c2_components =
+        c2->GetVectorComponents(const_mgr);
+    uint32_t resultVectorSize = result_type->AsVector()->element_count();
+
+    std::vector<uint32_t> ids;
+
+    if ((c1 && c1->IsZero()) || (c2 && c2->IsZero())) {
+      std::vector<uint32_t> words(float_type->width() / 32, 0);
+      for (uint32_t i = 0; i < resultVectorSize; ++i) {
+        const analysis::Constant* new_elem =
+            const_mgr->GetConstant(float_type, words);
+        ids.push_back(const_mgr->GetDefiningInstruction(new_elem)->result_id());
+      }
+      return const_mgr->GetConstant(vector_type, ids);
+    }
+
+    if (float_type->width() == 32) {
+      for (uint32_t i = 0; i < resultVectorSize; ++i) {
+        float result_scalar = 0.0f;
+        for (uint32_t j = 0; j < c1_components.size(); ++j) {
+          float c1_scalar = c1_components[j]
+                                ->AsVectorConstant()
+                                ->GetComponents()[i]
+                                ->GetFloat();
+          float c2_scalar = c2_components[j]->GetFloat();
+          result_scalar += c1_scalar * c2_scalar;
+        }
+        utils::FloatProxy<float> result(result_scalar);
+        std::vector<uint32_t> words = result.GetWords();
+        const analysis::Constant* new_elem =
+            const_mgr->GetConstant(float_type, words);
+        ids.push_back(const_mgr->GetDefiningInstruction(new_elem)->result_id());
+      }
+      return const_mgr->GetConstant(vector_type, ids);
+    } else if (float_type->width() == 64) {
+      for (uint32_t i = 0; i < resultVectorSize; ++i) {
+        double result_scalar = 0.0;
+        for (uint32_t j = 0; j < c1_components.size(); ++j) {
+          double c1_scalar = c1_components[j]
+                                 ->AsVectorConstant()
+                                 ->GetComponents()[i]
+                                 ->GetDouble();
+          double c2_scalar = c2_components[j]->GetDouble();
+          result_scalar += c1_scalar * c2_scalar;
+        }
+        utils::FloatProxy<double> result(result_scalar);
+        std::vector<uint32_t> words = result.GetWords();
+        const analysis::Constant* new_elem =
+            const_mgr->GetConstant(float_type, words);
+        ids.push_back(const_mgr->GetDefiningInstruction(new_elem)->result_id());
+      }
+      return const_mgr->GetConstant(vector_type, ids);
+    }
+    return nullptr;
+  };
+}
+
 ConstantFoldingRule FoldCompositeWithConstants() {
   // Folds an OpCompositeConstruct where all of the inputs are constants to a
   // constant.  A new constant is created if necessary.
@@ -1288,6 +1475,8 @@
 
   rules_[SpvOpVectorShuffle].push_back(FoldVectorShuffleWithConstants());
   rules_[SpvOpVectorTimesScalar].push_back(FoldVectorTimesScalar());
+  rules_[SpvOpVectorTimesMatrix].push_back(FoldVectorTimesMatrix());
+  rules_[SpvOpMatrixTimesVector].push_back(FoldMatrixTimesVector());
 
   rules_[SpvOpFNegate].push_back(FoldFNegate());
   rules_[SpvOpQuantizeToF16].push_back(FoldQuantizeToF16());

diff --git a/test/opt/fold_test.cpp b/test/opt/fold_test.cpp
index 8ee14a9..f6c94ca 100644
--- a/test/opt/fold_test.cpp
+++ b/test/opt/fold_test.cpp

@@ -148,6 +148,8 @@
 %v2half = OpTypeVector %half 2
 %v2bool = OpTypeVector %bool 2
 %m2x2int = OpTypeMatrix %v2int 2
+%mat4v4float = OpTypeMatrix %v4float 4
+%mat4v4double = OpTypeMatrix %v4double 4
 %struct_v2int_int_int = OpTypeStruct %v2int %int %int
 %_ptr_int = OpTypePointer Function %int
 %_ptr_uint = OpTypePointer Function %uint
@@ -276,13 +278,20 @@
 %v4float_0_0_0_1 = OpConstantComposite %v4float %float_0 %float_0 %float_0 %float_1
 %v4float_0_1_0_0 = OpConstantComposite %v4float %float_0 %float_1 %float_null %float_0
 %v4float_1_1_1_1 = OpConstantComposite %v4float %float_1 %float_1 %float_1 %float_1
+%v4float_1_2_3_4 = OpConstantComposite %v4float %float_1 %float_2 %float_3 %float_4
+%v4float_null = OpConstantNull %v4float
+%mat4v4float_null = OpConstantComposite %mat4v4float %v4float_null %v4float_null %v4float_null %v4float_null
+%mat4v4float_1_2_3_4 = OpConstantComposite %mat4v4float %v4float_1_2_3_4 %v4float_1_2_3_4 %v4float_1_2_3_4 %v4float_1_2_3_4
 %107 = OpConstantComposite %v4double %double_0 %double_0 %double_0 %double_0
 %v4double_0_0_0_0 = OpConstantComposite %v4double %double_0 %double_0 %double_0 %double_0
 %v4double_0_0_0_1 = OpConstantComposite %v4double %double_0 %double_0 %double_0 %double_1
 %v4double_0_1_0_0 = OpConstantComposite %v4double %double_0 %double_1 %double_null %double_0
 %v4double_1_1_1_1 = OpConstantComposite %v4double %double_1 %double_1 %double_1 %double_1
+%v4double_1_2_3_4 = OpConstantComposite %v4double %double_1 %double_2 %double_3 %double_4
 %v4double_1_1_1_0p5 = OpConstantComposite %v4double %double_1 %double_1 %double_1 %double_0p5
 %v4double_null = OpConstantNull %v4double
+%mat4v4double_null = OpConstantComposite %mat4v4double %v4double_null %v4double_null %v4double_null %v4double_null
+%mat4v4double_1_2_3_4 = OpConstantComposite %mat4v4double %v4double_1_2_3_4 %v4double_1_2_3_4 %v4double_1_2_3_4 %v4double_1_2_3_4
 %v4float_n1_2_1_3 = OpConstantComposite %v4float %float_n1 %float_2 %float_1 %float_3
 %uint_0x3f800000 = OpConstant %uint 0x3f800000
 %uint_0xbf800000 = OpConstant %uint 0xbf800000
@@ -912,7 +921,61 @@
            "%2 = OpBitcast %v2double %v4int_0x3FF00000_0x00000000_0xC05FD666_0x66666666\n" +
            "OpReturn\n" +
            "OpFunctionEnd",
-       2, {1.0,-127.35})
+       2, {1.0,-127.35}),
+   // Test case 1: OpVectorTimesMatrix Non-Zero Zero {{0.0, 0.0, 0.0, 0.0}, {0.0, 0.0, 0.0, 0.0}, {0.0, 0.0, 0.0, 0.0}, {0.0, 0.0, 0.0, 0.0}} {1.0, 2.0, 3.0, 4.0} {0.0, 0.0, 0.0, 0.0}
+   InstructionFoldingCase<std::vector<double>>(
+       Header() +
+       "%main = OpFunction %void None %void_func\n" +
+       "%main_lab = OpLabel\n" +
+       "%2 = OpVectorTimesMatrix %v4double %v4double_1_2_3_4 %mat4v4double_null\n" +
+       "OpReturn\n" +
+       "OpFunctionEnd",
+       2, {0.0,0.0,0.0,0.0}),
+   // Test case 2: OpVectorTimesMatrix Zero Non-Zero {{1.0, 2.0, 3.0, 4.0}, {1.0, 2.0, 3.0, 4.0}, {1.0, 2.0, 3.0, 4.0}, {1.0, 2.0, 3.0, 4.0}} {0.0, 0.0, 0.0, 0.0} {0.0, 0.0, 0.0, 0.0}
+   InstructionFoldingCase<std::vector<double>>(
+       Header() +
+       "%main = OpFunction %void None %void_func\n" +
+       "%main_lab = OpLabel\n" +
+       "%2 = OpVectorTimesMatrix %v4double %v4double_null %mat4v4double_1_2_3_4\n" +
+       "OpReturn\n" +
+       "OpFunctionEnd",
+       2, {0.0,0.0,0.0,0.0}),
+   // Test case 3: OpVectorTimesMatrix Non-Zero Non-Zero {{1.0, 2.0, 3.0, 4.0}, {1.0, 2.0, 3.0, 4.0}, {1.0, 2.0, 3.0, 4.0}, {1.0, 2.0, 3.0, 4.0}} {1.0, 2.0, 3.0, 4.0} {30.0, 30.0, 30.0, 30.0}
+   InstructionFoldingCase<std::vector<double>>(
+       Header() +
+       "%main = OpFunction %void None %void_func\n" +
+       "%main_lab = OpLabel\n" +
+       "%2 = OpVectorTimesMatrix %v4double %v4double_1_2_3_4 %mat4v4double_1_2_3_4\n" +
+       "OpReturn\n" +
+       "OpFunctionEnd",
+       2, {30.0,30.0,30.0,30.0}),
+   // Test case 4: OpMatrixTimesVector Zero Non-Zero {1.0, 2.0, 3.0, 4.0} {{0.0, 0.0, 0.0, 0.0}, {0.0, 0.0, 0.0, 0.0}, {0.0, 0.0, 0.0, 0.0}, {0.0, 0.0, 0.0, 0.0}} {0.0, 0.0, 0.0, 0.0}
+   InstructionFoldingCase<std::vector<double>>(
+       Header() +
+       "%main = OpFunction %void None %void_func\n" +
+       "%main_lab = OpLabel\n" +
+       "%2 = OpMatrixTimesVector %v4double %mat4v4double_null %v4double_1_2_3_4\n" +
+       "OpReturn\n" +
+       "OpFunctionEnd",
+       2, {0.0,0.0,0.0,0.0}),
+   // Test case 5: OpMatrixTimesVector Non-Zero Zero {{1.0, 2.0, 3.0, 4.0}, {1.0, 2.0, 3.0, 4.0}, {1.0, 2.0, 3.0, 4.0}, {1.0, 2.0, 3.0, 4.0}} {0.0, 0.0, 0.0, 0.0} {0.0, 0.0, 0.0, 0.0}
+   InstructionFoldingCase<std::vector<double>>(
+       Header() +
+       "%main = OpFunction %void None %void_func\n" +
+       "%main_lab = OpLabel\n" +
+       "%2 = OpMatrixTimesVector %v4double %mat4v4double_1_2_3_4 %v4double_null\n" +
+       "OpReturn\n" +
+       "OpFunctionEnd",
+       2, {0.0,0.0,0.0,0.0}),
+   // Test case 6: OpMatrixTimesVector Non-Zero Non-Zero {1.0, 2.0, 3.0, 4.0} {{1.0, 2.0, 3.0, 4.0}, {1.0, 2.0, 3.0, 4.0}, {1.0, 2.0, 3.0, 4.0}, {1.0, 2.0, 3.0, 4.0}} {10.0, 20.0, 30.0, 40.0}
+   InstructionFoldingCase<std::vector<double>>(
+       Header() +
+       "%main = OpFunction %void None %void_func\n" +
+       "%main_lab = OpLabel\n" +
+       "%2 = OpMatrixTimesVector %v4double %mat4v4double_1_2_3_4 %v4double_1_2_3_4\n" +
+       "OpReturn\n" +
+       "OpFunctionEnd",
+       2, {10.0,20.0,30.0,40.0})
 ));
 
 using FloatVectorInstructionFoldingTest =
@@ -981,7 +1044,61 @@
            "%2 = OpBitcast %v2float %long_0xbf8000003f800000\n" +
            "OpReturn\n" +
            "OpFunctionEnd",
-       2, {1.0f,-1.0f})
+       2, {1.0f,-1.0f}),
+   // Test case 3: OpVectorTimesMatrix Non-Zero Zero {{0.0, 0.0, 0.0, 0.0}, {0.0, 0.0, 0.0, 0.0}, {0.0, 0.0, 0.0, 0.0}, {0.0, 0.0, 0.0, 0.0}} {1.0, 2.0, 3.0, 4.0} {0.0, 0.0, 0.0, 0.0}
+   InstructionFoldingCase<std::vector<float>>(
+       Header() +
+       "%main = OpFunction %void None %void_func\n" +
+       "%main_lab = OpLabel\n" +
+       "%2 = OpVectorTimesMatrix %v4float %v4float_1_2_3_4 %mat4v4float_null\n" +
+       "OpReturn\n" +
+       "OpFunctionEnd",
+       2, {0.0f,0.0f,0.0f,0.0f}),
+   // Test case 4: OpVectorTimesMatrix Zero Non-Zero {{1.0, 2.0, 3.0, 4.0}, {1.0, 2.0, 3.0, 4.0}, {1.0, 2.0, 3.0, 4.0}, {1.0, 2.0, 3.0, 4.0}} {0.0, 0.0, 0.0, 0.0} {0.0, 0.0, 0.0, 0.0}
+   InstructionFoldingCase<std::vector<float>>(
+       Header() +
+       "%main = OpFunction %void None %void_func\n" +
+       "%main_lab = OpLabel\n" +
+       "%2 = OpVectorTimesMatrix %v4float %v4float_null %mat4v4float_1_2_3_4\n" +
+       "OpReturn\n" +
+       "OpFunctionEnd",
+       2, {0.0f,0.0f,0.0f,0.0f}),
+   // Test case 5: OpVectorTimesMatrix Non-Zero Non-Zero {{1.0, 2.0, 3.0, 4.0}, {1.0, 2.0, 3.0, 4.0}, {1.0, 2.0, 3.0, 4.0}, {1.0, 2.0, 3.0, 4.0}} {1.0, 2.0, 3.0, 4.0} {30.0, 30.0, 30.0, 30.0}
+   InstructionFoldingCase<std::vector<float>>(
+       Header() +
+       "%main = OpFunction %void None %void_func\n" +
+       "%main_lab = OpLabel\n" +
+       "%2 = OpVectorTimesMatrix %v4float %v4float_1_2_3_4 %mat4v4float_1_2_3_4\n" +
+       "OpReturn\n" +
+       "OpFunctionEnd",
+       2, {30.0f,30.0f,30.0f,30.0f}),
+   // Test case 6: OpMatrixTimesVector Zero Non-Zero {1.0, 2.0, 3.0, 4.0} {{0.0, 0.0, 0.0, 0.0}, {0.0, 0.0, 0.0, 0.0}, {0.0, 0.0, 0.0, 0.0}, {0.0, 0.0, 0.0, 0.0}} {0.0, 0.0, 0.0, 0.0}
+   InstructionFoldingCase<std::vector<float>>(
+       Header() +
+       "%main = OpFunction %void None %void_func\n" +
+       "%main_lab = OpLabel\n" +
+       "%2 = OpMatrixTimesVector %v4float %mat4v4float_null %v4float_1_2_3_4\n" +
+       "OpReturn\n" +
+       "OpFunctionEnd",
+       2, {0.0f,0.0f,0.0f,0.0f}),
+   // Test case 7: OpMatrixTimesVector Non-Zero Zero {{1.0, 2.0, 3.0, 4.0}, {1.0, 2.0, 3.0, 4.0}, {1.0, 2.0, 3.0, 4.0}, {1.0, 2.0, 3.0, 4.0}} {0.0, 0.0, 0.0, 0.0} {0.0, 0.0, 0.0, 0.0}
+   InstructionFoldingCase<std::vector<float>>(
+       Header() +
+       "%main = OpFunction %void None %void_func\n" +
+       "%main_lab = OpLabel\n" +
+       "%2 = OpMatrixTimesVector %v4float %mat4v4float_1_2_3_4 %v4float_null\n" +
+       "OpReturn\n" +
+       "OpFunctionEnd",
+       2, {0.0f,0.0f,0.0f,0.0f}),
+   // Test case 8: OpMatrixTimesVector Non-Zero Non-Zero {1.0, 2.0, 3.0, 4.0} {{1.0, 2.0, 3.0, 4.0}, {1.0, 2.0, 3.0, 4.0}, {1.0, 2.0, 3.0, 4.0}, {1.0, 2.0, 3.0, 4.0}} {10.0, 20.0, 30.0, 40.0}
+   InstructionFoldingCase<std::vector<float>>(
+       Header() +
+       "%main = OpFunction %void None %void_func\n" +
+       "%main_lab = OpLabel\n" +
+       "%2 = OpMatrixTimesVector %v4float %mat4v4float_1_2_3_4 %v4float_1_2_3_4\n" +
+       "OpReturn\n" +
+       "OpFunctionEnd",
+       2, {10.0f,20.0f,30.0f,40.0f})
 ));
 // clang-format on
 using BooleanInstructionFoldingTest =