Skip to content

Commit 4186f6c

Browse files
authored
Fix security issues after LibingAI scan (#438)
1 parent cc00faa commit 4186f6c

3 files changed

Lines changed: 10 additions & 2 deletions

File tree

csrc/assign_cache_op/op_host/assign_cache.cpp

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -36,6 +36,12 @@ HOST_API size_t GetElementByteSize(const at::Tensor &tensor)
3636
HOST_API bool assign_cache_op(at::Tensor &dstTensor, const at::Tensor &srcTensor, const at::Tensor &dstStartIdx,
3737
const at::Tensor &dstEndIdx, const at::Tensor &srcStartIdx, const at::Tensor &srcEndIdx)
3838
{
39+
OP_CHECK(dstTensor.defined() && dstTensor.dim() >= 1, "dstTensor is invalid or 0-dimensional", return false);
40+
OP_CHECK(srcTensor.defined() && srcTensor.dim() >= 1, "srcTensor is invalid or 0-dimensional", return false);
41+
OP_CHECK(dstStartIdx.defined() && dstStartIdx.dim() >= 1, "dstStartIdx is invalid or 0-dimensional", return false);
42+
OP_CHECK(dstEndIdx.defined() && dstEndIdx.dim() >= 1, "dstEndIdx is invalid or 0-dimensional", return false);
43+
OP_CHECK(srcStartIdx.defined() && srcStartIdx.dim() >= 1, "srcStartIdx is invalid or 0-dimensional", return false);
44+
OP_CHECK(srcEndIdx.defined() && srcEndIdx.dim() >= 1, "srcEndIdx is invalid or 0-dimensional", return false);
3945
auto dstShape = dstTensor.sizes(), dstStartShape = dstStartIdx.sizes(), dstEndShape = dstEndIdx.sizes();
4046
auto srcShape = srcTensor.sizes(), srcStartShape = srcStartIdx.sizes(), srcEndShape = srcEndIdx.sizes();
4147
OP_CHECK(dstShape[0] == srcShape[0] && dstStartShape[0] == srcStartShape[0] && dstEndShape[0] == srcEndShape[0],

csrc/lightning_indexer/op_host/tiling/lightning_indexer_tiling.cpp

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -69,7 +69,7 @@ ge::graphStatus LIInfoParser::GetNpuInfo()
6969
auto ascendcPlatform = *platform_ascendc::PlatformAscendCManager::GetInstance();
7070
uint32_t aivNum = ascendcPlatform.GetCoreNumAiv();
7171
uint32_t aicNum = ascendcPlatform.GetCoreNumAic();
72-
TORCH_CHECK(aivNum != 0 && aivNum != 0, OPS_LOG_E(opName_, "num of core obtained is 0"));
72+
TORCH_CHECK(aivNum != 0 && aicNum != 0, OPS_LOG_E(opName_, "num of core obtained is 0"));
7373

7474
socVersion_ = ascendcPlatform.GetSocVersion();
7575
TORCH_CHECK(socVersion_ == platform_ascendc::SocVersion::ASCEND910B ||
@@ -409,6 +409,8 @@ ge::graphStatus LIInfoParser::ValidateInputShapesMatch()
409409
if (qLayout_ == DataLayout::TND) {
410410
// -----------------------check BatchSize-------------------
411411
// bSize_ 来源于act_seq_q
412+
TORCH_CHECK(opParamInfo_.actualSeqLengths.tensor != nullptr, opName_, ": actualSeqLengths tensor is null");
413+
412414
TORCH_CHECK((opParamInfo_.actualSeqLengths.tensor->GetShapeSize() == bSize_) &&
413415
(opParamInfo_.blockTable.tensor == nullptr ||
414416
opParamInfo_.blockTable.tensor->GetStorageShape().GetDim(0) == bSize_),

csrc/lightning_indexer/op_kernel/lightning_indexer_vector.h

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -119,7 +119,7 @@ __aicore__ inline uint64_t FindNearestPower2(uint64_t value)
119119
return value;
120120
} else {
121121
const uint64_t pow = 63 - clz(value); // clz返回前导0的个数,对于64位整数,最大有效位位置 = 63 - 前导0个数
122-
return (1 << pow);
122+
return (1ULL << pow);
123123
}
124124
}
125125

0 commit comments

Comments
 (0)