Skip to content

Commit 2d7060b

Browse files
jeffbolznvjeffbolznv
authored
Validate SPV_NV_cooperative_matrix_decode_vector (KhronosGroup#6693)
1 parent 8b5f796 commit 2d7060b

6 files changed

Lines changed: 973 additions & 25 deletions

File tree

‎DEPS‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,7 @@ vars = {
1414

1515
're2_revision': '972a15cedd008d846f1a39b2e88ce48d7f166cbd',
1616

17-
'spirv_headers_revision': '82dfca37143878c848d8df1f9e4d746f65e6827e',
17+
'spirv_headers_revision': 'aaffbc59b41bf8faea671b0d1c6b34a584c45171',
1818

1919
'mimalloc_revision': 'fef6b0dd70f9d7fa0750b0d0b9fbb471203b94cd',
2020
}

‎source/opt/ir_context.cpp‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -996,6 +996,11 @@ void IRContext::AddCalls(const Function* func, std::queue<uint32_t>* todo) {
996996

997997
if (mask & uint32_t(spv::TensorAddressingOperandsMask::DecodeFunc)) {
998998
todo->push(ii->GetSingleWordInOperand(tensor_operands_index + count));
999+
++count;
1000+
}
1001+
if (mask &
1002+
uint32_t(spv::TensorAddressingOperandsMask::DecodeVectorFunc)) {
1003+
todo->push(ii->GetSingleWordInOperand(tensor_operands_index + count));
9991004
}
10001005
}
10011006
}

‎source/val/validate_memory.cpp‎

Lines changed: 95 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -2765,6 +2765,9 @@ int TensorAddressingOperandsNumWords(spv::TensorAddressingOperandsMask mask) {
27652765
if ((mask & spv::TensorAddressingOperandsMask::DecodeFunc) !=
27662766
spv::TensorAddressingOperandsMask::MaskNone)
27672767
++result;
2768+
if ((mask & spv::TensorAddressingOperandsMask::DecodeVectorFunc) !=
2769+
spv::TensorAddressingOperandsMask::MaskNone)
2770+
++result;
27682771
return result;
27692772
}
27702773

@@ -2893,32 +2896,74 @@ spv_result_t ValidateCooperativeMatrixLoadStoreTensorNV(
28932896
tensor_operand_index++;
28942897
}
28952898

