Skip to content

Commit b711387

Browse files
committed
CANN: Fix ggml_cann_set_device to avoid redundant device switches
- Added a check to skip aclrtSetDevice if the current device is already set. - Prevents unnecessary context switches while keeping thread/device consistency.
1 parent 704d90c commit b711387

File tree

2 files changed

+9
-7
lines changed

2 files changed

+9
-7
lines changed

ggml/src/ggml-cann/common.h

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -526,7 +526,10 @@ struct ggml_backend_cann_context {
526526
*/
527527
aclrtStream stream(int stream) {
528528
if (streams[stream] == nullptr) {
529-
ggml_cann_set_device(device);
529+
// If the device is not set here, destroying the stream later may cause a mismatch
530+
// between the thread contexts where the stream was created and destroyed.
531+
// However, I printed the device_id, thread_id, and stream, and they are all consistent.
532+
ACL_CHECK(aclrtSetDevice(device));
530533
ACL_CHECK(aclrtCreateStream(&streams[stream]));
531534
}
532535
return streams[stream];

ggml/src/ggml-cann/ggml-cann.cpp

Lines changed: 5 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -75,13 +75,12 @@
7575
* @param device The device ID to set.
7676
*/
7777
void ggml_cann_set_device(const int32_t device) {
78-
// TODO: uncomment these lines after empty context has fixed.
79-
// int current_device;
80-
// ACL_CHECK(aclrtGetDevice(&current_device));
78+
int current_device;
79+
aclrtGetDevice(&current_device);
8180

82-
// if (device == current_device) {
83-
// return;
84-
// }
81+
if (device == current_device) {
82+
return;
83+
}
8584
ACL_CHECK(aclrtSetDevice(device));
8685
}
8786

0 commit comments

Comments
 (0)