Spaces:
Running
on
Zero
Running
on
Zero
File size: 1,833 Bytes
28c256d |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 |
#ifndef _MPS_UTILS_H_
#define _MPS_UTILS_H_
#include <torch/extension.h>
#ifdef __OBJC__
#include <Foundation/Foundation.h>
#include <Metal/Metal.h>
#include <MetalPerformanceShaders/MetalPerformanceShaders.h>
typedef id<MTLBuffer> MTLBuffer_t;
typedef id<MTLComputeCommandEncoder> MTLComputeCommandEncoder_t;
#else
typedef void* MTLBuffer;
typedef void* MTLBuffer_t;
typedef void* MTLComputeCommandEncoder;
typedef void* MTLComputeCommandEncoder_t;
#endif
// utils
static inline MTLBuffer_t getMTLBufferStorage(const at::Tensor& tensor) {
return __builtin_bit_cast(MTLBuffer_t, tensor.storage().data());
}
template <typename T,
std::enable_if_t<!std::is_same<std::decay_t<T>, at::Tensor>::value, bool> = true>
void setMTLArg(MTLComputeCommandEncoder_t encoder, int index, T&& t);
template <typename T,
std::enable_if_t<std::is_same<std::decay_t<T>, at::Tensor>::value, bool> = true>
void setMTLArg(MTLComputeCommandEncoder_t encoder, int index, T&& t) {
[encoder setBuffer:getMTLBufferStorage(t) offset:0 atIndex:index];
}
template <typename T, std::enable_if_t<!std::is_same<std::decay_t<T>, at::Tensor>::value, bool>>
void setMTLArg(MTLComputeCommandEncoder_t encoder, int index, T&& t) {
[encoder setBytes:&t length:sizeof(t) atIndex:index];
}
inline void setMTLArgsImpl(MTLComputeCommandEncoder_t, int) {}
template <typename T, typename... Args>
void setMTLArgsImpl(MTLComputeCommandEncoder_t encoder, int index, T&& t, Args&&... args) {
setMTLArg(encoder, index, std::forward<T>(t));
setMTLArgsImpl(encoder, index + 1, std::forward<Args>(args)...);
}
template <typename... Args>
void setMTLArgs(MTLComputeCommandEncoder_t encoder, MTLComputePipelineState_t pso, Args&&... args) {
[encoder setComputePipelineState:pso];
setMTLArgsImpl(encoder, 0, std::forward<Args>(args)...);
}
#endif
|