[ZBAL] 修复代码安全编码问题-memfabric-hybrid-AtomGit

Mapp/zbal/src/csrc/adaptor/pytorch_npu/zbal_pytorch_process_group.cpp+5-3

	@@ -980,8 +980,8 @@ ProcessGroupZBAL::~ProcessGroupZBAL()
980	if (groupComm_ != nullptr) {	980	if (groupComm_ != nullptr) {
981	auto result = zbal_comm_destroy(groupComm_, 0);	981	auto result = zbal_comm_destroy(groupComm_, 0);
982	if (result != Z_OK) {	982	if (result != Z_OK) {
		983	+ /* log and continue: aborting here would leak every remaining comm below */
983	ZBAL_LOG_WARN("~ process group: " << groupName_ << " on rank " << myWorldRank_ << " result " << result);	984	ZBAL_LOG_WARN("~ process group: " << groupName_ << " on rank " << myWorldRank_ << " result " << result);
984	- return;
985	}	985	}
986	groupComm_ = nullptr;	986	groupComm_ = nullptr;
987	ZBAL_LOG_DEBUG("~ process group success: " << groupName_ << " on rank " << myWorldRank_);	987	ZBAL_LOG_DEBUG("~ process group success: " << groupName_ << " on rank " << myWorldRank_);
	@@ -990,10 +990,12 @@ ProcessGroupZBAL::~ProcessGroupZBAL()
990	for (auto &[groupName, groupComm] : groupP2pComms_) {	990	for (auto &[groupName, groupComm] : groupP2pComms_) {
991	auto result = zbal_comm_destroy(groupComm, 0);	991	auto result = zbal_comm_destroy(groupComm, 0);
992	if (result != Z_OK) {	992	if (result != Z_OK) {
		993	+ /* keep destroying the rest, destructor is the last cleanup chance
		994	+ * and partial abort would leak all remaining comms */
993	ZBAL_LOG_WARN("~ process group: " << groupName << " on rank " << myWorldRank_ << " result " << result);	995	ZBAL_LOG_WARN("~ process group: " << groupName << " on rank " << myWorldRank_ << " result " << result);
994	- return;	996	+ } else {
		997	+ ZBAL_LOG_DEBUG("~ process group success: " << groupName << " on rank " << myWorldRank_);
995	}	998	}
996	- ZBAL_LOG_DEBUG("~ process group success: " << groupName << " on rank " << myWorldRank_);
997	}	999	}
998	groupP2pComms_.clear();	1000	groupP2pComms_.clear();
999	}	1001	}

Mapp/zbal/src/csrc/bootstrap/memory/memfabric/zbal_mem_mf_bootstrap.cpp+6-1

	@@ -155,7 +155,12 @@ ZResult MemFabricBoostrap::CreateSHMSpace() noexcept
155	return Z_INIT_BOOTSTRAP_FAILED;	155	return Z_INIT_BOOTSTRAP_FAILED;
156	}	156	}
157		157
158	- ZBAL_ASSERT_RETURN(options_.totalMemSize <= spaceSize, Z_ERROR);	158	+ if (options_.totalMemSize > spaceSize) {
		159	+ ZBAL_LOG_AND_SET_LAST_ERROR("Requested memSize exceeds symmetric space, request=" << options_.totalMemSize
		160	+ << ", space=" << spaceSize);
		161	+ DlMfApi::SmemShmDestroy(tmpShmHandle, 0);
		162	+ return Z_ERROR;
		163	+ }
159	ZBAL_LOG_DEBUG("Space size: " << spaceSize << ", actual device memory size: " << options_.totalMemSize);	164	ZBAL_LOG_DEBUG("Space size: " << spaceSize << ", actual device memory size: " << options_.totalMemSize);
160		165
161	/* get local GVA */	166	/* get local GVA */

Mapp/zbal/src/csrc/operators/npu/aicore/ops/fused_deep_moe/raw_distributed/zbal_moe_distribute_combine.h+7-6

	@@ -196,8 +196,8 @@ private:
