Use the new top-k implementation in Accelerate #21
Conversation
- added accuracy to testing method - added tests
There was a problem hiding this comment.
This works great!
Testing in a development branch, I got the following results:
- Greedy generation baseline: 58.53 tok/s
- Top-10 sampling before this PR: 36.40 tok/s
- Top-10 sampling with this PR: 57.71 tok/s
It's 59% improvement over the previous implementation, and a degradation of just 1.4% with respect to greedy generation. Fantastic!
I only have a comment / question about whether to control errors in the call to the accelerate function. I'm fine merging like this for now.
Sources/TensorUtils/Math.swift
Outdated
| try? Accelerate.BNNS.applyTopK( | ||
| k: k, | ||
| input: arrDescriptor, | ||
| bestValues: bestValues, | ||
| bestIndices: bestIndices, | ||
| axis: 0, | ||
| batchSize: 1, | ||
| filterParameters: nil | ||
| ) | ||
| let distances = bestValues.data!.withMemoryRebound(to: Float.self, capacity: k) { ptr in |
There was a problem hiding this comment.
Is there anything sensible to do if the operation fails rather than crashing? I'm fine merging as is, just wandering if you have any ideas that don't complicate things too much.
There was a problem hiding this comment.
I'm ok with merging it as is for now as well. Documentation is rather silent about the error being thrown, so I suspect that the swift function definition might be autogenerated from some lower level C code. Making this function throwable might just overcomplicate it too much with not much of a gain.
There was a problem hiding this comment.
@pcuenca do you think it'd be better to use try! here?
There was a problem hiding this comment.
Nah, I guess it will fail in the following line anyway.
There was a problem hiding this comment.
Actually, yes, let's do that just in case.
There was a problem hiding this comment.
done, forgot to call deallocate on BNNSNDArrayDescriptor so added it as well
| func XCTAssertEqual<T: FloatingPoint>( | ||
| _ expression1: @autoclosure () throws -> [T], | ||
| _ expression2: @autoclosure () throws -> [T], | ||
| accuracy: T, | ||
| _ message: @autoclosure () -> String = "", | ||
| file: StaticString = #filePath, | ||
| line: UInt = #line | ||
| ) { | ||
| do { | ||
| let lhsEvaluated = try expression1() | ||
| let rhsEvaluated = try expression2() | ||
| XCTAssertEqual(lhsEvaluated.count, rhsEvaluated.count, file: file, line: line) | ||
| for (lhs, rhs) in zip(lhsEvaluated, rhsEvaluated) { | ||
| XCTAssertEqual(lhs, rhs, accuracy: accuracy, file: file, line: line) | ||
| } | ||
| } catch { | ||
| XCTFail("Unexpected error: \(error)", file: file, line: line) | ||
| } | ||
| } |
- try! instead of try? when applying topk function
|
Amazing work, thanks again! |
In this PR:
Accelerate.BNNS.applyTopKmethod