gaspersic/triton
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

[examples/python/pytorch] added batchnorm cpp extension

Browse Source
This commit is contained in:
Philippe Tillet
2019-07-09 20:59:04 -07:00
parent b7986baffa
commit 63b249c1d6
6 changed files with 214 additions and 206 deletions
Download Patch File Download Diff File Expand all files Collapse all files

164
examples/python/pytorch/shift.cpp
View File

@@ -12,103 +12,111 @@
#define CHECK_CONTIGUOUS(x) AT_CHECK(x.is_contiguous(), #x " must be contiguous")
#define CHECK_INPUT(x) CHECK_CUDA(x); CHECK_CONTIGUOUS(x)
typedef std::tuple<int32_t, int32_t, int32_t, int32_t, int32_t,
int32_t, int32_t, int32_t, int32_t,
int32_t*, int32_t*,
triton::dnn::shift::type, bool> shift_key_t;
static std::map<CUstream, std::unique_ptr<triton::driver::stream>> m_shift_stream;
static std::map<shift_key_t, std::unique_ptr<triton::jit>> m_shift_jit;
static std::map<shift_key_t, std::unique_ptr<triton::dnn::shift>> m_shift_config;
torch::Tensor shift_common(
int32_t B, int32_t C, int32_t D, int32_t H, int32_t W,
int32_t T, int32_t R, int32_t S, int32_t F,
std::vector<int32_t> shift_h, std::vector<int32_t> shift_w,
int32_t stride_h, int32_t stride_w,
int32_t* shift_h, int32_t* shift_w,
triton::dnn::shift::type ty,
torch::Tensor torcha, torch::Tensor torchb, torch::Tensor torchbias,
bool autotune = false
) {
// Wrap CUDA handles
c10::DeviceIndex device = torcha.storage().device().index();
// Get stream
CUstream custream = (CUstream)at::cuda::getCurrentCUDAStream(device).stream();
triton::driver::stream* stream;
if(m_shift_stream.find(custream) == m_shift_stream.end())
stream = m_shift_stream.emplace(custream, new triton::driver::cu_stream(custream, false)).first->second.get();
else
stream = m_shift_stream.at(custream).get();
// Get context
triton::driver::context* ctx = stream->context();
triton::driver::cu_stream stream(custream, false);
triton::driver::context* ctx = stream.context();
// Get configuration
bool has_bias = torchbias.storage().size() > 0;
shift_key_t key = {B, C, D, H, W, T, R, S, F, shift_h.data(), shift_w.data(), ty, has_bias};
triton::dnn::shift* configuration;
if(m_shift_config.find(key) == m_shift_config.end())
configuration = m_shift_config.emplace(key, new triton::dnn::shift(
B, C, D, H, W, T, R, S, F,
shift_h, shift_w, "fp32", "fp32",
ty, has_bias)).first->second.get();
else
configuration = m_shift_config.at(key).get();
triton::dnn::shift shift(B, C, D, H, W, T, R, S, F,
stride_h, stride_w,
shift_h, shift_w, "fp32", "fp32",
ty, has_bias);
// Bind memory
triton::driver::cu_buffer a(ctx, (CUdeviceptr)torcha.storage().data(), false);
triton::driver::cu_buffer b(ctx, (CUdeviceptr)torchb.storage().data(), false);
triton::driver::cu_buffer cubias(ctx, (CUdeviceptr)torchbias.storage().data(), false);
triton::driver::buffer* bias = has_bias ? &cubias : nullptr;
// Allocate output
std::vector<int32_t> c_shapes = configuration->c_shapes();
std::vector<int32_t> c_shapes = shift.c_shapes();
torch::Tensor torchc = torch::empty({c_shapes[0], c_shapes[1], c_shapes[2], c_shapes[3]}).cuda();
triton::driver::cu_buffer c(ctx, (CUdeviceptr)torchc.storage().data(), false);
// Get JIT
triton::jit* jit;
if(m_shift_jit.find(key) == m_shift_jit.end()){
jit = m_shift_jit.emplace(key, new triton::jit(ctx)).first->second.get();
std::ostringstream oss;
configuration->triton_c_src(oss);
std::string src = oss.str();
// benchmark a given shiftolution kernel
auto benchmark = [&](triton::driver::kernel* kernel,
triton::jit::launch_information info) {
configuration->init_impl(stream, (triton::driver::cu_module*)kernel->module());
unsigned TM = info.global_range_size[0];
unsigned TN = info.global_range_size[1];
unsigned nthreads = info.num_threads;
configuration->enqueue_impl(stream, kernel, &a, &b, &c, TM, TN, nthreads);
stream->synchronize();
double ts = triton::tools::bench([&](){ configuration->enqueue_impl(stream, kernel, &a, &b, &c, TM, TN, nthreads); },
[&](){ stream->synchronize(); }, stream->context()->device());
return configuration->num_flops() / ts * 1e-3;
};
// auto-tune and save result
if(autotune) {
triton::jit::tune_res_t best = jit->autotune("shift", src.c_str(), benchmark);
jit->add_module("shift", src.c_str(), best.params);
}
else {
jit->add_module("shift", src.c_str(), jit->get_valid("shift", src.c_str()));
}
triton::driver::kernel* kernel = jit->get_function("shift");
configuration->init_impl(stream, (triton::driver::cu_module*)kernel->module());
}
else
jit = m_shift_jit.at(key).get();
// Run
triton::driver::kernel* kernel = jit->get_function("shift");
triton::jit::launch_information info = jit->get_launch_info("shift");
// launch info
unsigned TM = info.global_range_size[0];
unsigned TN = info.global_range_size[1];
unsigned nthreads = info.num_threads;
// enqueue
configuration->enqueue_impl(stream, kernel, &a, &b, &c, TM, TN, nthreads);
// Enqueue
shift.enqueue(&stream, {&a, &b, &c});
return torchc;
}
torch::Tensor shift_y(
const torch::Tensor x,
const torch::Tensor w,
const torch::Tensor bias,
int32_t R, int32_t S,
int32_t stride_h, int32_t stride_w,
int32_t* shift_h, int32_t* shift_w) {
// shapes for a
int64_t Ca = x.size(0);
int64_t H = x.size(1);
int64_t W = x.size(2);
int64_t B = x.size(3);
// shapes for b
int64_t Cb = w.size(0);
int64_t F = w.size(1);
AT_CHECK(Ca == Cb, "operands must have the same number of channels");
int64_t C = Ca;
// run
shift_common(B, C, 1, H, W, 1, R, S, F, stride_h, stride_w, shift_h, shift_w, triton::dnn::shift::FPROP, x, w, bias);
}
torch::Tensor shift_dx(
const torch::Tensor dy,
const torch::Tensor w,
const torch::Tensor bias,
int32_t R, int32_t S,
int32_t stride_h, int32_t stride_w,
int32_t* shift_h, int32_t* shift_w) {
// shapes for a
int64_t Ca = dy.size(0);
int64_t H = dy.size(1);
int64_t W = dy.size(2);
int64_t B = dy.size(3);
H *= stride_h;
W *= stride_w;
// shapes for b
int64_t Cb = w.size(0);
int64_t F = w.size(1);
std::swap(Cb, F);
// checks
AT_CHECK(Ca == Cb, "operands must have the same number of channels");
int64_t C = Ca;
std::swap(C, F);
// run
shift_common(B, C, 1, H, W, 1, R, S, F, stride_h, stride_w, shift_h, shift_w, triton::dnn::shift::BPROP, dy, w, bias);
}
torch::Tensor shift_dw(
const torch::Tensor dy,
const torch::Tensor x,
const torch::Tensor bias,
int32_t R, int32_t S,
int32_t stride_h, int32_t stride_w,
int32_t* shift_h, int32_t* shift_w) {
// shapes for a
int64_t F = dy.size(0);
int64_t Ha = dy.size(1);
int64_t Wa = dy.size(2);
int64_t Ba = dy.size(3);
// shapes for b
int64_t C = x.size(0);
int64_t Hb = x.size(1);
int64_t Wb = x.size(2);
int64_t Bb = x.size(3);
// check
AT_CHECK(Ha*stride_h == Hb, "operands must have the same image height");
AT_CHECK(Wa*stride_w == Wb, "operands must have the same image width");
AT_CHECK(Ba == Bb, "operands must have the same batch size");
int64_t H = Hb;
int64_t W = Wb;
int64_t B = Bb;
// run
shift_common(B, C, 1, H, W, 1, R, S, F, stride_h, stride_w, shift_h, shift_w, triton::dnn::shift::WGRAD, dy, x, bias);
}