Source ml/tensor_vdsp.sc

// SafeC Standard Library — vDSP-accelerated elementwise Tensor ops
// implementation (see tensor_vdsp.h).
// safeguard: link macos -framework Accelerate
#pragma once
#include <std/ml/tensor_vdsp.h>
#include <std/ml/tensor.h>

namespace std {

// Same Accelerate framework as tensor_blas.sc's cblas_sgemm, a different
// corner of it: vectorized elementwise float32 ops instead of matrix
// multiply. Verified directly (not assumed) before using any of these:
// vDSP_vsub computes C = B - A (second argument minus first — the
// reverse of what the argument order suggests), the other two are the
// ordinary commutative operations their names imply.
#ifdef __APPLE__
extern void vDSP_vadd(const float* A, long IA, const float* B, long IB,
                       float* C, long IC, unsigned long N);
extern void vDSP_vsub(const float* A, long IA, const float* B, long IB,
                       float* C, long IC, unsigned long N);
extern void vDSP_vmul(const float* A, long IA, const float* B, long IB,
                       float* C, long IC, unsigned long N);
extern void vDSP_sve(const float* A, long IA, float* result, unsigned long N);

&FloatTensor tensor_add_vdsp(const &FloatTensor a, const &FloatTensor b) {
    &FloatTensor out = __tensor_alloc_uninit_like(a);
    unsafe {
        vDSP_vadd((const float*)a->data, 1L, (const float*)b->data, 1L,
                  (float*)out->data, 1L, out->size);
    }
    __tensor_link2(out, a, b, __add_backward);
    return out;
}

&FloatTensor tensor_sub_vdsp(const &FloatTensor a, const &FloatTensor b) {
    &FloatTensor out = __tensor_alloc_uninit_like(a);
    unsafe {
        // out = a - b: vDSP_vsub(X, ..., Y, ..., Z, ...) computes Z = Y - X,
        // so X=b, Y=a gives Z=a-b.
        vDSP_vsub((const float*)b->data, 1L, (const float*)a->data, 1L,
                  (float*)out->data, 1L, out->size);
    }
    __tensor_link2(out, a, b, __sub_backward);
    return out;
}

&FloatTensor tensor_mul_vdsp(const &FloatTensor a, const &FloatTensor b) {
    &FloatTensor out = __tensor_alloc_uninit_like(a);
    unsafe {
        vDSP_vmul((const float*)a->data, 1L, (const float*)b->data, 1L,
                  (float*)out->data, 1L, out->size);
    }
    __tensor_link2(out, a, b, __mul_backward);
    return out;
}

&FloatTensor FloatTensor::sum_vdsp() const {
    unsigned long one[1];
    unsafe { one[0] = 1UL; }
    &FloatTensor out = __tensor_alloc(one, 1UL, 0);
    unsafe {
        float acc = (float)0.0;
        vDSP_sve((const float*)self->data, 1L, &acc, self->size);
        out->data[0] = acc;
    }
    __tensor_link1(out, self, __sum_backward);
    return out;
}

#else
// Accelerate is Apple-only. Elsewhere the same functions are std's own
// vectorized, threaded tensor ops (tensor.h), so tensor_auto.h and direct
// callers work -- and link -- on every hosted platform.
&FloatTensor tensor_add_vdsp(const &FloatTensor a, const &FloatTensor b) { return tensor_add(a, b); }
&FloatTensor tensor_sub_vdsp(const &FloatTensor a, const &FloatTensor b) { return tensor_sub(a, b); }
&FloatTensor tensor_mul_vdsp(const &FloatTensor a, const &FloatTensor b) { return tensor_mul(a, b); }
&FloatTensor FloatTensor::sum_vdsp() const { return self.sum(); }
#endif

} // namespace std