| 101 | } |
| 102 | |
| 103 | inline int32_t getMHAMaskPackedSize(int32_t smVersion, nvinfer1::DataType dataType, int32_t sequenceLength) |
| 104 | { |
| 105 | // this code must match EmbLayerNormPluginDynamic::getOutputDimensions in embLayerNormPlugin.cpp |
| 106 | int32_t packedSize = unfusedMaskSize; |
| 107 | bool isSmOK = (smVersion == kSM_75 || smVersion == kSM_80 || smVersion == kSM_86 || smVersion == kSM_87 |
| 108 | || smVersion == kSM_90); |
| 109 | bool isPrecisionOK = (dataType == nvinfer1::DataType::kINT8 || dataType == nvinfer1::DataType::kHALF); |
| 110 | if (isSmOK && isPrecisionOK) |
| 111 | { |
| 112 | if (sequenceLength == 64) |
| 113 | { |
| 114 | packedSize = packedMaskSize64; |
| 115 | } |
| 116 | else if (sequenceLength == 96) |
| 117 | { |
| 118 | packedSize = packedMaskSize96; |
| 119 | } |
| 120 | else if (sequenceLength == 128) |
| 121 | { |
| 122 | packedSize = packedMaskSize128; |
| 123 | } |
| 124 | else if (sequenceLength == 384) |
| 125 | { |
| 126 | packedSize = packedMaskSize384; |
| 127 | } |
| 128 | } |
| 129 | return packedSize; |
| 130 | } |
| 131 | |
| 132 | inline uint32_t getElementSize(nvinfer1::DataType t) noexcept |
| 133 | { |
no outgoing calls
no test coverage detected