<commit_before>//
//  DeviceOperations.swift
//  CUDA
//
//  Created by <NAME> on 11/5/16.
//
//

import CuBLAS
import CUDADriver

infix operator • : MultiplicationPrecedence

extension DeviceArray {
    var kernelManager: KernelManager {
        return KernelManager.shared(on: device)
    }
}

public extension DeviceArray where Element : BLASDataProtocol & FloatingPoint {

    public func dotProduct(with other: DeviceArray) -> Element {
        return withUnsafeDevicePointer { ptr in
            other.withUnsafeDevicePointer { otherPtr in
                BLAS.global(on: self.device).dot(
                    x: ptr, stride: 1, y: otherPtr, stride: 1,
                    count: Int32(Swift.min(self.count, other.count))
                )
            }
        }
    }

    @inline(__always)
    public static func •(lhs: DeviceArray<Element>, rhs: DeviceArray<Element>) -> Element {
        return lhs.dotProduct(with: rhs)
    }

}

public extension DeviceArray where Element : KernelDataProtocol {

    public mutating func add(_ other: DeviceArray<Element>, multipliedBy alpha: Element = 1) {
        let axpy = kernelManager.kernel(.axpy, forType: Element.self)
        device.sync {
            let blockCount = (count+127)/128
            try! axpy<<<(blockCount, 128)>>>[
                .value(alpha),
                .constPointer(to: other),
                .pointer(to: &self),
                .longLong(Int64(count))
            ]
        }
    }

    public func adding(_ other: DeviceArray<Element>, multipliedBy alpha: Element = 1) -> DeviceArray<Element> {
        var copy = self
        copy.add(other, multipliedBy: alpha)
        return copy
    }

    public static func +=(lhs: inout DeviceArray<Element>, rhs: DeviceArray<Element>) {
        return lhs.add(rhs)
    }

    public mutating func scale(by alpha: Element) {
        let scale = kernelManager.kernel(.scale, forType: Element.self)
        device.sync {
            let blockCount = (count+127)/128
            try! scale<<<(blockCount, 128)>>>[
                .pointer(to: &self), .value(alpha), .longLong(Int64(count))
            ]
        }
    }

    @inline(__always)
    public static func *=(lhs: inout DeviceArray<Element>, rhs: Element) {
        lhs.scale(by: rhs)
    }

    public func scaled(by alpha: Element) -> DeviceArray {
        var copy = self
        copy.scale(by: alpha)
        return copy
    }

    public func reduced() -> Element {
        var result = DeviceValue<Element>()
        let sum = kernelManager.kernel(.sum, forType: Element.self)<commit_msg>Let blockSize be min(128, count)<commit_after>//
//  DeviceOperations.swift
//  CUDA
//
//  Created by <NAME> on 11/5/16.
//
//

import CuBLAS
import CUDADriver

infix operator • : MultiplicationPrecedence

extension DeviceArray {
    var kernelManager: KernelManager {
        return KernelManager.shared(on: device)
    }
}

public extension DeviceArray where Element : BLASDataProtocol & FloatingPoint {

    public func dotProduct(with other: DeviceArray) -> Element {
        return withUnsafeDevicePointer { ptr in
            other.withUnsafeDevicePointer { otherPtr in
                BLAS.global(on: self.device).dot(
                    x: ptr, stride: 1, y: otherPtr, stride: 1,
                    count: Int32(Swift.min(self.count, other.count))
                )
            }
        }
    }

    @inline(__always)
    public static func •(lhs: DeviceArray<Element>, rhs: DeviceArray<Element>) -> Element {
        return lhs.dotProduct(with: rhs)
    }

}

public extension DeviceArray where Element : KernelDataProtocol {

    public mutating func add(_ other: DeviceArray<Element>, multipliedBy alpha: Element = 1) {
        let axpy = kernelManager.kernel(.axpy, forType: Element.self)
        device.sync {
            let blockSize = Swift.min(128, count)
            let blockCount = (count+blockSize-1)/blockSize
            try! axpy<<<(blockCount, blockSize)>>>[
                .value(alpha),
                .constPointer(to: other),
                .pointer(to: &self),
                .longLong(Int64(count))
            ]
        }
    }

    public func adding(_ other: DeviceArray<Element>, multipliedBy alpha: Element = 1) -> DeviceArray<Element> {
        var copy = self
        copy.add(other, multipliedBy: alpha)
        return copy
    }

    public static func +=(lhs: inout DeviceArray<Element>, rhs: DeviceArray<Element>) {
        return lhs.add(rhs)
    }

    public mutating func scale(by alpha: Element) {
        let scale = kernelManager.kernel(.scale, forType: Element.self)
        device.sync {
            let blockSize = Swift.min(128, count)
            let blockCount = (count+blockSize-1)/blockSize
            try! scale<<<(blockCount, blockSize)>>>[
                .pointer(to: &self), .value(alpha), .longLong(Int64(count))
            ]
        }
    }

    @inline(__always)
    public static func *=(lhs: inout DeviceArray<Element>, rhs: Element) {
        lhs.scale(by: rhs)
    }

    public func scaled(by alpha: Element) -> DeviceArray {
        var copy = self
        copy.scale(by: alpha)
        return copy
    }

    public func reduced() -> Element {
        var result = DeviceValue<Element>()
        let sum = kernelManager.kernel(.sum, forType: Element.self)
{'fraction_non_alphanumeric': 0.11119473189087488, 'fraction_numerical': 0.01110065851364064, 'mean_word_length': 2.1813285457809695, 'pattern_counts': {'":': 0, '<': 35, '<?xml version=': 0, '>': 45, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '6326052', 'n_tokens_mistral': 1594, 'n_tokens_neox': 1497, 'n_words': 389}