diff --git a/WORKSPACE b/WORKSPACE index b92252e759c..1c7a45f8c48 100644 --- a/WORKSPACE +++ b/WORKSPACE @@ -5,9 +5,9 @@ load("@bazel_tools//tools/build_defs/repo:http.bzl", "http_archive") git_repository( name = "ccv", - commit = "d2622a324f45c243c69fe8c2019eeaca9162ca0a", + commit = "ef7bd537aea035ab04e4f7d0d817f56f7b0492fc", remote = "https://github.com/liuliu/ccv.git", - shallow_since = "1726462252 -0400", + shallow_since = "1727501268 -0400", ) load("@ccv//config:ccv.bzl", "ccv_deps", "ccv_setting") diff --git a/deps.bzl b/deps.bzl index fa44d795ffe..1ce124f1e45 100644 --- a/deps.bzl +++ b/deps.bzl @@ -17,8 +17,8 @@ def s4nnc_deps(): git_repository, name = "ccv", remote = "https://github.com/liuliu/ccv.git", - commit = "d2622a324f45c243c69fe8c2019eeaca9162ca0a", - shallow_since = "1726462252 -0400", + commit = "ef7bd537aea035ab04e4f7d0d817f56f7b0492fc", + shallow_since = "1727501268 -0400", ) _maybe( diff --git a/nnc/ModelAddons.swift b/nnc/ModelAddons.swift index 6ed90e75bbe..a78b4a0dfc0 100644 --- a/nnc/ModelAddons.swift +++ b/nnc/ModelAddons.swift @@ -1498,6 +1498,40 @@ public final class ScaledDotProductAttention: Model { } } +/// Debug model. +public final class Debug: Model { + required init(_ model: OpaquePointer) { + super.init(model) + } + + public init( + name: String = "", + _ callback: @escaping ([AnyTensor?], StreamContext?) -> Void + ) { + let underlying = Wrapped(callback) + super.init( + ccv_cnnp_debug( + { inputs, inputSize, stream, context in + let callback = Unmanaged Void>>.fromOpaque( + context! + ).takeUnretainedValue().value + callback( + (0.. Void>>.fromOpaque(context).release() + }, + { context in + guard let context = context else { return nil } + return Unmanaged Void>>.fromOpaque(context) + .retain().toOpaque() + }, name)) + } +} + /// Custom model. public final class CustomModel: Model { required init(_ model: OpaquePointer) { diff --git a/nnc/StreamContext.swift b/nnc/StreamContext.swift index 0f1182d6699..8a61493deaa 100644 --- a/nnc/StreamContext.swift +++ b/nnc/StreamContext.swift @@ -23,8 +23,14 @@ public final class StreamContext { } } + let selfOwned: Bool let _stream: OpaquePointer + init(stream: OpaquePointer, selfOwned: Bool) { + _stream = stream + self.selfOwned = selfOwned + } + /** * Create a new stream context. * @@ -39,6 +45,7 @@ public final class StreamContext { type = Int32((ordinal << 8) | CCV_STREAM_CONTEXT_GPU) } _stream = ccv_nnc_stream_context_new(type)! + selfOwned = true } /** @@ -69,6 +76,7 @@ public final class StreamContext { } deinit { + guard selfOwned else { return } ccv_nnc_stream_context_free(_stream) } } diff --git a/test/model.swift b/test/model.swift index d683ceef1d5..2d63377d577 100644 --- a/test/model.swift +++ b/test/model.swift @@ -171,6 +171,33 @@ final class ModelTests: XCTestCase { XCTAssertEqual(tv3s[1].rawValue[0], 0.5 / 0.2, accuracy: 1e-5) } + func testModelDebug() throws { + let dynamicGraph = DynamicGraph() + + var value: Float = 0 + func DivRec() -> Model { + let i0 = Input() + let i1 = Input() + let i2 = i0 ./ i1 + let debugi2 = + (Debug { tensors, _ in + let tensor = Tensor(tensors[0]!) + value = tensor[0] + })(i2) + let i3 = Input() + let i4 = debugi2 ./ i3 + return Model([i0, i1, i3], [i4]) + } + + let div = DivRec() + let tv0 = dynamicGraph.variable(Tensor([1.1], .CPU, .C(1))) + let tv1 = dynamicGraph.variable(Tensor([2.2], .CPU, .C(1))) + let tv2 = dynamicGraph.variable(Tensor([0.2], .CPU, .C(1))) + let tv3s = div(inputs: tv0, tv1, tv2).map { $0.as(of: Float32.self) } + XCTAssertEqual(tv3s[0].rawValue[0], 0.5 / 0.2, accuracy: 1e-5) + XCTAssertEqual(value, 0.5, accuracy: 1e-5) + } + func testModelScaledDotProductAttention() throws { let dynamicGraph = DynamicGraph() let q = dynamicGraph.variable(Tensor(.CPU, .NHWC(1, 10, 8, 20))) @@ -458,6 +485,7 @@ final class ModelTests: XCTestCase { ("testModelWithScalar", testModelWithScalar), ("testModelWithParameter", testModelWithParameter), ("testModelDiv", testModelDiv), + ("testModelDebug", testModelDebug), ("testModelScaledDotProductAttention", testModelScaledDotProductAttention), ("testCustomModel", testCustomModel), ("testModelShareWeights", testModelShareWeights),