196	uint32_t firstTpTokenEndIdx_{0};	196	uint32_t firstTpTokenEndIdx_{0};
197	uint32_t firstTpTokenEndOffset_{0};	197	uint32_t firstTpTokenEndOffset_{0};
198	uint32_t endTok_{0};	198	uint32_t endTok_{0};
199	- uint32_t epDataOffsetOnWin_{0};	199	+ uint64_t epDataOffsetOnWin_{0};
200	- uint32_t tpDataOffsetOnWin_{0};	200	+ uint64_t tpDataOffsetOnWin_{0};
201	uint32_t epStateOffsetOnWin_{0};	201	uint32_t epStateOffsetOnWin_{0};
202	uint32_t tpStateOffsetOnWin_{0};	202	uint32_t tpStateOffsetOnWin_{0};
203	uint32_t axisHFloatSize_{0};	203	uint32_t axisHFloatSize_{0};
	@@ -307,11 +307,12 @@ __aicore__ inline void CamMoeDistributeCombine<TemplateMC2TypeFunc>::Init(
307	epWindowGM_ = combineWs; // local combine workspace replaces data window	307	epWindowGM_ = combineWs; // local combine workspace replaces data window
308	epStatusSpaceGm_ = GetWinStateAddrByRankId(epRankId_, EP_DOMAIN);	308	epStatusSpaceGm_ = GetWinStateAddrByRankId(epRankId_, EP_DOMAIN);
309	epStatusSpaceGlobalTensor_.SetGlobalBuffer((__gm__ float *)epStatusSpaceGm_);	309	epStatusSpaceGlobalTensor_.SetGlobalBuffer((__gm__ float *)epStatusSpaceGm_);
310	- epDataOffsetOnWin_ = epRankId_ * moeExpertPerRankNum_ * static_cast<uint32_t>(expertPerSizeOnWin_);	310	+ epDataOffsetOnWin_ =
		311	+ static_cast<uint64_t>(epRankId_) * moeExpertPerRankNum_ * expertPerSizeOnWin_; /* keep uint64, no truncation */
311	epStateOffsetOnWin_ = epRankId_ * stateOffset_;	312	epStateOffsetOnWin_ = epRankId_ * stateOffset_;
312	- axisHFloatSize_ = axisH_ * sizeof(float);	313	+ axisHFloatSize_ = static_cast<uint64_t>(axisH_) * sizeof(float);
313	- axisHExpandXTypeSize_ = axisH_ * sizeof(ExpandXType);	314	+ axisHExpandXTypeSize_ = static_cast<uint64_t>(axisH_) * sizeof(ExpandXType);
314	- bsKNum_ = axisBS_ * axisK_;	315	+ bsKNum_ = static_cast<uint64_t>(axisBS_) * axisK_;
315		316
316	if constexpr (IsNeedReduceScatter) {	317	if constexpr (IsNeedReduceScatter) {
317	tpSendCountGM_.SetGlobalBuffer((__gm__ int32_t *)tpSendCount);	318	tpSendCountGM_.SetGlobalBuffer((__gm__ int32_t *)tpSendCount);

Mapp/zbal/src/csrc/operators/npu/zbal_npu_aicpu_launcher.cpp+4-0

	@@ -105,6 +105,10 @@ ZResult NpuAicpuLauncher::Launch(const AicpuWorkDesc &desc, void *stream)
105	/* Select per-op function handle for profiling visibility */	105	/* Select per-op function handle for profiling visibility */
106	uint32_t commIdx = param.commType < kMaxCommType ? param.commType : 0;	106	uint32_t commIdx = param.commType < kMaxCommType ? param.commType : 0;
107	aclrtFuncHandle launchHandle = opFuncHandles_[commIdx];	107	aclrtFuncHandle launchHandle = opFuncHandles_[commIdx];
		108	+ if (launchHandle == nullptr) {
		109	+ ZBAL_LOG_ERROR("AICPU Launch: no function handle for commType=" << param.commType << ", commIdx=" << commIdx);
		110	+ return Z_INVALID_PARAM;
		111	+ }
108		112
109	/* V2 API: single call, passes host buffer directly (no ArgsInit/Append/Finalize) */	113	/* V2 API: single call, passes host buffer directly (no ArgsInit/Append/Finalize) */
110	int32_t aclRet =	114	int32_t aclRet =

Mapp/zbal/src/csrc/operators/npu/zbal_npu_communicator_aicpu.h+8-0

	@@ -22,6 +22,14 @@ class NpuCommunicatorAICPU : public NpuCommunicatorBase {
22	public:	22	public:
23	using NpuCommunicatorBase::NpuCommunicatorBase;	23	using NpuCommunicatorBase::NpuCommunicatorBase;
24		24
		25	+ ~NpuCommunicatorAICPU() override
		26	+ {
		27	+ /* Must call here: virtual dispatch from ~NpuCommunicatorBase resolves to the
		28	+ * base version, skipping AICPU resource cleanup (workspace/binary/cclBuf).
		29	+ * Idempotent when UnInitialize was already called explicitly. */
		30	+ UnInitialize();
		31	+ }
		32	+
25	ZResult Initialize() noexcept override;	33	ZResult Initialize() noexcept override;
26	void UnInitialize() noexcept override;	34	void UnInitialize() noexcept override;
27		35

Mapp/zbal/src/csrc/operators/npu/zbal_npu_communicator_base.cpp+5-0

	@@ -215,6 +215,11 @@ void NpuCommunicatorBase::DestroyProfMemory()
215	ZBAL_LOG_WARN("UnRegister host memory failed, result: " << result);	215	ZBAL_LOG_WARN("UnRegister host memory failed, result: " << result);
216	}	216	}
217		217
		218	+ result = DlCannApi::AclrtFreeHost(perfHostMemory_);
		219	+ if (result != Z_OK) {
		220	+ ZBAL_LOG_WARN("Free host memory failed, result: " << result);
		221	+ }
		222	+
218	perfHostMemory_ = nullptr;	223	perfHostMemory_ = nullptr;
219	}	224	}
220		225

