Fix prefast warnings (#15340)

Fix prefast warnings: 
(1) Arithmetic overflow: Using operator '*' on a 4 byte value and then
casting the result to a 8 byte value. Cast the value to the wider type
before calling operator '*' to avoid overflow (io.2).
(2) Dereferencing NULL pointer 'key'.
This commit is contained in:
Tianlei Wu 2023-04-03 22:29:13 -07:00 committed by GitHub
parent dec11afb83
commit 3cf3fa0467
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
4 changed files with 29 additions and 27 deletions

View file

@ -202,7 +202,7 @@ Status CheckInputs(const T* query,
if (mask_dims.size() == 1) {
if (mask_dims[0] == static_cast<int64_t>(batch_size)) {
mask_type = AttentionMaskType::MASK_1D_KEY_SEQ_LEN;
} else if (mask_dims[0] == static_cast<int64_t>(3 * batch_size + 2)) {
} else if (mask_dims[0] == static_cast<int64_t>(3) * static_cast<int64_t>(batch_size) + static_cast<int64_t>(2)) {
mask_type = AttentionMaskType::MASK_1D_KEY_SEQ_LEN_START;
}
} else if (mask_dims.size() == 2 && mask_dims[0] == static_cast<int64_t>(batch_size) && mask_dims[1] == static_cast<int64_t>(kv_sequence_length)) {

View file

@ -207,25 +207,26 @@ Status LongformerAttention<T>::ComputeInternal(OpKernelContext* context) const {
input_data, k,
&zero, global_q, n, device_prop));
} else {
CUBLAS_RETURN_IF_ERROR(cublasGemmStridedBatchedHelper(cublas,
CUBLAS_OP_N,
CUBLAS_OP_N,
hidden_size, // m
max_num_global, // n
hidden_size, // k
&one, // alpha
global_q_weight, // A
hidden_size, // lda
0, // strideA
input_data, // B
hidden_size, // ldb
sequence_length * hidden_size, // strideB
&zero, // beta
global_q, // C
hidden_size, // ldc
max_num_global * hidden_size, // strideC
batch_size, // batch count
device_prop));
CUBLAS_RETURN_IF_ERROR(cublasGemmStridedBatchedHelper(
cublas,
CUBLAS_OP_N,
CUBLAS_OP_N,
hidden_size, // m
max_num_global, // n
hidden_size, // k
&one, // alpha
global_q_weight, // A
hidden_size, // lda
0, // strideA
input_data, // B
hidden_size, // ldb
static_cast<int64_t>(sequence_length) * hidden_size, // strideB
&zero, // beta
global_q, // C
hidden_size, // ldc
static_cast<int64_t>(max_num_global) * hidden_size, // strideC
batch_size, // batch count
device_prop));
}
// global k
const CudaT* global_k_weight = global_weights_data + static_cast<int64_t>(hidden_size) * hidden_size;

View file

@ -220,12 +220,13 @@ Status MultiHeadAttention<T>::ComputeInternal(OpKernelContext* context) const {
data.mask_index = (nullptr == key_padding_mask) ? nullptr : key_padding_mask->Data<int>();
data.mask_index_dims = (nullptr == key_padding_mask) ? gsl::span<const int64_t>() : key_padding_mask->Shape().GetDims();
data.past = nullptr;
data.past_key = (parameters.pass_past_in_kv) ? reinterpret_cast<const CudaT*>(key->Data<T>())
: (nullptr == past_key) ? nullptr
: reinterpret_cast<const CudaT*>(past_key->Data<T>());
data.past_value = (parameters.pass_past_in_kv) ? reinterpret_cast<const CudaT*>(value->Data<T>())
: (nullptr == past_value) ? nullptr
: reinterpret_cast<const CudaT*>(past_value->Data<T>());
const bool pass_key_value_as_past = (parameters.pass_past_in_kv && nullptr != key && nullptr != value);
data.past_key = pass_key_value_as_past ? reinterpret_cast<const CudaT*>(key->Data<T>())
: (nullptr == past_key) ? nullptr
: reinterpret_cast<const CudaT*>(past_key->Data<T>());
data.past_value = pass_key_value_as_past ? reinterpret_cast<const CudaT*>(value->Data<T>())
: (nullptr == past_value) ? nullptr
: reinterpret_cast<const CudaT*>(past_value->Data<T>());
data.relative_position_bias = (nullptr == relative_position_bias) ? nullptr : reinterpret_cast<const CudaT*>(relative_position_bias->Data<T>());
data.has_qkv_workspace = !no_qkv_workspace;
data.workspace = reinterpret_cast<CudaT*>(work_space.get());

View file

@ -217,7 +217,7 @@ static Fused_multihead_attention_params_mhca getMHCAParams(
params.gmem_q_params.cu_seqlens = static_cast<int32_t*>(cu_seqlens_q_d);
params.gmem_kv_params.ptr = const_cast<void*>(kv_packed_d);
params.gmem_kv_params.stride_in_bytes = h * 2 * d * sizeof(half);
params.gmem_kv_params.stride_in_bytes = static_cast<int64_t>(h) * 2 * d * static_cast<int64_t>(sizeof(half));
params.gmem_kv_params.h = h;
params.gmem_kv_params.d = d;
params.gmem_kv_params.cu_seqlens = static_cast<int32_t*>(cu_seqlens_kv_d);