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

more work on heuristics

Browse Source
This commit is contained in:
Philippe Tillet
2019-07-21 18:11:54 -07:00
parent 484e3871cf
commit b1d81a5802
17 changed files with 268 additions and 99 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
examples/python/tensorflow/dot.cpp
View File

@@ -49,7 +49,7 @@ class DotOp : public OpKernel {
triton::driver::cu_buffer db(ctx, (CUdeviceptr)b.flat<Eigen::half>().data(), false);
triton::driver::cu_buffer dc(ctx, (CUdeviceptr)c->flat<float>().data(), false);
// template
triton::dnn::gemm dot(M, N, K, false, false, "fp16", "fp16", 4, 4);
triton::dnn::dot dot(M, N, K, false, false, "fp16", "fp16", 8, 8);
dot.enqueue(stream, {&da, &db, &dc});
}

2
examples/python/tensorflow/shift.cpp
View File

@@ -19,7 +19,7 @@
using namespace tensorflow;
using GPUDevice = Eigen::GpuDevice;
template<triton::dnn::shift::type OP>
template<triton::dnn::shift::op_t OP>
class ShiftConvOp : public OpKernel {
public:
explicit ShiftConvOp(OpKernelConstruction* context) : OpKernel(context), layout_(triton::dnn::shift::NCHW) {