Mapp/zbal/src/csrc/operators/zbal_communicator.cpp+13-1

	@@ -75,7 +75,14 @@ ZResult Communicator::Create(const zbal_comm_options_t &options, zbal_comm_t *co
75	}	75	}
76		76
77	result = commInner->AssignGatherGroupId(tmpGroupId);	77	result = commInner->AssignGatherGroupId(tmpGroupId);
78	- ZBAL_VALIDATE_RETURN(result == Z_OK, "Assign group id failed, result: " << result, result);	78	+ if (result != Z_OK) {
		79	+ ZBAL_LOG_AND_SET_LAST_ERROR("Assign group id failed, result: " << result);
		80	+ /* roll back: erase from lookup maps, release group id and device resources,
		81	+ * otherwise the half-created comm blocks re-create with the same name and
		82	+ * leaks the cross-rank group id */
		83	+ (void)DestroyInner(commInner);
		84	+ return result;
		85	+ }
79		86
80	*comm = commInner.Get();	87	*comm = commInner.Get();
81		88
	@@ -224,6 +231,9 @@ CommunicatorPtr Communicator::CreateInner(zbal_backend_t backendType, const Comm
224	comm->ConstructCommGroupInfo(options);	231	comm->ConstructCommGroupInfo(options);
225	if (comm->Initialize()) {	232	if (comm->Initialize()) {
226	ZBAL_LOG_AND_SET_LAST_ERROR("Initialize communicator failed");	233	ZBAL_LOG_AND_SET_LAST_ERROR("Initialize communicator failed");
		234	+ /* drop the creator's reference: local ZRef's DecreaseRef then triggers
		235	+ * delete (refCount underflows to delete), releasing all resources */
		236	+ comm->DecreaseRef();
227	return nullptr;	237	return nullptr;
228	}	238	}
229		239
	@@ -243,6 +253,7 @@ CommunicatorPtr Communicator::CreateInner(zbal_backend_t backendType, const Comm
243	* return nullptr directly as its already created	253	* return nullptr directly as its already created
244	*/	254	*/
245	ZBAL_LOG_AND_SET_LAST_ERROR("Create communicator failed as world group already created");	255	ZBAL_LOG_AND_SET_LAST_ERROR("Create communicator failed as world group already created");
		256	+ comm->DecreaseRef();
246	return nullptr;	257	return nullptr;
247	} else if (!isWorldGroup && gWorldCommunicator == nullptr) {	258	} else if (!isWorldGroup && gWorldCommunicator == nullptr) {
248	/*	259	/*
	@@ -251,6 +262,7 @@ CommunicatorPtr Communicator::CreateInner(zbal_backend_t backendType, const Comm
251	* return nullptr	262	* return nullptr
252	*/	263	*/
253	ZBAL_LOG_AND_SET_LAST_ERROR("Create communicator failed as world group not created");	264	ZBAL_LOG_AND_SET_LAST_ERROR("Create communicator failed as world group not created");
		265	+ comm->DecreaseRef();
254	return nullptr;	266	return nullptr;
255	} else {	267	} else {
256	/*	268	/*