Skip to content

Commit c9dcafa

Browse files
committed
cann: implement backend memset_tensor interface
Add ggml_backend_cann_buffer_memset_tensor and wire it into `ggml_backend_cann_buffer_interface`. This ensures backend tensor memset operations are supported and avoids incorrect behavior when tensors need explicit zero-initialization (e.g. cache buffers).
1 parent 6861f65 commit c9dcafa

1 file changed

Lines changed: 17 additions & 1 deletion

File tree

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

Lines changed: 17 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1327,6 +1327,22 @@ static bool ggml_backend_cann_buffer_cpy_tensor(ggml_backend_buffer_t buffer,
13271327
return false;
13281328
}
13291329

1330+
/**
1331+
* @brief Set a region of a tensor's device memory to a specified value.
1332+
*
1333+
* @param buffer The CANN buffer containing the tensor.
1334+
* @param tensor Pointer to the tensor whose memory will be set.
1335+
* @param value The value to which each byte in the region will be set.
1336+
* @param offset Byte offset within the tensor's data to start setting.
1337+
* @param size Number of bytes to set.
1338+
*/
1339+
static void ggml_backend_cann_buffer_memset_tensor(ggml_backend_buffer_t buffer, ggml_tensor * tensor, uint8_t value, size_t offset, size_t size) {
1340+
ggml_backend_cann_buffer_context * ctx = (ggml_backend_cann_buffer_context *) buffer->context;
1341+
1342+
ggml_cann_set_device(ctx->device);
1343+
ACL_CHECK(aclrtMemset((char *) tensor->data + offset, size, value, size));
1344+
}
1345+
13301346
/**
13311347
* @brief Clear a CANN buffer by setting all its memory to a specified value.
13321348
*
@@ -1353,7 +1369,7 @@ static const ggml_backend_buffer_i ggml_backend_cann_buffer_interface = {
13531369
/* .free_buffer = */ ggml_backend_cann_buffer_free_buffer,
13541370
/* .get_base = */ ggml_backend_cann_buffer_get_base,
13551371
/* .init_tensor = */ ggml_backend_cann_buffer_init_tensor,
1356-
/* .memset_tensor = */ NULL,
1372+
/* .memset_tensor = */ ggml_backend_cann_buffer_memset_tensor,
13571373
/* .set_tensor = */ ggml_backend_cann_buffer_set_tensor,
13581374
/* .get_tensor = */ ggml_backend_cann_buffer_get_tensor,
13591375
/* .cpy_tensor = */ ggml_backend_cann_buffer_cpy_tensor,

0 commit comments

Comments
 (0)