Skip to content

Commit 4805f4c

Browse files
stikvesstikves
committed
Fix Swift format
1 parent 6921969 commit 4805f4c

1 file changed

Lines changed: 45 additions & 33 deletions

File tree

‎swift/Sources/CoreAILanguageModels/Samplers/MPSGraphSamplers.swift‎

Lines changed: 45 additions & 33 deletions
Original file line numberDiff line numberDiff line change
@@ -316,9 +316,11 @@ final class MPSGraphArgmaxSampler: @unchecked Sendable {
316316
get throws {
317317
if let buf = constrainedBitmaskBuffer { return buf }
318318
let byteCount = max(bitmaskSize * MemoryLayout<Int32>.size, 64)
319-
guard let buf = device.makeBuffer(
320-
length: byteCount, options: .storageModeShared
321-
) else {
319+
guard
320+
let buf = device.makeBuffer(
321+
length: byteCount, options: .storageModeShared
322+
)
323+
else {
322324
throw MPSGraphSamplerError.bufferAllocationFailed
323325
}
324326
constrainedBitmaskBuffer = buf
@@ -403,9 +405,10 @@ final class MPSGraphArgmaxSampler: @unchecked Sendable {
403405
with: queue, inputs: [inputData, bitmaskTensorData],
404406
results: [outputData], executionDescriptor: execDesc)
405407
} else {
406-
encode(to: queue, logitsBuffer: logitsBuffer, logitsOffset: logitsOffset,
407-
outputBuffer: outputBuffer, outputOffset: outputOffset,
408-
completion: completion)
408+
encode(
409+
to: queue, logitsBuffer: logitsBuffer, logitsOffset: logitsOffset,
410+
outputBuffer: outputBuffer, outputOffset: outputOffset,
411+
completion: completion)
409412
}
410413
}
411414

@@ -420,29 +423,32 @@ final class MPSGraphArgmaxSampler: @unchecked Sendable {
420423
completion: @escaping (Int32) -> Void
421424
) {
422425
if queryLength == 1 {
423-
encode(to: queue, logitsBuffer: logitsBuffer, logitsOffset: 0,
424-
outputBuffer: outputBuffer, outputOffset: outputOffset,
425-
applyBitmask: applyBitmask, completion: completion)
426+
encode(
427+
to: queue, logitsBuffer: logitsBuffer, logitsOffset: 0,
428+
outputBuffer: outputBuffer, outputOffset: outputOffset,
429+
applyBitmask: applyBitmask, completion: completion)
426430
return
427431
}
428432
// For prefill with bitmask, slice last token then apply constrained path
429433
let logitsOffset = (queryLength - 1) * vocabSize * MemoryLayout<UInt16>.size
430434
let sliceSize = vocabSize * MemoryLayout<UInt16>.size
431435
guard let tempBuffer = device.makeBuffer(length: sliceSize, options: .storageModeShared),
432-
let blitCmdBuffer = queue.makeCommandBuffer(),
433-
let blitEncoder = blitCmdBuffer.makeBlitCommandEncoder()
436+
let blitCmdBuffer = queue.makeCommandBuffer(),
437+
let blitEncoder = blitCmdBuffer.makeBlitCommandEncoder()
434438
else {
435439
completion(0)
436440
return
437441
}
438-
blitEncoder.copy(from: logitsBuffer, sourceOffset: logitsOffset,
439-
to: tempBuffer, destinationOffset: 0, size: sliceSize)
442+
blitEncoder.copy(
443+
from: logitsBuffer, sourceOffset: logitsOffset,
444+
to: tempBuffer, destinationOffset: 0, size: sliceSize)
440445
blitEncoder.endEncoding()
441446
blitCmdBuffer.commit()
442447

443-
encode(to: queue, logitsBuffer: tempBuffer, logitsOffset: 0,
444-
outputBuffer: outputBuffer, outputOffset: outputOffset,
445-
applyBitmask: applyBitmask, completion: completion)
448+
encode(
449+
to: queue, logitsBuffer: tempBuffer, logitsOffset: 0,
450+
outputBuffer: outputBuffer, outputOffset: outputOffset,
451+
applyBitmask: applyBitmask, completion: completion)
446452
}
447453

448454
/// Encode argmax sampling.
@@ -895,10 +901,12 @@ final class MPSGraphCompositeSampler: @unchecked Sendable {
895901
var bitmaskBuffer: MTLBuffer {
896902
get throws {
897903
if let buf = constrainedBitmaskBuffer { return buf }
898-
guard let buf = device.makeBuffer(
899-
length: bitmaskSize * MemoryLayout<Int32>.size,
900-
options: .storageModeShared
901-
) else {
904+
guard
905+
let buf = device.makeBuffer(
906+
length: bitmaskSize * MemoryLayout<Int32>.size,
907+
options: .storageModeShared
908+
)
909+
else {
902910
throw MPSGraphSamplerError.bufferAllocationFailed
903911
}
904912
constrainedBitmaskBuffer = buf
@@ -1021,9 +1029,10 @@ final class MPSGraphCompositeSampler: @unchecked Sendable {
10211029
inputs: [logitsData, temperatureData, randomData, topPData, minPData, bitmaskTensorData],
10221030
results: [outputData], executionDescriptor: execDesc)
10231031
} else {
1024-
encode(to: queue, logitsBuffer: logitsBuffer, logitsOffset: logitsOffset,
1025-
outputBuffer: outputBuffer, outputOffset: outputOffset,
1026-
completion: completion)
1032+
encode(
1033+
to: queue, logitsBuffer: logitsBuffer, logitsOffset: logitsOffset,
1034+
outputBuffer: outputBuffer, outputOffset: outputOffset,
1035+
completion: completion)
10271036
}
10281037
}
10291038

@@ -1038,28 +1047,31 @@ final class MPSGraphCompositeSampler: @unchecked Sendable {
10381047
completion: @escaping (Int32) -> Void
10391048
) {
10401049
if queryLength == 1 {
1041-
encode(to: queue, logitsBuffer: logitsBuffer, logitsOffset: 0,
1042-
outputBuffer: outputBuffer, outputOffset: outputOffset,
1043-
applyBitmask: applyBitmask, completion: completion)
1050+
encode(
1051+
to: queue, logitsBuffer: logitsBuffer, logitsOffset: 0,
1052+
outputBuffer: outputBuffer, outputOffset: outputOffset,
1053+
applyBitmask: applyBitmask, completion: completion)
10441054
return
10451055
}
10461056
let logitsOffset = (queryLength - 1) * vocabSize * MemoryLayout<UInt16>.size
10471057
let sliceSize = vocabSize * MemoryLayout<UInt16>.size
10481058
guard let tempBuffer = device.makeBuffer(length: sliceSize, options: .storageModeShared),
1049-
let blitCmdBuffer = queue.makeCommandBuffer(),
1050-
let blitEncoder = blitCmdBuffer.makeBlitCommandEncoder()
1059+
let blitCmdBuffer = queue.makeCommandBuffer(),
1060+
let blitEncoder = blitCmdBuffer.makeBlitCommandEncoder()
10511061
else {
10521062
completion(0)
10531063
return
10541064
}
1055-
blitEncoder.copy(from: logitsBuffer, sourceOffset: logitsOffset,
1056-
to: tempBuffer, destinationOffset: 0, size: sliceSize)
1065+
blitEncoder.copy(
1066+
from: logitsBuffer, sourceOffset: logitsOffset,
1067+
to: tempBuffer, destinationOffset: 0, size: sliceSize)
10571068
blitEncoder.endEncoding()
10581069
blitCmdBuffer.commit()
10591070

1060-
encode(to: queue, logitsBuffer: tempBuffer, logitsOffset: 0,
1061-
outputBuffer: outputBuffer, outputOffset: outputOffset,
1062-
applyBitmask: applyBitmask, completion: completion)
1071+
encode(
1072+
to: queue, logitsBuffer: tempBuffer, logitsOffset: 0,
1073+
outputBuffer: outputBuffer, outputOffset: outputOffset,
1074+
applyBitmask: applyBitmask, completion: completion)
10631075
}
10641076

10651077
/// Encode composite sampling asynchronously (protocol conformance).

0 commit comments

Comments
 (0)