2896-
if ((tensor_operands & spv::TensorAddressingOperandsMask::DecodeFunc) !=
2897-
spv::TensorAddressingOperandsMask::MaskNone) {
2899+
const bool has_decode_func =
2900+
(tensor_operands & spv::TensorAddressingOperandsMask::DecodeFunc) !=
2901+
spv::TensorAddressingOperandsMask::MaskNone;
2902+
const bool has_decode_vector_func =
2903+
(tensor_operands & spv::TensorAddressingOperandsMask::DecodeVectorFunc) !=
2904+
spv::TensorAddressingOperandsMask::MaskNone;
2905+
2906+
if (has_decode_func || has_decode_vector_func) {
28982907
if (inst->opcode() == spv::Op::OpCooperativeMatrixStoreTensorNV) {
28992908
return _.diag(SPV_ERROR_INVALID_ID, inst)
2900-
<< "OpCooperativeMatrixStoreTensorNV does not support DecodeFunc.";
2909+
<< "OpCooperativeMatrixStoreTensorNV does not support DecodeFunc "
2910+
"or DecodeVectorFunc.";
29012911
}
2902-
const auto decode_func_id =
2903-
inst->GetOperandAs<uint32_t>(tensor_operand_index);
2912+
}
2913+
2914+
const auto component_type_index = 1;
2915+
const auto component_type_id =
2916+
matrix_type->GetOperandAs<uint32_t>(component_type_index);
2917+
const auto tensor_layout_type = _.FindDef(tensor_layout->type_id());
2918+
2919+
// Validate one decode-function operand (scalar DecodeFunc or vector
2920+
// DecodeVectorFunc). `expect_vector` selects which return-type rule to
2921+
// enforce.
2922+
auto validate_decode_function = [&](uint32_t decode_func_id,
2923+
const char* operand_name,
2924+
bool expect_vector) -> spv_result_t {
29042925
const auto decode_func = _.FindDef(decode_func_id);
29052926

29062927
if (!decode_func || decode_func->opcode() != spv::Op::OpFunction) {
29072928
return _.diag(SPV_ERROR_INVALID_ID, inst)
2908-
<< opname << " DecodeFunc <id> " << _.getIdName(decode_func_id)
2909-
<< " is not a function.";
2929+
<< opname << " " << operand_name << " <id> "
2930+
<< _.getIdName(decode_func_id) << " is not a function.";
29102931
}
29112932

2912-
const auto component_type_index = 1;
2913-
const auto component_type_id =
2914-
matrix_type->GetOperandAs<uint32_t>(component_type_index);
2915-
29162933
const auto function_type =
29172934
_.FindDef(decode_func->GetOperandAs<uint32_t>(3));
2918-
if (function_type->GetOperandAs<uint32_t>(1) != component_type_id) {
2919-
return _.diag(SPV_ERROR_INVALID_ID, inst)
2920-
<< opname << " DecodeFunc <id> " << _.getIdName(decode_func_id)
2921-
<< " return type must match matrix component type.";
2935+
const auto return_type_id = function_type->GetOperandAs<uint32_t>(1);
2936+
const auto return_type = _.FindDef(return_type_id);
2937+
const bool return_is_scalar_match = (return_type_id == component_type_id);
2938+
const bool return_is_vector = _.IsVectorType(return_type_id);
2939+
const uint32_t return_vec_component_type_id =
2940+
return_is_vector ? return_type->GetOperandAs<uint32_t>(1) : 0;
2941+
2942+
if (!expect_vector) {
2943+
if (!return_is_scalar_match) {
2944+
return _.diag(SPV_ERROR_INVALID_ID, inst)
2945+
<< opname << " " << operand_name << " <id> "
2946+
<< _.getIdName(decode_func_id)
2947+
<< " return type must match matrix component type.";
2948+
}
2949+
} else {
2950+
if (!return_is_vector ||
2951+
return_vec_component_type_id != component_type_id) {
2952+
return _.diag(SPV_ERROR_INVALID_ID, inst)
2953+
<< opname << " " << operand_name << " <id> "
2954+
<< _.getIdName(decode_func_id)
2955+
<< " return type must be a vector of the matrix component "
2956+
"type.";
2957+
}
2958+
// GetDimension returns 0 for OpTypeVectorIdEXT whose count is a
2959+
// spec constant; skip the static check in that case.
2960+
const uint32_t v = _.GetDimension(return_type_id);
2961+
if (v != 0 && v != 2 && v != 4 && v != 8) {
2962+
return _.diag(SPV_ERROR_INVALID_ID, inst)
2963+
<< opname << " " << operand_name << " <id> "
2964+
<< _.getIdName(decode_func_id)
2965+
<< " return vector length must be 2, 4, or 8.";
2966+
}
29222967
}
29232968

29242969
const auto decode_ptr_type_id = function_type->GetOperandAs<uint32_t>(2);
@@ -2928,21 +2973,20 @@ spv_result_t ValidateCooperativeMatrixLoadStoreTensorNV(
29282973

29292974
if (decode_storage_class != spv::StorageClass::PhysicalStorageBuffer) {
29302975
return _.diag(SPV_ERROR_INVALID_ID, inst)
2931-
<< opname << " DecodeFunc <id> " << _.getIdName(decode_func_id)
2976+
<< opname << " " << operand_name << " <id> "
2977+
<< _.getIdName(decode_func_id)
29322978
<< " first parameter must be pointer to PhysicalStorageBuffer.";
29332979
}
29342980

2935-
const auto tensor_layout_type = _.FindDef(tensor_layout->type_id());
2936-
29372981
for (uint32_t param = 3; param < 5; ++param) {
29382982
const auto param_type_id = function_type->GetOperandAs<uint32_t>(param);
29392983
const auto param_type = _.FindDef(param_type_id);
29402984
if (param_type->opcode() != spv::Op::OpTypeArray) {
29412985
return _.diag(SPV_ERROR_INVALID_ID, inst)
2942-
<< opname << " DecodeFunc <id> " << _.getIdName(decode_func_id)
2986+
<< opname << " " << operand_name << " <id> "
2987+
<< _.getIdName(decode_func_id)
29432988
<< " second/third parameter must be array of 32-bit integer "
2944-
"with "
2945-
<< " dimension equal to the tensor dimension.";
2989+
"with dimension equal to the tensor dimension.";
29462990
}
29472991
const auto length_index = 2u;
29482992
uint64_t array_length;
@@ -2955,16 +2999,43 @@ spv_result_t ValidateCooperativeMatrixLoadStoreTensorNV(
29552999
if (_.EvalConstantValUint64(tensor_layout_dim_id, &dim_value)) {
29563000
if (array_length != dim_value) {
29573001
return _.diag(SPV_ERROR_INVALID_ID, inst)
2958-
<< opname << " DecodeFunc <id> "
3002+
<< opname << " " << operand_name << " <id> "
29593003
<< _.getIdName(decode_func_id)
29603004
<< " second/third parameter must be array of 32-bit integer "
2961-
"with "
2962-
<< " dimension equal to the tensor dimension.";
3005+
"with dimension equal to the tensor dimension.";
29633006
}
29643007
}
29653008
}
29663009
}
29673010

3011+
return SPV_SUCCESS;
3012+
};
3013+
3014+
if (has_decode_func) {
3015+
const uint32_t decode_func_id =
3016+
inst->GetOperandAs<uint32_t>(tensor_operand_index);
3017+
if (auto error = validate_decode_function(decode_func_id, "DecodeFunc",
3018+
/*expect_vector=*/false)) {
3019+
return error;
3020+
}
3021+
tensor_operand_index++;
3022+
}
3023+
3024+
if (has_decode_vector_func) {
3025+
if (!has_decode_func) {
3026+
return _.diag(SPV_ERROR_INVALID_ID, inst)
3027+
<< opname
3028+
<< " DecodeVectorFunc requires DecodeFunc to also be specified.";
3029+
}
3030+
3031+
const uint32_t decode_vector_func_id =
3032+
inst->GetOperandAs<uint32_t>(tensor_operand_index);
3033+
if (auto error =
3034+
validate_decode_function(decode_vector_func_id, "DecodeVectorFunc",
3035+
/*expect_vector=*/true)) {
3036+
return error;
3037+
}
3038+
29683039
tensor_operand_index++;
29693040
}
29703041

‎test/opt/CMakeLists.txt‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,7 @@ add_spvtools_unittest(TARGET opt
2727
cfg_test.cpp
2828
code_sink_test.cpp
2929
combine_access_chains_test.cpp
30+
coop_matrix_decode_vector_opt_test.cpp
3031
compact_ids_test.cpp
3132
constants_test.cpp
3233
constant_manager_test.cpp

0 commit comments

Comments
 (0)