Skip to content

Commit d5fabe3

Browse files
authored
CANN: Optimize ggml_cann_set_device (ggml-org#15935)
* 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. * CANN: add device default id
1 parent 8ff2060 commit d5fabe3

File tree

2 files changed

+10
-7
lines changed

2 files changed

+10
-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: 6 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 = -1;
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

@@ -1729,6 +1728,7 @@ static bool ggml_cann_compute_forward(ggml_backend_cann_context& ctx,
17291728
ggml_cann_get_rows(ctx, dst);
17301729
break;
17311730
case GGML_OP_SET_ROWS:
1731+
std::cout << "lcg GGML_OP_SET_ROWS"<< std::endl;
17321732
ggml_cann_set_rows(ctx, dst);
17331733
break;
17341734
case GGML_OP_DUP:

0 commit comments

Comments
 (0)