293 lines
10 KiB
C++
293 lines
10 KiB
C++
#include <torch/torch.h>
|
|
#include <torch/script.h>
|
|
#include "ATen/cuda/CUDAContext.h"
|
|
#include <vector>
|
|
#include "triton/jit.h"
|
|
#include "triton/driver/stream.h"
|
|
|
|
#define CHECK_CUDA(x) AT_CHECK(x.type().is_cuda(), #x " must be a CUDA tensor")
|
|
#define CHECK_CONTIGUOUS(x) AT_CHECK(x.is_contiguous(), #x " must be contiguous")
|
|
#define CHECK_INPUT(x) CHECK_CUDA(x); CHECK_CONTIGUOUS(x)
|
|
|
|
const char* src =
|
|
R"(
|
|
const tunable int32 TM = {16, 32, 64};
|
|
const tunable int32 TN = {16, 32, 64};
|
|
const tunable int32 TK = {8};
|
|
|
|
__constant__ int32* delta = alloc_const int32[18];
|
|
__constant__ int32* masks = alloc_const int32[1024];
|
|
|
|
void conv(read_only restrict fp32 *a,
|
|
read_only restrict fp32 *b,
|
|
fp32 *c,
|
|
int32 M, int32 N, int32 K,
|
|
int32 AN, int32 AH, int32 AW,
|
|
int32 CN, int32 CK, int32 CP, int32 CQ,
|
|
int32 AC, int32 AR, int32 AS,
|
|
int32 lda_n, int32 lda_c, int32 lda_h, int32 lda_w,
|
|
int32 ldc_n, int32 ldc_k, int32 ldc_p, int32 ldc_q,
|
|
int32 pad_h, int32 pad_w,
|
|
int32 bound){
|
|
int32 rxa[TM] = get_global_range[TM](0);
|
|
int32 rb0[TN] = get_global_range[TN](1);
|
|
int32 rka[TK] = 0 ... TK;
|
|
int32 rb1[TK] = 0 ... TK;
|
|
fp32 C[TM, TN] = 0;
|
|
int32 ranh[TM] = rxa / CQ;
|
|
int32 raw[TM] = rxa % CQ - pad_w;
|
|
int32 ran[TM] = ranh / CP;
|
|
int32 rah[TM] = ranh % CP - pad_h;
|
|
int32 ra0[TM] = ran*lda_n + rah*lda_h + raw*lda_w;
|
|
int32 racr[TK] = rka / AS;
|
|
int32 ras[TK] = rka % AS;
|
|
int32 rac[TK] = racr / AR;
|
|
int32 rar[TK] = racr % AR;
|
|
int32 ra1[TK] = rac*lda_c + rar*lda_h + ras*lda_w;
|
|
fp32* pa[TM, TK] = a + ra1[newaxis, :] + ra0[:, newaxis];
|
|
fp32* pb[TN, TK] = b + rb1[newaxis, :]*CK + rb0[:, newaxis];
|
|
__constant__ int32* pincd[TK] = delta + rka;
|
|
__constant__ int32* pd[TK] = delta + AR*AS + rka;
|
|
int32 d[TK] = *pd;
|
|
int32 incd[TK] = *pincd;
|
|
int32 maskh[TM] = pad_h + min(rah, 0) + max(rah + AR - AH, 0);
|
|
int32 maskw[TM] = pad_w + min(raw, 0) + max(raw + AS - AW, 0);
|
|
__constant__ int32* pm[TM] = masks + AR*AS + maskw*AR*AS + maskh*AR*AS*(2*pad_w + 1);
|
|
__constant__ int32* pincm[TM] = delta;
|
|
int32 incm[TM] = *pincm;
|
|
int32 checka0[TM] = *pm;
|
|
int32 checka1[TK] = 1 << rka;
|
|
int1 checka[TM, TK] = (checka0[:, newaxis] & checka1[newaxis, :]) > 0;
|
|
fp32 a[TM, TK] = checka ? *pa : 0;
|
|
fp32 b[TN, TK] = *pb;
|
|
for(int32 k = K; k > 0; k = k - TK){
|
|
C = dot(a, trans(b), C);
|
|
pb = pb + TK*CK;
|
|
pa = pa + d[newaxis, :];
|
|
b = *pb;
|
|
pd = pd + incd;
|
|
pincd = pincd + incd;
|
|
d = *pd;
|
|
incd = *pincd;
|
|
pm = pm + incm;
|
|
pincm = pincm + incm;
|
|
incm = *pincm;
|
|
checka0 = *pm;
|
|
checka = (checka0[:, newaxis] & checka1[newaxis, :]) > 0;
|
|
a = checka ? *pa : 0;
|
|
}
|
|
int32 rxc[TM] = get_global_range[TM](0);
|
|
int32 rc1[TN] = get_global_range[TN](1);
|
|
int32 rcn[TM] = rxc / (CP*CQ);
|
|
int32 rcpq[TM] = rxc % (CP*CQ);
|
|
int32 rc0[TM] = rcn * ldc_n + rcpq;
|
|
fp32* pc[TM, TN] = c + rc1[newaxis, :]*ldc_k + rc0[:, newaxis];
|
|
int1 checkc0[TM] = rxc < M;
|
|
int1 checkc1[TN] = rc1 < N;
|
|
int1 checkc[TM, TN] = checkc0[:, newaxis] && checkc1[newaxis, :];
|
|
@checkc *pc = C;
|
|
})";
|
|
|
|
void build_conv_lut(int TK,
|
|
int stride_d, int stride_h, int stride_w, int stride_c,
|
|
int pad_d, int pad_h, int pad_w,
|
|
int T, int R, int S,
|
|
std::vector<int>& res, std::vector<int>& masks) {
|
|
/* convolution parameters */
|
|
int F = T * R * S;
|
|
int Nlut = (TK + F - 1) / F * F;
|
|
int upsample_w = 1;
|
|
int upsample_h = 1;
|
|
int upsample_d = 1;
|
|
/* unpack index wrt filters */
|
|
auto unpack = [&](int32_t trs){
|
|
int32_t tr = trs / S;
|
|
int32_t s = trs - tr*S;
|
|
int32_t t = tr / R;
|
|
int32_t r = tr - t*R;
|
|
return std::make_tuple(t, r, s);
|
|
};
|
|
/* increments */
|
|
for(size_t i = 0; i < Nlut; ++i)
|
|
res[i] = (((i + TK) % Nlut) - i);
|
|
/* deltas */
|
|
size_t Ds0 = Nlut;
|
|
size_t Ds1 = upsample_w;
|
|
size_t Ds2 = upsample_h;
|
|
size_t Ds3 = upsample_d;
|
|
for(size_t pd = 0; pd < Ds3; ++pd)
|
|
for(size_t ph = 0; ph < Ds2; ++ph)
|
|
for(size_t pw = 0; pw < Ds1; ++pw){
|
|
int32_t* deltas_ptr = &res[Nlut + pw*Ds0 + ph*Ds0*Ds1 + pd*Ds0*Ds1*Ds2];
|
|
// cumulative increments
|
|
for(size_t i = 0; i < Ds0; ++i){
|
|
int32_t ctrs = i;
|
|
int32_t c = ctrs / F;
|
|
int32_t t, r, s;
|
|
std::tie(t, r, s) = unpack(ctrs % F);
|
|
// next indices
|
|
int32_t nextctrs = ctrs + TK;
|
|
int32_t nextc = nextctrs / F;
|
|
int32_t nextt, nextr, nexts;
|
|
std::tie(nextt, nextr, nexts) = unpack(nextctrs % F);
|
|
// diffs
|
|
int32_t cdiff = nextc - c;
|
|
int32_t tdiff = (nextt + pd)/upsample_d - (t + pd)/upsample_d;
|
|
int32_t rdiff = (nextr + ph)/upsample_h - (r + ph)/upsample_h;
|
|
int32_t sdiff = (nexts + pw)/upsample_w - (s + pw)/upsample_w;
|
|
// delta pointers
|
|
deltas_ptr[i] = cdiff*stride_c + sdiff*stride_w + rdiff*stride_h + tdiff*stride_d;
|
|
}
|
|
}
|
|
|
|
/* Masks */
|
|
size_t Ms0 = Nlut;
|
|
size_t Ms1 = 2*pad_w + 1;
|
|
size_t Ms2 = 2*pad_h + 1;
|
|
size_t Ms3 = 2*pad_d + 1;
|
|
|
|
for(size_t pd = 0; pd < Ms3; ++pd)
|
|
for(size_t ph = 0; ph < Ms2; ++ph)
|
|
for(size_t pw = 0; pw < Ms1; ++pw){
|
|
int32_t* masks_ptr = &masks[Nlut + pw*Ms0 + ph*Ms0*Ms1 + pd*Ms0*Ms1*Ms2];
|
|
for(size_t i = 0; i < Ms0; ++i){
|
|
int32_t t, r, s;
|
|
int32_t mask = 0x0;
|
|
for(size_t j = 0; j < TK; ++j){
|
|
std::tie(t, r, s) = unpack((i + j) % F);
|
|
bool in_bounds_d = (t + pd) >= pad_d && (t + pd) < (T + pad_d);
|
|
bool in_bounds_h = (r + ph) >= pad_h && (r + ph) < (R + pad_h);
|
|
bool in_bounds_w = (s + pw) >= pad_w && (s + pw) < (S + pad_w);
|
|
mask |= (in_bounds_d && in_bounds_h && in_bounds_w) << j;
|
|
}
|
|
masks_ptr[i] = mask;
|
|
}
|
|
}
|
|
for(size_t i = 0; i < Nlut; ++i)
|
|
masks[i] = 0x0;
|
|
}
|
|
|
|
torch::Tensor conv_forward(
|
|
const torch::Tensor data,
|
|
const torch::Tensor weight) {
|
|
// Check
|
|
CHECK_INPUT(data);
|
|
CHECK_INPUT(weight);
|
|
// Unpack data shapes
|
|
const int32_t B = data.size(0);
|
|
const int32_t Ci = data.size(1);
|
|
const int32_t H = data.size(2);
|
|
const int32_t W = data.size(3);
|
|
// Unpack weight shapes
|
|
const int32_t Cf = weight.size(0);
|
|
const int32_t T = 1;
|
|
const int32_t R = weight.size(1);
|
|
const int32_t S = weight.size(2);
|
|
const int32_t NF = weight.size(3);
|
|
// Conv parameters
|
|
int32_t upsample_d = 1, upsample_h = 1, upsample_w = 1;
|
|
int32_t pad_d = 0, pad_h = 0, pad_w = 0;
|
|
int32_t stride_h = 1, stride_w = 1;
|
|
// Output shapes
|
|
int32_t P = (H*upsample_h - R + 1 + 2*pad_h + stride_h - 1)/stride_h;
|
|
int32_t Q = (W*upsample_w - S + 1 + 2*pad_w + stride_w - 1)/stride_w;
|
|
// Allocate output
|
|
AT_CHECK(Ci == Cf, "Number of channels in data and weights must match");
|
|
torch::Tensor output = torch::empty({B, NF, P, Q}, torch::kFloat).cuda();
|
|
// Wrap CUDA handles
|
|
c10::DeviceIndex device = output.storage().device().index();
|
|
triton::driver::cu_stream sstream((CUstream)at::cuda::getCurrentCUDAStream(device).stream(), false);
|
|
triton::driver::stream* stream = &sstream;
|
|
triton::driver::context* ctx = stream->context();
|
|
triton::driver::cu_buffer d(ctx, (CUdeviceptr)data.storage().data(), false);
|
|
triton::driver::cu_buffer w(ctx, (CUdeviceptr)weight.storage().data(), false);
|
|
triton::driver::cu_buffer a(ctx, (CUdeviceptr)output.storage().data(), false);
|
|
// Create JIT
|
|
triton::jit jit(ctx);
|
|
std::vector<unsigned> params = {
|
|
16, 2, 64,
|
|
32, 2, 64,
|
|
16, 8, 2, 2,
|
|
8, 1, 8,
|
|
4
|
|
};
|
|
jit.add_module("conv", src, params);
|
|
triton::driver::kernel* kernel = jit.get_function("conv");
|
|
triton::jit::launch_information info = jit.get_launch_info("conv");
|
|
// launch info
|
|
unsigned TM = info.global_range_size[0];
|
|
unsigned TN = info.global_range_size[1];
|
|
unsigned TK = jit.get_int("TK");
|
|
// initialize constant memory
|
|
int FS = T*R*S;
|
|
int nlut = (TK + FS - 1) / FS * FS;
|
|
std::vector<int> h_delta(nlut + upsample_d*upsample_h*upsample_w*nlut);
|
|
std::vector<int> h_masks(nlut + (2*pad_h+1)*(2*pad_w+1)*(2*pad_d+1)*nlut);
|
|
// memory stride for images
|
|
int32_t stride_i_w = 1;
|
|
int32_t stride_i_h = W*stride_i_w;
|
|
int32_t stride_i_d = H*stride_i_h;
|
|
int32_t stride_i_c = 1*stride_i_d;
|
|
int32_t stride_i_n = Ci*stride_i_c;
|
|
// memory stride for activations
|
|
int32_t stride_o_q = 1;
|
|
int32_t stride_o_p = Q*stride_o_q;
|
|
int32_t stride_o_m = P*stride_o_p;
|
|
int32_t stride_o_k = 1*stride_o_m;
|
|
int32_t stride_o_n = NF*stride_o_k;
|
|
build_conv_lut(TK, stride_i_d, stride_i_h, stride_i_w, stride_i_c, pad_d, pad_h, pad_w, T, R, S, h_delta, h_masks);
|
|
// equivalent matmul dimensions
|
|
int32_t M = B*P*Q;
|
|
int32_t N = NF;
|
|
int32_t K = Ci*R*S;
|
|
triton::driver::buffer* delta = jit.get_buffer("delta");
|
|
triton::driver::buffer* masks = jit.get_buffer("masks");
|
|
stream->write(delta, false, 0, h_delta.size()*4, h_delta.data());
|
|
stream->write(masks, false, 0, h_masks.size()*4, h_masks.data());
|
|
// launch info
|
|
unsigned nthreads = info.num_threads;
|
|
std::array<size_t, 3> grid = {(M + TM - 1)/TM, (N + TN - 1)/TN, 1};
|
|
// fast bounds-checking
|
|
unsigned lasti = (grid[0]*TM - 1)*TM + TM - 1;
|
|
unsigned lastj = (grid[1]*TN - 1)*TN + TN - 1;
|
|
unsigned lastk = TK - 1;
|
|
bool AT = false;
|
|
bool BT = true;
|
|
unsigned last_safe_a = (AT==false)?(M*K - 1 - lasti)/M - lastk : M*K - 1 - lasti*K - lastk;
|
|
unsigned last_safe_b = (BT==true)?(N*K - 1 - lastj)/N - lastk : N*K - 1 - lastj*K - lastk;
|
|
int32_t bound = std::max<unsigned>(1, std::max(K - last_safe_a, K - last_safe_b));
|
|
// set arguments
|
|
kernel->setArg(0, *d.cu());
|
|
kernel->setArg(1, *w.cu());
|
|
kernel->setArg(2, *a.cu());
|
|
kernel->setArg(3, M);
|
|
kernel->setArg(4, N);
|
|
kernel->setArg(5, K);
|
|
kernel->setArg(6, B);
|
|
kernel->setArg(7, H);
|
|
kernel->setArg(8, W);
|
|
kernel->setArg(9, NF);
|
|
kernel->setArg(10, P);
|
|
kernel->setArg(11, Q);
|
|
kernel->setArg(12, Ci);
|
|
kernel->setArg(13, R);
|
|
kernel->setArg(14, S);
|
|
kernel->setArg(15, stride_i_n);
|
|
kernel->setArg(16, stride_i_c);
|
|
kernel->setArg(17, stride_i_h);
|
|
kernel->setArg(18, stride_i_w);
|
|
kernel->setArg(19, stride_o_n);
|
|
kernel->setArg(20, stride_o_k);
|
|
kernel->setArg(21, stride_o_p);
|
|
kernel->setArg(22, stride_o_q);
|
|
kernel->setArg(23, pad_h);
|
|
kernel->setArg(24, pad_w);
|
|
kernel->setArg(25, bound);
|
|
// // dry run
|
|
stream->enqueue(kernel, grid, {nthreads, 1, 1});
|
|
return output;
|
|
}
|
|
|
|
static auto registry =
|
|
torch::jit::RegisterOperators("triton::conv_forward", &conv_forward);
|