Returns whether the provided operation class supports the provided data type combination :param op_class: operation class to consider :type op_class: cutlass_library.OpcodeClass :param element_a: data type of operand A :type element_a: cutlass_library.DataTy
(
self,
op_class: cutlass_library.OpcodeClass,
element_a: cutlass_library.DataType,
element_b: cutlass_library.DataType,
element_accumulator: cutlass_library.DataType,
layout_a: cutlass_library.LayoutType,
layout_b: cutlass_library.LayoutType,
math_operation: cutlass_library.MathOperation,
)
| 503 | return supporting_op_classes |
| 504 | |
| 505 | def operations( |
| 506 | self, |
| 507 | op_class: cutlass_library.OpcodeClass, |
| 508 | element_a: cutlass_library.DataType, |
| 509 | element_b: cutlass_library.DataType, |
| 510 | element_accumulator: cutlass_library.DataType, |
| 511 | layout_a: cutlass_library.LayoutType, |
| 512 | layout_b: cutlass_library.LayoutType, |
| 513 | math_operation: cutlass_library.MathOperation, |
| 514 | ) -> KernelsForDataType: |
| 515 | """ |
| 516 | Returns whether the provided operation class supports the provided data type combination |
| 517 | |
| 518 | :param op_class: operation class to consider |
| 519 | :type op_class: cutlass_library.OpcodeClass |
| 520 | :param element_a: data type of operand A |
| 521 | :type element_a: cutlass_library.DataType |
| 522 | :param element_b: data type of operand B |
| 523 | :type element_b: cutlass_library.DataType |
| 524 | :param element_accumulator: data type of accumulator |
| 525 | :type element_accumulator: cutlass_library.DataType |
| 526 | :param layout_a: layout of operand A |
| 527 | :type layout_a: cutlass_library.LayoutType |
| 528 | :param layout_b: layout of operand B |
| 529 | :type layout_b: cutlass_library.LayoutType |
| 530 | :param math_operation: math operation to consider |
| 531 | :type math_operation: cutlass_cppgen.MathOperation |
| 532 | |
| 533 | :return: container of kernels by alignment supported by the provided combination of parameters |
| 534 | :rtype: KernelsForDataType |
| 535 | """ |
| 536 | datatype_comb = (element_a, element_b, element_accumulator) |
| 537 | layout_comb = (layout_a, layout_b) |
| 538 | if not self.opclass_supports_combination(op_class, datatype_comb, layout_comb, math_operation): |
| 539 | raise Exception( |
| 540 | f"Data type layout combination {datatype_comb}, {layout_comb} " |
| 541 | f"is not supported by opcode class {op_class} on CC {self.cc}." |
| 542 | ) |
| 543 | return self.operations_by_opclass[op_class][(datatype_comb, layout_comb)] |
| 544 | |
| 545 | |
| 546 | class OptionRegistry: |