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

[codegen][selection] bugfix in scanline dot lowering

Browse Source
This commit is contained in:
Philippe Tillet
2019-09-01 16:30:53 -04:00
parent 2d4ddab4d0
commit 90d80c3b2e
5 changed files with 59 additions and 63 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
lib/codegen/analysis/tune.cc
View File

@@ -308,7 +308,6 @@ void grids::run(ir::module &mod) {
std::string str_d = std::to_string(d);
effective_num_threads *= params_.at(i).at("mts.d" + str_d)->get_value();
}
if(num_threads != effective_num_threads)
throw std::runtime_error("cannot create a kernel with this amount of warps");
}

2
lib/codegen/selection/selection.cc
View File

@@ -1209,7 +1209,7 @@ void selection::lower_scanline_dot(ir::dot_inst *dot, LLVMContext &ctx, Function
TA->set_vector_size(TC->axis(0).contiguous);
TB->set_vector_size(TC->axis(1).contiguous);
TC->for_each([&](indices_t idx){
Value *res = TC->get_value(idx);
Value *res = TD->get_value(idx);
for(unsigned K = 0; K < NK; ++K){
// input indices
indices_t a_idx = {idx[0], builder.getInt32(K)};

2
lib/runtime/function.cc
View File

@@ -217,7 +217,7 @@ std::unique_ptr<driver::module> function::make_bin(ir::module &module, driver::c
dce.run(module);
vectorize.run(module);
dce.run(module);
ir::print(module, std::cout);
// ir::print(module, std::cout);
// generate llvm code
llvm::LLVMContext ctx;
std::unique_ptr<llvm::Module> llvm(new llvm::Module(module.get_name(), ctx));

30
tests/bench/dot.cc
View File

@@ -27,8 +27,8 @@ inline rt::function::grid_fn_ty grid(size_t M, size_t N) {
std::vector<double> do_bench(drv::stream* stream, bool AT, bool BT, int32_t M, int32_t N, int32_t K){
typedef half_float::half NumericT;
std::string ty = "half";
typedef float NumericT;
std::string ty = "float";
size_t dt_nbytes = sizeof(NumericT);
drv::context* context = stream->context();
// leading dimensions
@@ -46,25 +46,25 @@ std::vector<double> do_bench(drv::stream* stream, bool AT, bool BT, int32_t M, i
opt.defines.push_back({"AT", {""}});
if(BT)
opt.defines.push_back({"BT", {""}});
opt.defines.push_back({"TM", {"64", "128"}});
opt.defines.push_back({"TN", {"128"}});
opt.defines.push_back({"TK", {"32"}});
opt.defines.push_back({"TM", {"64"}});
opt.defines.push_back({"TN", {"64"}});
opt.defines.push_back({"TK", {"8"}});
opt.num_warps = {4};
// create function
rt::function function(src::dot, opt);
// benchmark available libraries
std::vector<double> result;
auto tflops = [&](double nanosec) { return 2.*M*N*K / nanosec * 1e-3; };
// cublas
if(cublas::cublasinit()){
NumericT alpha(static_cast<double>(1));
NumericT beta(static_cast<double>(0));
cublasGemmAlgo_t fastest;
cublasGemm(CUDA_R_16F, stream, AT, BT, M, N, K, &alpha, &*da, lda, &*db, ldb, &beta, &*dc, ldc, &fastest);
double cublas_ms = triton::tools::bench([&]() { cublasGemm(CUDA_R_16F, stream, AT, BT, M, N, K,
&alpha, &*da, lda, &*db, ldb, &beta, &*dc, ldc, nullptr, fastest); }, stream);
result.push_back(tflops(cublas_ms));
}
// // cublas
// if(cublas::cublasinit()){
// NumericT alpha(static_cast<double>(1));
// NumericT beta(static_cast<double>(0));
// cublasGemmAlgo_t fastest;
// cublasGemm(CUDA_R_16F, stream, AT, BT, M, N, K, &alpha, &*da, lda, &*db, ldb, &beta, &*dc, ldc, &fastest);
// double cublas_ms = triton::tools::bench([&]() { cublasGemm(CUDA_R_16F, stream, AT, BT, M, N, K,
// &alpha, &*da, lda, &*db, ldb, &beta, &*dc, ldc, nullptr, fastest); }, stream);
// result.push_back(tflops(cublas_ms));
// }
// triton
double triton_ms = triton::tools::bench([&]() { function({&*da, &*db, &*dc, M, N, K, lda, ldb, ldc}, grid(M, N), stream);}, stream);
result.push_back(tflops(triton_ms));

87
tests/common/src/dot.h
View File

@@ -4,74 +4,71 @@ namespace src {
R"(
#ifdef AT
#define USEA ^a
#define STRIDE_AK lda
#define STRIDE_AM 1
#define BROADCAST_AK :, newaxis
#define BROADCAST_AM newaxis, :
#define SHAPE_A TK, TM
#else
#define USEA a
#define STRIDE_AK 1
#define STRIDE_AM lda
#define BROADCAST_AK newaxis, :
#define BROADCAST_AM :, newaxis
#define SHAPE_A TM, TK
#endif
#ifdef BT
#define USEB ^b
#define STRIDE_BK 1
#define STRIDE_BN ldb
#define BROADCAST_BK newaxis, :
#define BROADCAST_BN :, newaxis
#define SHAPE_B TN, TK
#else
#define USEB b
#define STRIDE_BK ldb
#define STRIDE_BN 1
#define BROADCAST_BK :, newaxis
#define BROADCAST_BN newaxis, :
#define SHAPE_B TK, TN
#endif
void dot(TYPE * A __noalias __readonly __aligned(16),
TYPE * B __noalias __readonly __aligned(16),
TYPE * C __noalias __readonly __aligned(16),
void dot(TYPE * A, TYPE * B, TYPE * C,
int M, int N, int K,
int lda __multipleof(8),
int ldb __multipleof(8),
int ldc) {
// prologue
int ridx = get_program_id(0);
int ridy = get_program_id(1);
int rxa[TM] = ridx * TM + 0 ... TM;
int ryb[TN] = ridy * TN + 0 ... TN;
int rka[TK] = 0 ... TK;
int rkb[TK] = 0 ... TK;
float xc[TM, TN] = 0;
#ifdef AT
TYPE* pa[TK, TM] = A + rka[:, newaxis] + rxa[newaxis, :]*lda;
bool checka[TK, TM] = rka[:, newaxis] < TK;
TYPE a[TK, TM] = checka ? *pa : 0;
#else
TYPE* pa[TM, TK] = A + rka[newaxis, :]*lda + rxa[:, newaxis];
bool checka[TM, TK] = rka[newaxis, :] < TK;
TYPE a[TM, TK] = checka ? *pa : 0;
#endif
#ifdef BT
TYPE* pb[TN, TK] = B + rkb[newaxis, :]*ldb + ryb[:, newaxis];
bool checkb[TN, TK] = rkb[newaxis, :] < TK;
TYPE b[TN, TK] = checkb ? *pb : 0;
#else
TYPE* pb[TK, TN] = B + rkb[:, newaxis] + ryb[newaxis, :]*ldb;
bool checkb[TK, TN] = rkb[:, newaxis] < TK;
TYPE b[TK, TN] = checkb ? *pb : 0;
#endif
for(int k = K; k > 0; k = k - TK){
xc = USEA @ USEB + xc;
#ifdef AT
pa = pa + TK;
#else
pa = pa + TK*lda;
#endif
#ifdef BT
pb = pb + TK*ldb;
#else
pb = pb + TK;
#endif
checka = k > TK;
checkb = k > TK;
a = checka ? *pa : 0;
b = checkb ? *pb : 0;
float c[TM, TN] = 0;
// pointers to operands
TYPE* pa[SHAPE_A] = A + rka[BROADCAST_AK] * STRIDE_AK + rxa[BROADCAST_AM] * STRIDE_AM;
TYPE* pb[SHAPE_B] = B + rkb[BROADCAST_BK] * STRIDE_BK + ryb[BROADCAST_BN] * STRIDE_BN;
// prefetches operands
TYPE a[SHAPE_A] = *pa;
TYPE b[SHAPE_B] = *pb;
// reduction loop
for(int k = K; k > 0; k-= TK){
c += USEA @ USEB;
pa = pa + TK * STRIDE_AK;
pb = pb + TK * STRIDE_BK;
a = *pa;
b = *pb;
}
int rxc[TM] = ridx * TM + (0 ... TM);
int ryc[TN] = ridy * TN + (0 ... TN);
TYPE* pc[TM, TN] = C + ryc[newaxis, :]*ldc + rxc[:, newaxis];
TYPE c[TM, TN] = xc;
bool checkc0[TM] = rxc < M;
bool checkc1[TN] = ryc < N;
bool checkc[TM, TN] = checkc0[:, newaxis] && checkc1[newaxis, :];
// epilogue
int rxc[TM] = ridx * TM + 0 ... TM;
int ryc[TN] = ridy * TN + 0 ... TN;
TYPE* pc[TM, TN] = C + ryc[newaxis, :] + rxc[:, newaxis] * ldc;
bool checkc[TM, TN] = (rxc < M)[:, newaxis] && (ryc < N)[newaxis, :];
*?(checkc) pc = c;
}
)";
}