Skip to content

Commit b941c3c

Browse files
stikvesstikves
authored
Improve Swift lifetime handling of new states
## Summary Fixes two ownership regressions introduced by #132, and removes dead code 1. `StateHandler+MTLBuffer.swift` — `insertAll(into: inout AsyncMutableViews)` modifies structures outside of their compiler known lifetime. Removed the helper; state bindings are now inlined at each call site with all values in a flat scope. 2. `CoreAISequentialEngine.swift` — Dead methods (`ensureKVCapacity`, `copyCache`, `zeroFill`) reference properties removed by the state handler refactor. Deleted.
1 parent 6329412 commit b941c3c

4 files changed

Lines changed: 193 additions & 142 deletions

File tree

‎swift/Sources/CoreAILanguageModels/Handlers/StateHandler+MTLBuffer.swift‎

Lines changed: 0 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -48,16 +48,6 @@ public struct FixedMTLBufferState {
4848
get { bindings[index] }
4949
}
5050

51-
/// Insert all states into AsyncMutableViews for pipelined encoding.
52-
public func insertAll(into views: inout InferenceFunction.AsyncMutableViews) {
53-
for (name, buffer, scalarType, shape, strides) in bindings {
54-
var value = unsafe InferenceFunction.AsyncMutableValue(
55-
unsafeBuffer: buffer, byteOffset: 0,
56-
scalarType: scalarType, shape: shape, strides: strides)
57-
views.insert(&value, for: name)
58-
}
59-
}
60-
6151
/// Zero all state buffers. Caller must ensure no in-flight GPU work references these.
6252
public mutating func reset() {
6353
for (_, buffer, _, _, _) in bindings {

‎swift/Sources/CoreAILanguageModels/Handlers/StateHandler+NDArray.swift‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -176,7 +176,7 @@ public struct GrowingNDArrayState: SyncStateHandler {
176176
// MARK: - Shared Utilities
177177

178178
/// Zero-initialize an NDArray, dispatching on scalar type.
179-
private func zeroFillNDArray(_ array: inout NDArray) {
179+
func zeroFillNDArray(_ array: inout NDArray) {
180180
let count = array.shape.reduce(1, *)
181181
switch array.scalarType {
182182
case .float16, .bfloat16:

‎swift/Sources/CoreAILanguageModels/InferenceEngines/CoreAIPipelinedEngine.swift‎

Lines changed: 192 additions & 54 deletions
Original file line numberDiff line numberDiff line change
@@ -826,26 +826,11 @@ private struct EngineImpl: ~Copyable {
826826
strides: valStrides
827827
)
828828

829-
var asyncStates = InferenceFunction.AsyncMutableViews()
830-
asyncStates.insert(&keyState, for: keyCacheName)
831-
asyncStates.insert(&valState, for: valueCacheName)
832-
additionalStates?.insertAll(into: &asyncStates)
833-
834829
// Build Output as AsyncMutableValue (logits)
835830
// Decode uses per-step rotating buffer; prefill uses the shared growing buffer.
836831
let logitsOutputBuffer = tokens.isEmpty ? decodeLogitsBuffers[step % pipelineDepth] : logits.metalBuffer
837832
let logitsShape = [1, queryLength, vocabSize]
838833
let logitsStrides = try resolvedStrides(descriptor: logitsBaseDesc, shape: logitsShape)
839-
var logitsOutput = unsafe InferenceFunction.AsyncMutableValue(
840-
unsafeBuffer: logitsOutputBuffer,
841-
byteOffset: 0,
842-
scalarType: .float16,
843-
shape: logitsShape,
844-
strides: logitsStrides
845-
)
846-
847-
var asyncOutputs = InferenceFunction.AsyncMutableViews()
848-
asyncOutputs.insert(&logitsOutput, for: logitsOutputName)
849834

850835
prepareSpan.end()
851836

@@ -856,12 +841,73 @@ private struct EngineImpl: ~Copyable {
856841
// This commits + uses runAfterSyncPoint (no stream wait) — enables true pipelining.
857842
let logitsSpan = InstrumentsProfiler.beginLogitsInference(
858843
step: currentStep, tokens: queryLength, engine: "CoreAI-Pipelined")
859-
let _ = try function.encode(
860-
inputs: asyncInputs,
861-
states: consume asyncStates,
862-
outputViews: consume asyncOutputs,
863-
to: computeStream
864-
)
844+
845+
// Swift 6 lifetime safety: AsyncMutableViews uses @lifetime(self: &mutableValue)
846+
// on insert(), so all inserts + consume must be in the same scope without branching.
847+
// We switch on state count to avoid the optional-chain branch that the checker rejects.
848+
switch additionalStates?.stateCount ?? 0 {
849+
case 1:
850+
let extra0 = additionalStates![stateIndex: 0]
851+
var extraState0 = unsafe InferenceFunction.AsyncMutableValue(
852+
unsafeBuffer: extra0.buffer, byteOffset: 0,
853+
scalarType: extra0.scalarType, shape: extra0.shape, strides: extra0.strides)
854+
var asyncStates = InferenceFunction.AsyncMutableViews()
855+
asyncStates.insert(&keyState, for: keyCacheName)
856+
asyncStates.insert(&valState, for: valueCacheName)
857+
asyncStates.insert(&extraState0, for: extra0.name)
858+
var logitsOutput = unsafe InferenceFunction.AsyncMutableValue(
859+
unsafeBuffer: logitsOutputBuffer, byteOffset: 0,
860+
scalarType: .float16, shape: logitsShape, strides: logitsStrides)
861+
var asyncOutputs = InferenceFunction.AsyncMutableViews()
862+
asyncOutputs.insert(&logitsOutput, for: logitsOutputName)
863+
let _ = try function.encode(
864+
inputs: asyncInputs,
865+
states: consume asyncStates,
866+
outputViews: consume asyncOutputs,
867+
to: computeStream
868+
)
869+
case 2:
870+
let extra0 = additionalStates![stateIndex: 0]
871+
let extra1 = additionalStates![stateIndex: 1]
872+
var extraState0 = unsafe InferenceFunction.AsyncMutableValue(
873+
unsafeBuffer: extra0.buffer, byteOffset: 0,
874+
scalarType: extra0.scalarType, shape: extra0.shape, strides: extra0.strides)
875+
var extraState1 = unsafe InferenceFunction.AsyncMutableValue(
876+
unsafeBuffer: extra1.buffer, byteOffset: 0,
877+
scalarType: extra1.scalarType, shape: extra1.shape, strides: extra1.strides)
878+
var asyncStates = InferenceFunction.AsyncMutableViews()
879+
asyncStates.insert(&keyState, for: keyCacheName)
880+
asyncStates.insert(&valState, for: valueCacheName)
881+
asyncStates.insert(&extraState0, for: extra0.name)
882+
asyncStates.insert(&extraState1, for: extra1.name)
883+
var logitsOutput = unsafe InferenceFunction.AsyncMutableValue(
884+
unsafeBuffer: logitsOutputBuffer, byteOffset: 0,
885+
scalarType: .float16, shape: logitsShape, strides: logitsStrides)
886+
var asyncOutputs = InferenceFunction.AsyncMutableViews()
887+
asyncOutputs.insert(&logitsOutput, for: logitsOutputName)
888+
let _ = try function.encode(
889+
inputs: asyncInputs,
890+
states: consume asyncStates,
891+
outputViews: consume asyncOutputs,
892+
to: computeStream
893+
)
894+
default:
895+
// case 0: no additional states
896+
var asyncStates = InferenceFunction.AsyncMutableViews()
897+
asyncStates.insert(&keyState, for: keyCacheName)
898+
asyncStates.insert(&valState, for: valueCacheName)
899+
var logitsOutput = unsafe InferenceFunction.AsyncMutableValue(
900+
unsafeBuffer: logitsOutputBuffer, byteOffset: 0,
901+
scalarType: .float16, shape: logitsShape, strides: logitsStrides)
902+
var asyncOutputs = InferenceFunction.AsyncMutableViews()
903+
asyncOutputs.insert(&logitsOutput, for: logitsOutputName)
904+
let _ = try function.encode(
905+
inputs: asyncInputs,
906+
states: consume asyncStates,
907+
outputViews: consume asyncOutputs,
908+
to: computeStream
909+
)
910+
}
865911
logitsSpan.end()
866912

867913
// GPU sampling via Metal queue
@@ -1138,25 +1184,71 @@ private struct EngineImpl: ~Copyable {
11381184
var valState = unsafe InferenceFunction.AsyncMutableValue(
11391185
unsafeBuffer: valBuffer, byteOffset: 0,
11401186
scalarType: valueCacheScalarType, shape: valShape, strides: valStrides)
1141-
var asyncStates = InferenceFunction.AsyncMutableViews()
1142-
asyncStates.insert(&keyState, for: keyCacheName)
1143-
asyncStates.insert(&valState, for: valueCacheName)
1144-
additionalStates?.insertAll(into: &asyncStates)
1145-
11461187
let logitsShape = [1, queryLength, vocabSize]
11471188
let logitsStrides = try resolvedStrides(descriptor: logitsBaseDesc, shape: logitsShape)
1148-
var logitsOutput = unsafe InferenceFunction.AsyncMutableValue(
1149-
unsafeBuffer: logits.metalBuffer, byteOffset: 0,
1150-
scalarType: .float16, shape: logitsShape, strides: logitsStrides)
1151-
var asyncOutputs = InferenceFunction.AsyncMutableViews()
1152-
asyncOutputs.insert(&logitsOutput, for: logitsOutputName)
1153-
1154-
let _ = try function.encode(
1155-
inputs: asyncInputs,
1156-
states: consume asyncStates,
1157-
outputViews: consume asyncOutputs,
1158-
to: computeStream
1159-
)
1189+
1190+
switch additionalStates?.stateCount ?? 0 {
1191+
case 1:
1192+
let extra0 = additionalStates![stateIndex: 0]
1193+
var extraState0 = unsafe InferenceFunction.AsyncMutableValue(
1194+
unsafeBuffer: extra0.buffer, byteOffset: 0,
1195+
scalarType: extra0.scalarType, shape: extra0.shape, strides: extra0.strides)
1196+
var asyncStates = InferenceFunction.AsyncMutableViews()
1197+
asyncStates.insert(&keyState, for: keyCacheName)
1198+
asyncStates.insert(&valState, for: valueCacheName)
1199+
asyncStates.insert(&extraState0, for: extra0.name)
1200+
var logitsOutput = unsafe InferenceFunction.AsyncMutableValue(
1201+
unsafeBuffer: logits.metalBuffer, byteOffset: 0,
1202+
scalarType: .float16, shape: logitsShape, strides: logitsStrides)
1203+
var asyncOutputs = InferenceFunction.AsyncMutableViews()
1204+
asyncOutputs.insert(&logitsOutput, for: logitsOutputName)
1205+
let _ = try function.encode(
1206+
inputs: asyncInputs,
1207+
states: consume asyncStates,
1208+
outputViews: consume asyncOutputs,
1209+
to: computeStream
1210+
)
1211+
case 2:
1212+
let extra0 = additionalStates![stateIndex: 0]
1213+
let extra1 = additionalStates![stateIndex: 1]
1214+
var extraState0 = unsafe InferenceFunction.AsyncMutableValue(
1215+
unsafeBuffer: extra0.buffer, byteOffset: 0,
1216+
scalarType: extra0.scalarType, shape: extra0.shape, strides: extra0.strides)
1217+
var extraState1 = unsafe InferenceFunction.AsyncMutableValue(
1218+
unsafeBuffer: extra1.buffer, byteOffset: 0,
1219+
scalarType: extra1.scalarType, shape: extra1.shape, strides: extra1.strides)
1220+
var asyncStates = InferenceFunction.AsyncMutableViews()
1221+
asyncStates.insert(&keyState, for: keyCacheName)
1222+
asyncStates.insert(&valState, for: valueCacheName)
1223+
asyncStates.insert(&extraState0, for: extra0.name)
1224+
asyncStates.insert(&extraState1, for: extra1.name)
1225+
var logitsOutput = unsafe InferenceFunction.AsyncMutableValue(
1226+
unsafeBuffer: logits.metalBuffer, byteOffset: 0,
1227+
scalarType: .float16, shape: logitsShape, strides: logitsStrides)
1228+
var asyncOutputs = InferenceFunction.AsyncMutableViews()
1229+
asyncOutputs.insert(&logitsOutput, for: logitsOutputName)
1230+
let _ = try function.encode(
1231+
inputs: asyncInputs,
1232+
states: consume asyncStates,
1233+
outputViews: consume asyncOutputs,
1234+
to: computeStream
1235+
)
1236+
default:
1237+
var asyncStates = InferenceFunction.AsyncMutableViews()
1238+
asyncStates.insert(&keyState, for: keyCacheName)
1239+
asyncStates.insert(&valState, for: valueCacheName)
1240+
var logitsOutput = unsafe InferenceFunction.AsyncMutableValue(
1241+
unsafeBuffer: logits.metalBuffer, byteOffset: 0,
1242+
scalarType: .float16, shape: logitsShape, strides: logitsStrides)
1243+
var asyncOutputs = InferenceFunction.AsyncMutableViews()
1244+
asyncOutputs.insert(&logitsOutput, for: logitsOutputName)
1245+
let _ = try function.encode(
1246+
inputs: asyncInputs,
1247+
states: consume asyncStates,
1248+
outputViews: consume asyncOutputs,
1249+
to: computeStream
1250+
)
1251+
}
11601252

11611253
processedTokenCount += queryLength
11621254
step += 1
@@ -1245,25 +1337,71 @@ private struct EngineImpl: ~Copyable {
12451337
var valState = unsafe InferenceFunction.AsyncMutableValue(
12461338
unsafeBuffer: valBuffer, byteOffset: 0,
12471339
scalarType: valueCacheScalarType, shape: vShape, strides: vStrides)
1248-
var asyncStates = InferenceFunction.AsyncMutableViews()
1249-
asyncStates.insert(&keyState, for: keyCacheName)
1250-
asyncStates.insert(&valState, for: valueCacheName)
1251-
additionalStates?.insertAll(into: &asyncStates)
1252-
12531340
let lShape = [1, shape, vocabSize]
12541341
let lStrides = try resolvedStrides(descriptor: logitsBaseDesc, shape: lShape)
1255-
var logitsOutput = unsafe InferenceFunction.AsyncMutableValue(
1256-
unsafeBuffer: logits.metalBuffer, byteOffset: 0,
1257-
scalarType: .float16, shape: lShape, strides: lStrides)
1258-
var asyncOutputs = InferenceFunction.AsyncMutableViews()
1259-
asyncOutputs.insert(&logitsOutput, for: logitsOutputName)
12601342

1261-
let _ = try function.encode(
1262-
inputs: asyncInputs,
1263-
states: consume asyncStates,
1264-
outputViews: consume asyncOutputs,
1265-
to: computeStream
1266-
)
1343+
switch additionalStates?.stateCount ?? 0 {
1344+
case 1:
1345+
let extra0 = additionalStates![stateIndex: 0]
1346+
var extraState0 = unsafe InferenceFunction.AsyncMutableValue(
1347+
unsafeBuffer: extra0.buffer, byteOffset: 0,
1348+
scalarType: extra0.scalarType, shape: extra0.shape, strides: extra0.strides)
1349+
var asyncStates = InferenceFunction.AsyncMutableViews()
1350+
asyncStates.insert(&keyState, for: keyCacheName)
1351+
asyncStates.insert(&valState, for: valueCacheName)
1352+
asyncStates.insert(&extraState0, for: extra0.name)
1353+
var logitsOutput = unsafe InferenceFunction.AsyncMutableValue(
1354+
unsafeBuffer: logits.metalBuffer, byteOffset: 0,
1355+
scalarType: .float16, shape: lShape, strides: lStrides)
1356+
var asyncOutputs = InferenceFunction.AsyncMutableViews()
1357+
asyncOutputs.insert(&logitsOutput, for: logitsOutputName)
1358+
let _ = try function.encode(
1359+
inputs: asyncInputs,
1360+
states: consume asyncStates,
1361+
outputViews: consume asyncOutputs,
1362+
to: computeStream
1363+
)
1364+
case 2:
1365+
let extra0 = additionalStates![stateIndex: 0]
1366+
let extra1 = additionalStates![stateIndex: 1]
1367+
var extraState0 = unsafe InferenceFunction.AsyncMutableValue(
1368+
unsafeBuffer: extra0.buffer, byteOffset: 0,
1369+
scalarType: extra0.scalarType, shape: extra0.shape, strides: extra0.strides)
1370+
var extraState1 = unsafe InferenceFunction.AsyncMutableValue(
1371+
unsafeBuffer: extra1.buffer, byteOffset: 0,
1372+
scalarType: extra1.scalarType, shape: extra1.shape, strides: extra1.strides)
1373+
var asyncStates = InferenceFunction.AsyncMutableViews()
1374+
asyncStates.insert(&keyState, for: keyCacheName)
1375+
asyncStates.insert(&valState, for: valueCacheName)
1376+
asyncStates.insert(&extraState0, for: extra0.name)
1377+
asyncStates.insert(&extraState1, for: extra1.name)
1378+
var logitsOutput = unsafe InferenceFunction.AsyncMutableValue(
1379+
unsafeBuffer: logits.metalBuffer, byteOffset: 0,
1380+
scalarType: .float16, shape: lShape, strides: lStrides)
1381+
var asyncOutputs = InferenceFunction.AsyncMutableViews()
1382+
asyncOutputs.insert(&logitsOutput, for: logitsOutputName)
1383+
let _ = try function.encode(
1384+
inputs: asyncInputs,
1385+
states: consume asyncStates,
1386+
outputViews: consume asyncOutputs,
1387+
to: computeStream
1388+
)
1389+
default:
1390+
var asyncStates = InferenceFunction.AsyncMutableViews()
1391+
asyncStates.insert(&keyState, for: keyCacheName)
1392+
asyncStates.insert(&valState, for: valueCacheName)
1393+
var logitsOutput = unsafe InferenceFunction.AsyncMutableValue(
1394+
unsafeBuffer: logits.metalBuffer, byteOffset: 0,
1395+
scalarType: .float16, shape: lShape, strides: lStrides)
1396+
var asyncOutputs = InferenceFunction.AsyncMutableViews()
1397+
asyncOutputs.insert(&logitsOutput, for: logitsOutputName)
1398+
let _ = try function.encode(
1399+
inputs: asyncInputs,
1400+
states: consume asyncStates,
1401+
outputViews: consume asyncOutputs,
1402+
to: computeStream
1403+
)
1404+
}
12671405

12681406
// Warm up argmax kernel using pipeline-matched decode buffers
12691407
let warmupLogitsBuffer = decodeLogitsBuffers[step % pipelineDepth]

0 commit comments

Comments
 (0)