diff --git a/ggml/src/ggml-metal/ggml-metal-context.m b/ggml/src/ggml-metal/ggml-metal-context.m index e9064666f..597c126ca 100644 --- a/ggml/src/ggml-metal/ggml-metal-context.m +++ b/ggml/src/ggml-metal/ggml-metal-context.m @@ -539,8 +539,14 @@ enum ggml_status ggml_metal_graph_compute(ggml_metal_t ctx, struct ggml_cgraph * descriptor.outputURL = [NSURL fileURLWithPath:path]; NSError * error = nil; - if (![[MTLCaptureManager sharedCaptureManager] startCaptureWithDescriptor:descriptor error:&error]) { + const bool ok = [[MTLCaptureManager sharedCaptureManager] startCaptureWithDescriptor:descriptor error:&error]; + + [descriptor release]; + + if (!ok) { GGML_LOG_ERROR("%s: error: unable to start capture '%s' (did you set METAL_CAPTURE_ENABLED=1 ?)\n", __func__, [[error localizedDescription] UTF8String]); + [ctx->capture_scope release]; + ctx->capture_scope = nil; } else { [ctx->capture_scope beginScope]; ctx->capture_started = true; @@ -600,6 +606,8 @@ enum ggml_status ggml_metal_graph_compute(ggml_metal_t ctx, struct ggml_cgraph * // enter here only when capturing in order to wait for all computation to finish // otherwise, we leave the graph to compute asynchronously if (use_capture && ctx->capture_started) { + enum ggml_status res = GGML_STATUS_SUCCESS; + // wait for completion and check status of each command buffer // needed to detect if the device ran out-of-memory for example (#1881) { @@ -613,11 +621,11 @@ enum ggml_status ggml_metal_graph_compute(ggml_metal_t ctx, struct ggml_cgraph * GGML_LOG_INFO("error: %s\n", [[cmd_buf error].localizedDescription UTF8String]); } - return GGML_STATUS_FAILED; + res = GGML_STATUS_FAILED; } } - for (int i = 0; i < n_cb; ++i) { + for (int i = 0; i < n_cb && res == GGML_STATUS_SUCCESS; ++i) { id cmd_buf = ctx->cmd_bufs[i].obj; [cmd_buf waitUntilCompleted]; @@ -627,8 +635,8 @@ enum ggml_status ggml_metal_graph_compute(ggml_metal_t ctx, struct ggml_cgraph * if (status == MTLCommandBufferStatusError) { GGML_LOG_INFO("error: %s\n", [[cmd_buf error].localizedDescription UTF8String]); } - - return GGML_STATUS_FAILED; + res = GGML_STATUS_FAILED; + break; } id next_buffer = (i + 1 < n_cb ? ctx->cmd_bufs[i + 1].obj : nil); @@ -643,7 +651,8 @@ enum ggml_status ggml_metal_graph_compute(ggml_metal_t ctx, struct ggml_cgraph * if (ctx->abort_callback && ctx->abort_callback(ctx->abort_callback_data)) { GGML_LOG_INFO("%s: command buffer %d aborted", __func__, i); - return GGML_STATUS_ABORTED; + res = GGML_STATUS_ABORTED; + break; } [next_buffer commit]; @@ -652,7 +661,14 @@ enum ggml_status ggml_metal_graph_compute(ggml_metal_t ctx, struct ggml_cgraph * [ctx->capture_scope endScope]; [[MTLCaptureManager sharedCaptureManager] stopCapture]; + [ctx->capture_scope release]; + ctx->capture_scope = nil; + ctx->capture_started = false; + + if (res != GGML_STATUS_SUCCESS) { + return res; + } } }