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

[python][examples] some more cleaning of dot product example

Browse Source
This commit is contained in:
Philippe Tillet
2019-08-30 17:05:03 -07:00
parent 7e0af2118c
commit 5db3a7adfe
2 changed files with 14 additions and 20 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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); dce.run(module);
vectorize.run(module); vectorize.run(module);
dce.run(module); dce.run(module);
// ir::print(module, std::cout); ir::print(module, std::cout);
// generate llvm code // generate llvm code
llvm::LLVMContext ctx; llvm::LLVMContext ctx;
std::unique_ptr<llvm::Module> llvm(new llvm::Module(module.get_name(), ctx)); std::unique_ptr<llvm::Module> llvm(new llvm::Module(module.get_name(), ctx));

32
python/examples/dot.py
View File

@@ -35,45 +35,39 @@ src = """
#define SHAPE_B TK, TN #define SHAPE_B TK, TN
#endif #endif
void dot(TYPE * A, void dot(TYPE * A, TYPE * B, TYPE * C,
TYPE * B,
TYPE * C,
int M, int N, int K, int M, int N, int K,
int lda __multipleof(8), int lda __multipleof(8),
int ldb __multipleof(8), int ldb __multipleof(8),
int ldc) { int ldc) {
// prologue
/* prologue */
int ridx = get_program_id(0); int ridx = get_program_id(0);
int ridy = get_program_id(1); int ridy = get_program_id(1);
int rxa[TM] = ridx * TM + 0 ... TM; int rxa[TM] = ridx * TM + 0 ... TM;
int ryb[TN] = ridy * TN + 0 ... TN; int ryb[TN] = ridy * TN + 0 ... TN;
int rka[TK] = 0 ... TK; int rka[TK] = 0 ... TK;
int rkb[TK] = 0 ... TK; int rkb[TK] = 0 ... TK;
float xc[TM, TN] = 0; float c[TM, TN] = 0;
/* pointers for operands */ // pointers to operands
TYPE* pa[SHAPE_A] = A + rka[BROADCAST_AK] * STRIDE_AK + rxa[BROADCAST_AM] * STRIDE_AM; 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; TYPE* pb[SHAPE_B] = B + rkb[BROADCAST_BK] * STRIDE_BK + ryb[BROADCAST_BN] * STRIDE_BN;
/* prefetches operands */ // prefetches operands
TYPE a[SHAPE_A] = *pa; TYPE a[SHAPE_A] = *pa;
TYPE b[SHAPE_B] = *pb; TYPE b[SHAPE_B] = *pb;
/* reduction loop */ // reduction loop
for(int k = K; k > 0; k = k - TK){ for(int k = K; k > 0; k-= TK){
xc = USEA @ USEB + xc; c += USEA @ USEB;
pa = pa + TK * STRIDE_AK; pa = pa + TK * STRIDE_AK;
pb = pb + TK * STRIDE_BK; pb = pb + TK * STRIDE_BK;
a = *pa; a = *pa;
b = *pb; b = *pb;
} }
/* epilogue */ // epilogue
int rxc[TM] = ridx * TM + (0 ... TM); int rxc[TM] = ridx * TM + 0 ... TM;
int ryc[TN] = ridy * TN + (0 ... TN); int ryc[TN] = ridy * TN + 0 ... TN;
TYPE* pc[TM, TN] = C + ryc[newaxis, :] + rxc[:, newaxis] * ldc; TYPE* pc[TM, TN] = C + ryc[newaxis, :] + rxc[:, newaxis] * ldc;
TYPE c[TM, TN] = xc; bool checkc[TM, TN] = (rxc < M)[:, newaxis] && (ryc < N)[newaxis, :];
bool checkc0[TM] = rxc < M; *?(checkc) pc = c;
bool checkc1[TN] = ryc < N;
bool checkc[TM, TN] = checkc0[:, newaxis] && checkc1[newaxis, :];
*pc = c;
} }
""" """