From 66006ec26e75957464ba1867a90dc771350b7c7f Mon Sep 17 00:00:00 2001 From: Rawn Henry Date: Fri, 9 Oct 2020 16:14:32 -0700 Subject: [PATCH 01/19] Added new inference operator that peforms bias addition and optionally RELU for inference. When upgrading to cuda 11, the bias and relu can be fused into the matrix multiply with cublasLt. --- src/CMakeLists.txt | 1 + src/graph/expression_operators.cpp | 39 +++++++++---- src/graph/expression_operators.h | 3 +- src/graph/node_operators_binary.h | 92 ++++++++++++++++++++++++++++++ src/layers/generic.h | 14 +++++ src/models/transformer.h | 12 +++- src/tensors/cpu/prod.cpp | 14 +++++ src/tensors/dispatch.h | 24 ++++++++ src/tensors/gpu/prod.cpp | 17 ++++++ src/tensors/gpu/prod.cu | 69 ++++++++++++++++++++++ src/tensors/gpu/prod.h | 15 +++++ src/tensors/tensor_operators.h | 2 + 12 files changed, 287 insertions(+), 15 deletions(-) create mode 100644 src/tensors/gpu/prod.cu diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index 6dcf7fd89..426bf7e07 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -164,6 +164,7 @@ if(CUDA_FOUND) tensors/gpu/device.cu tensors/gpu/algorithm.cu tensors/gpu/prod.cpp + tensors/gpu/prod.cu tensors/gpu/topk.cu tensors/gpu/element.cu tensors/gpu/add.cu diff --git a/src/graph/expression_operators.cpp b/src/graph/expression_operators.cpp index f571f9f8b..f700d75e7 100644 --- a/src/graph/expression_operators.cpp +++ b/src/graph/expression_operators.cpp @@ -521,6 +521,13 @@ Expr bdot(Expr a, Expr b, bool transA, bool transB, float scale) { return Expression(a, b, transA, transB, scale); } +// A fused version of affine for GPU inference +static Expr affineFused(Expr a, Expr b, Expr bias, bool transA, bool transB, float scale, bool do_relu) { + float clipValue = a->graph()->getBackend()->getClip(); + std::vector nodes = { clip(a, clipValue), clip(b, clipValue), bias}; + return Expression(nodes, transA, transB, scale, do_relu); +} + static Expr affineDefault(Expr a, Expr b, Expr bias, bool transA, bool transB, float scale) { // general version, MKL, CBlas or CUDA @@ -542,7 +549,7 @@ static Expr affineDefault(Expr a, Expr b, Expr bias, bool transA, bool transB, f // youki/packed-model-pr-backup1031 // https://machinetranslation.visualstudio.com/Marian/_git/marian-dev?version=GByouki%2Fpacked-model-pr-backup1031 // SHA: 3456a7ed1d1608cfad74cd2c414e7e8fe141aa52 -Expr affine(Expr a, Expr b, Expr bias, bool transA, bool transB, float scale) { +Expr affine(Expr a, Expr b, Expr bias, bool transA, bool transB, float scale, bool do_relu) { auto device = a->graph()->getDeviceId().type; float clipValue = a->graph()->getBackend()->getClip(); @@ -550,16 +557,17 @@ Expr affine(Expr a, Expr b, Expr bias, bool transA, bool transB, float scale) { Type bElementType = b->value_type(); if(device == DeviceType::cpu) { + Expr affineTransform; if(isFloat(aElementType) && isFloat(bElementType)) { if(a->graph()->getBackend()->isOptimized()) { // cpu int16 version - return cpu::int16::affine( + affineTransform = cpu::int16::affine( cpu::int16::quantize(transA ? transpose(a) : a, clipValue), cpu::int16::quantize(transB ? b : transpose(b), clipValue), bias, scale); } else { - return affineDefault(a, b, bias, transA, transB, scale); + affineTransform = affineDefault(a, b, bias, transA, transB, scale); } } else if(isFloat(aElementType) && isPacked(bElementType)) { #if USE_FBGEMM @@ -570,13 +578,13 @@ Expr affine(Expr a, Expr b, Expr bias, bool transA, bool transB, float scale) { // and this cpu lookup is executed only once and the state is kept in FBGEMM. if(fbgemm::fbgemmHasAvx2Support()) { // This variant of affine product can handle matrix multiplications with packed8 and packed16 weight matrix (B). - return cpu::variant::affine(clip(a, clipValue), - b, - b->shape(), - bias, - transA, - transB, - scale); + affineTransform = cpu::variant::affine(clip(a, clipValue), + b, + b->shape(), + bias, + transA, + transB, + scale); } else { ABORT("AVX2 is not available. At least, AVX2 is needed to use fbgemm-based packed GEMM"); } @@ -586,12 +594,21 @@ Expr affine(Expr a, Expr b, Expr bias, bool transA, bool transB, float scale) { } else { ABORT("Combination of types A: {} B: {} not supported", aElementType, bElementType); } + if(do_relu) + return relu(affineTransform); + return affineTransform; } else { // Default GEMM ABORT_IF(!isFloat(aElementType) || !isFloat(bElementType), "GPU-based GEMM only supports float types, you have A: {} and B: {}", aElementType, bElementType); - return affineDefault(a, b, bias, transA, transB, scale); + if(a->graph()->isInference()) { + return affineFused(a, b, bias, transA, transB, scale, do_relu); + } + Expr affineTransform = affineDefault(a, b, bias, transA, transB, scale); + if (do_relu) + return relu(affineTransform); + return affineTransform; } } diff --git a/src/graph/expression_operators.h b/src/graph/expression_operators.h index a4a2eeee4..057daacfb 100755 --- a/src/graph/expression_operators.h +++ b/src/graph/expression_operators.h @@ -147,7 +147,8 @@ Expr affine(Expr a, Expr c, bool transA = false, bool transB = false, - float scalar = 1.f); + float scalar = 1.f, + bool do_relu = false); Expr csr_dot(const Shape& A_shape, Expr Avalues, Expr Aindices, Expr Aoffsets, Expr B, bool transA = false); Expr dot_csr(Expr A, const Shape& B_shape, Expr B_values, Expr B_indices, Expr B_offsets, bool transB = false); diff --git a/src/graph/node_operators_binary.h b/src/graph/node_operators_binary.h index d596619b2..46ec4173d 100644 --- a/src/graph/node_operators_binary.h +++ b/src/graph/node_operators_binary.h @@ -380,6 +380,98 @@ class AffineNodeOp : public NaryNodeOp { }; +class FusedAffineNodeOp : public NaryNodeOp { +private: + friend class SerializationHelpers; + bool transA_; + bool transB_; + float scalar_; + bool do_relu_; + +public: + FusedAffineNodeOp(const std::vector& nodes, + bool transA, + bool transB, + float scalar, + bool do_relu=false) + : NaryNodeOp(nodes, newShape(nodes[0], nodes[1], transA, transB)), + transA_(transA), + transB_(transB), + scalar_(scalar), + do_relu_(do_relu) {} + + Shape newShape(Expr a, Expr b, bool transA, bool transB) { + auto shapeA = a->shape(); + if(transA) { + shapeA.set(shapeA.size() - 2, a->shape()[shapeA.size() - 1]); + shapeA.set(shapeA.size() - 1, a->shape()[shapeA.size() - 2]); + } + + auto shapeB = b->shape(); + if(transB) { + shapeB.set(shapeB.size() - 2, b->shape()[shapeB.size() - 1]); + shapeB.set(shapeB.size() - 1, b->shape()[shapeB.size() - 2]); + } + + Shape outShape = shapeA; + outShape.set(outShape.size() - 1, shapeB[shapeB.size() - 1]); + ABORT_IF(shapeA[shapeA.size() - 1] != shapeB[shapeB.size() - 2], + "Matrix product requires inner dimensions to match in {}{} * {}{}", std::string(shapeA), transA, std::string(shapeB), transB); + return outShape; + } + + NodeOps forwardOps() override { + using namespace functional; + + return { + NodeOp( + Affine(val_, + graph()->allocator(), + child(0)->val(), + child(1)->val(), + child(2)->val(), + transA_, + transB_, + 0.f, + scalar_, + do_relu_)) + }; + } + + NodeOps backwardOps() override { + ABORT("Node only supports inference."); + } + + const std::string type() override { return "fusedAffine"; } + + virtual size_t hash() override { + size_t seed = NaryNodeOp::hash(); + util::hash_combine(seed, transA_); + util::hash_combine(seed, transB_); + util::hash_combine(seed, scalar_); + util::hash_combine(seed, do_relu_); + return seed; + } + + virtual bool equal(Expr node) override { + if(!NaryNodeOp::equal(node)) + return false; + auto cnode = std::dynamic_pointer_cast(node); + if(!cnode) + return false; + if(transA_ != cnode->transA_) + return false; + if(transB_ != cnode->transB_) + return false; + if(scalar_ != cnode->scalar_) + return false; + if(do_relu_ != cnode->do_relu_) + return false; + return true; + } + +}; + class DotBatchedNodeOp : public NaryNodeOp { private: friend class SerializationHelpers; diff --git a/src/layers/generic.h b/src/layers/generic.h index e83663357..14e385a6e 100755 --- a/src/layers/generic.h +++ b/src/layers/generic.h @@ -1,5 +1,7 @@ #pragma once +#include "common/definitions.h" +#include "graph/expression_operators.h" #include "marian.h" #include "data/shortlist.h" @@ -455,6 +457,18 @@ Expr denseInline(Expr x, std::string prefix, std::string suffix, int outDim, con return x; } +static inline +Expr denseInlineRelu(Expr x, std::string prefix, std::string suffix, int outDim, float dropProb = 0.0f) +{ + auto graph = x->graph(); + + auto W = graph->param(prefix + "_W" + suffix, { x->shape()[-1], outDim }, inits::glorotUniform()); + auto b = graph->param(prefix + "_b" + suffix, { 1, outDim }, inits::zeros()); + x = affine(x, W, b, false, false, 1.f, true); + x = dropout(x, dropProb); + return x; +} + static inline Expr layerNorm(Expr x, std::string prefix, std::string suffix = std::string()) { int dimModel = x->shape()[-1]; diff --git a/src/models/transformer.h b/src/models/transformer.h index 2d9ced33c..de6304002 100755 --- a/src/models/transformer.h +++ b/src/models/transformer.h @@ -411,17 +411,23 @@ class Transformer : public EncoderOrDecoderBase { auto opsPre = opt("transformer-preprocess"); auto output = preProcess(prefix + "_ffn", opsPre, input, dropProb); + auto actName = opt("transformer-ffn-activation"); int dimFfn = opt("transformer-dim-ffn"); int depthFfn = opt("transformer-ffn-depth"); - auto actFn = activationByName(opt("transformer-ffn-activation")); + auto actFn = activationByName(actName); float ffnDropProb = inference_ ? 0 : opt("transformer-dropout-ffn"); ABORT_IF(depthFfn < 1, "Filter depth {} is smaller than 1", depthFfn); // the stack of FF layers - for(int i = 1; i < depthFfn; ++i) - output = denseInline(output, prefix, /*suffix=*/std::to_string(i), dimFfn, actFn, ffnDropProb); + for(int i = 1; i < depthFfn; ++i) { + if (actName == "relu") { + output = denseInlineRelu(output, prefix, std::to_string(i), dimFfn, ffnDropProb); + } else { + output = denseInline(output, prefix, /*suffix=*/std::to_string(i), dimFfn, actFn, ffnDropProb); + } + } output = denseInline(output, prefix, /*suffix=*/std::to_string(depthFfn), dimModel); auto opsPost = opt("transformer-postprocess"); diff --git a/src/tensors/cpu/prod.cpp b/src/tensors/cpu/prod.cpp index 8529db8b5..6112db7ee 100755 --- a/src/tensors/cpu/prod.cpp +++ b/src/tensors/cpu/prod.cpp @@ -14,6 +14,20 @@ namespace marian { namespace cpu { +void Affine(marian::Tensor C, + Ptr allocator, + const marian::Tensor& A, + const marian::Tensor& B, + const marian::Tensor& bias, + bool transA, + bool transB, + float beta, + float scalar, + bool relu_postprocess) { + + ABORT("Not supported on CPU"); + } + void Prod(marian::Tensor C, const marian::Tensor& A, const marian::Tensor& B, diff --git a/src/tensors/dispatch.h b/src/tensors/dispatch.h index 094f156cb..4ec1a6fe6 100644 --- a/src/tensors/dispatch.h +++ b/src/tensors/dispatch.h @@ -152,6 +152,30 @@ cpu::Function(arg1, arg2, arg3, arg4, arg5, arg6, arg7, arg8, arg9); \ } +#define DISPATCH10( \ + Function, Arg1, Arg2, Arg3, Arg4, Arg5, Arg6, Arg7, Arg8, Arg9, Arg10) \ +namespace gpu { \ +void Function(Arg1, Arg2, Arg3, Arg4, Arg5, Arg6, Arg7, Arg8, Arg9, Arg10); \ +} \ +namespace cpu { \ +void Function(Arg1, Arg2, Arg3, Arg4, Arg5, Arg6, Arg7, Arg8, Arg9, Arg10); \ +} \ +static inline void Function(Arg1 arg1, \ + Arg2 arg2, \ + Arg3 arg3, \ + Arg4 arg4, \ + Arg5 arg5, \ + Arg6 arg6, \ + Arg7 arg7, \ + Arg8 arg8, \ + Arg9 arg9, \ + Arg10 arg10) { \ + if(arg1->getBackend()->getDeviceId().type == DeviceType::gpu) \ + gpu::Function(arg1, arg2, arg3, arg4, arg5, arg6, arg7, arg8, arg9, arg10); \ + else \ + cpu::Function(arg1, arg2, arg3, arg4, arg5, arg6, arg7, arg8, arg9, arg10); \ +} + #else #define DISPATCH1(Function, Arg1) \ diff --git a/src/tensors/gpu/prod.cpp b/src/tensors/gpu/prod.cpp index 071b228ff..d149f36e5 100755 --- a/src/tensors/gpu/prod.cpp +++ b/src/tensors/gpu/prod.cpp @@ -539,5 +539,22 @@ void CSRProd(marian::Tensor C, if(St_offsets) allocator->free(St_offsets); } +// NOTE: The allocator isn't used currently but is useful for CUDA >= 11.0.3 with cublasLt so that a workspace can be allocated. +// Earlier versions of cublasLT do not support bias addition for fp32 and fp16. +void Affine(marian::Tensor C, Ptr allocator, const marian::Tensor& A, const marian::Tensor& B, const marian::Tensor& bias, + bool transA, bool transB, float beta, float scalar, bool do_relu) { + + if(C->type() == Type::float32) { + ProdTyped(C, A, B, transA, transB, beta, scalar); +#if COMPILE_FP16 + } else if(C->type() == Type::float16) { + ProdTyped(C, A, B, transA, transB, __float2half(beta), __float2half(scalar)); +#endif + } else { + ABORT("Prod not implemented for type {}", C->type()); + } + BiasAdd(C, bias, do_relu); +} + } // namespace gpu } // namespace marian diff --git a/src/tensors/gpu/prod.cu b/src/tensors/gpu/prod.cu new file mode 100644 index 000000000..ec01d57e9 --- /dev/null +++ b/src/tensors/gpu/prod.cu @@ -0,0 +1,69 @@ +#include +#include "tensors/tensor.h" +#include "tensors/gpu/cuda_helpers.h" +#include "tensors/gpu/backend.h" + +namespace marian { +namespace gpu { + +template +__global__ static void gBiasAddFused(T* tensor, T* bias, size_t tensor_size, size_t bias_size, ActFunc f) { + const size_t row_start = blockIdx.x * bias_size; + for(int bias_offset = threadIdx.x; bias_offset < bias_size; bias_offset+=blockDim.x) { + size_t offset_into_tensor = row_start + bias_offset; + if(offset_into_tensor < tensor_size) { + T added_bias = tensor[offset_into_tensor] + bias[bias_offset]; + tensor[offset_into_tensor] = f(added_bias); + } + } +} + +struct identity { + template + __device__ constexpr T&& operator() (T&& t) const noexcept { + return std::forward(t); + } +}; + +struct reluAct { + template + __device__ T operator() (T t) const noexcept { + return t > (T) 0? t : (T) 0; + } +}; + +void BiasAdd(marian::Tensor C, const marian::Tensor& bias, bool do_relu) { + auto backend = std::static_pointer_cast(C->getBackend()); + CUDA_CHECK(cudaSetDevice(backend->getDeviceId().no)); + + size_t size = C->shape().elements(); + size_t bias_size = bias->shape().elements(); + + int m = C->shape().elements() / C->shape().back(); + int n = C->shape().back(); + + ABORT_IF(n != bias_size, "The number of elements in the bias must match the number of columns in C"); + + int threads_per_block = std::min(MAX_THREADS, n); + int blocks = m; + + if(C->type() == Type::float32) { + if (do_relu) + gBiasAddFused<<>>(C->data(), bias->data(), size, bias_size, reluAct()); + else + gBiasAddFused<<>>(C->data(), bias->data(), size, bias_size, identity()); + +#if COMPILE_FP16 + } else if(C->type() == Type::float16) { + if (do_relu) + gBiasAddFused<<>>(C->data(), bias->data(), size, bias_size, reluAct()); + else + gBiasAddFused<<>>(C->data(), bias->data(), size, bias_size, identity()); +#endif + } else { + ABORT("Prod not implemented for type {}", C->type()); + } +} + +} +} \ No newline at end of file diff --git a/src/tensors/gpu/prod.h b/src/tensors/gpu/prod.h index 9dc1220c4..ce7a35de9 100644 --- a/src/tensors/gpu/prod.h +++ b/src/tensors/gpu/prod.h @@ -8,6 +8,21 @@ namespace marian { namespace gpu { +void BiasAdd(marian::Tensor C, + const marian::Tensor& bias, + bool do_relu = false); + +void Affine(marian::Tensor C, + Ptr allocator, + const marian::Tensor& A, + const marian::Tensor& B, + const marian::Tensor& bias, + bool transA, + bool transB, + float beta = 0, + float scalar = 1, + bool do_relu = false); + void Prod(marian::Tensor C, const marian::Tensor& A, const marian::Tensor& B, diff --git a/src/tensors/tensor_operators.h b/src/tensors/tensor_operators.h index e075244f5..a6bf42d61 100644 --- a/src/tensors/tensor_operators.h +++ b/src/tensors/tensor_operators.h @@ -103,6 +103,8 @@ DISPATCH7(Prod, marian::Tensor, const marian::Tensor&, const marian::Tensor&, bo DISPATCH8(ProdBatched, marian::Tensor, Ptr, const marian::Tensor, const marian::Tensor, bool, bool, float, float) DISPATCH9(CSRProd, marian::Tensor, Ptr, const marian::Tensor&, const marian::Tensor&, const marian::Tensor&, const marian::Tensor&, bool, bool, float) +DISPATCH10(Affine, marian::Tensor, Ptr, const marian::Tensor&, const marian::Tensor&, const marian::Tensor&, bool, bool, float, float, bool) + DISPATCH2(Softmax, marian::Tensor, marian::Tensor) DISPATCH3(SoftmaxGrad, marian::Tensor, marian::Tensor, marian::Tensor) From 23a3223249c55c8e0353c53e09018ef2adf487a2 Mon Sep 17 00:00:00 2001 From: Rawn Henry Date: Sat, 10 Oct 2020 02:31:17 -0700 Subject: [PATCH 02/19] Fix compilation for CPU only. Missed DISPATCH10 call when CUDA was undefined --- src/tensors/dispatch.h | 18 ++++++++++++++++++ 1 file changed, 18 insertions(+) diff --git a/src/tensors/dispatch.h b/src/tensors/dispatch.h index 4ec1a6fe6..f71543511 100644 --- a/src/tensors/dispatch.h +++ b/src/tensors/dispatch.h @@ -272,4 +272,22 @@ static inline void Function(Arg1 arg1, cpu::Function(arg1, arg2, arg3, arg4, arg5, arg6, arg7, arg8, arg9); \ } +#define DISPATCH10( \ + Function, Arg1, Arg2, Arg3, Arg4, Arg5, Arg6, Arg7, Arg8, Arg9, Arg10) \ + namespace cpu { \ + void Function(Arg1, Arg2, Arg3, Arg4, Arg5, Arg6, Arg7, Arg8, Arg9, Arg10); \ + } \ + static inline void Function(Arg1 arg1, \ + Arg2 arg2, \ + Arg3 arg3, \ + Arg4 arg4, \ + Arg5 arg5, \ + Arg6 arg6, \ + Arg7 arg7, \ + Arg8 arg8, \ + Arg9 arg9, \ + Arg10 arg10) { \ + cpu::Function(arg1, arg2, arg3, arg4, arg5, arg6, arg7, arg8, arg9, arg10); \ + } + #endif From dbb1653a61109d906b5341230f430f202f3e403f Mon Sep 17 00:00:00 2001 From: Rawn Henry Date: Tue, 8 Dec 2020 20:03:47 -0800 Subject: [PATCH 03/19] Places fused affine under one node --- src/graph/expression_operators.cpp | 61 +++++++------- src/graph/node_operators_binary.h | 126 ++++++++--------------------- 2 files changed, 66 insertions(+), 121 deletions(-) diff --git a/src/graph/expression_operators.cpp b/src/graph/expression_operators.cpp index f700d75e7..d5d6c031a 100644 --- a/src/graph/expression_operators.cpp +++ b/src/graph/expression_operators.cpp @@ -1,4 +1,5 @@ #include "graph/expression_operators.h" +#include "common/definitions.h" #include "layers/constructors.h" #include "graph/node_operators.h" @@ -521,27 +522,35 @@ Expr bdot(Expr a, Expr b, bool transA, bool transB, float scale) { return Expression(a, b, transA, transB, scale); } -// A fused version of affine for GPU inference -static Expr affineFused(Expr a, Expr b, Expr bias, bool transA, bool transB, float scale, bool do_relu) { - float clipValue = a->graph()->getBackend()->getClip(); - std::vector nodes = { clip(a, clipValue), clip(b, clipValue), bias}; - return Expression(nodes, transA, transB, scale, do_relu); -} - -static Expr affineDefault(Expr a, Expr b, Expr bias, bool transA, bool transB, float scale) { +static Expr affineDefault(Expr a, Expr b, Expr bias, bool transA, bool transB, float scale, bool do_relu=false) { // general version, MKL, CBlas or CUDA // if clipValue > 0, the inputs will be clipped to range [-clipValue, // clipValue] This is meant to keep values at the same range as used during // training when optimizing for 8-bit integer products. Likely to be removed // in the future when we explore better ways to handle this. - float clipValue = a->graph()->getBackend()->getClip(); + auto g = a->graph(); + float clipValue = g->getBackend()->getClip(); + + std::vector nodes = { clip(a, clipValue), clip(b, clipValue), bias }; - int rows = a->shape().elements() / a->shape()[-1]; - Expr ones = a->graph()->ones({ rows, 1 }); - std::vector nodes - = { clip(a, clipValue), clip(b, clipValue), bias, ones }; - return Expression(nodes, transA, transB, scale); + // If we are using CPU, we broadcast the ones vector. On GPU, the bias addition can be fused into the GEMM with CUDA >= 11 + if(g->getBackend()->getDeviceId().type == DeviceType::cpu || !g->isInference()) { + int rows = a->shape().elements() / a->shape()[-1]; + Expr ones = g->ones({ rows, 1 }); + nodes.push_back(ones); + } + + if(do_relu) { + // For GPU inference, we can fuse the RELU into the bias addition. + if(g->isInference() && g->getBackend()->getDeviceId().type == DeviceType::gpu) { + return Expression(nodes, transA, transB, scale, do_relu); + } + Expr affineOp = Expression(nodes, transA, transB, scale, false); + return relu(affineOp); + } + + return Expression(nodes, transA, transB, scale, do_relu); } // This operation used to implement auto-tuning. We have removed it for now due to complexity, but plan to revisit it in the future. @@ -557,17 +566,19 @@ Expr affine(Expr a, Expr b, Expr bias, bool transA, bool transB, float scale, bo Type bElementType = b->value_type(); if(device == DeviceType::cpu) { - Expr affineTransform; if(isFloat(aElementType) && isFloat(bElementType)) { if(a->graph()->getBackend()->isOptimized()) { // cpu int16 version - affineTransform = cpu::int16::affine( + Expr affineTransform = cpu::int16::affine( cpu::int16::quantize(transA ? transpose(a) : a, clipValue), cpu::int16::quantize(transB ? b : transpose(b), clipValue), bias, scale); + if(do_relu) + affineTransform = relu(affineTransform); + return affineTransform; } else { - affineTransform = affineDefault(a, b, bias, transA, transB, scale); + return affineDefault(a, b, bias, transA, transB, scale, do_relu); } } else if(isFloat(aElementType) && isPacked(bElementType)) { #if USE_FBGEMM @@ -578,13 +589,16 @@ Expr affine(Expr a, Expr b, Expr bias, bool transA, bool transB, float scale, bo // and this cpu lookup is executed only once and the state is kept in FBGEMM. if(fbgemm::fbgemmHasAvx2Support()) { // This variant of affine product can handle matrix multiplications with packed8 and packed16 weight matrix (B). - affineTransform = cpu::variant::affine(clip(a, clipValue), + Expr affineTransform = cpu::variant::affine(clip(a, clipValue), b, b->shape(), bias, transA, transB, scale); + if(do_relu) + affineTransform = relu(affineTransform); + return affineTransform; } else { ABORT("AVX2 is not available. At least, AVX2 is needed to use fbgemm-based packed GEMM"); } @@ -594,21 +608,12 @@ Expr affine(Expr a, Expr b, Expr bias, bool transA, bool transB, float scale, bo } else { ABORT("Combination of types A: {} B: {} not supported", aElementType, bElementType); } - if(do_relu) - return relu(affineTransform); - return affineTransform; } else { // Default GEMM ABORT_IF(!isFloat(aElementType) || !isFloat(bElementType), "GPU-based GEMM only supports float types, you have A: {} and B: {}", aElementType, bElementType); - if(a->graph()->isInference()) { - return affineFused(a, b, bias, transA, transB, scale, do_relu); - } - Expr affineTransform = affineDefault(a, b, bias, transA, transB, scale); - if (do_relu) - return relu(affineTransform); - return affineTransform; + return affineDefault(a, b, bias, transA, transB, scale, do_relu); } } diff --git a/src/graph/node_operators_binary.h b/src/graph/node_operators_binary.h index 46ec4173d..a211783dd 100644 --- a/src/graph/node_operators_binary.h +++ b/src/graph/node_operators_binary.h @@ -218,16 +218,27 @@ class AffineNodeOp : public NaryNodeOp { bool transA_; bool transB_; float scalar_; + bool do_relu_; public: AffineNodeOp(const std::vector& nodes, bool transA, bool transB, - float scalar) + float scalar, + bool do_relu=false) : NaryNodeOp(nodes, newShape(nodes[0], nodes[1], transA, transB)), transA_(transA), transB_(transB), - scalar_(scalar) {} + scalar_(scalar), + do_relu_(do_relu) { + // ReLU fusion checks + ABORT_IF(do_relu && graph()->getBackend()->getDeviceId().type == DeviceType::cpu, "ReLU not fused for CPU backend"); + ABORT_IF(do_relu && !graph()->isInference(), "ReLU fusion is only supported for GPU inference."); + + // bias fusion checks + ABORT_IF(nodes.size() == 3 && graph()->getBackend()->getDeviceId().type == DeviceType::cpu, "Bias addition not specialized for CPU backend"); + // No check for if inference since the forward computation will be the same. + } Shape newShape(Expr a, Expr b, bool transA, bool transB) { auto shapeA = a->shape(); @@ -252,6 +263,24 @@ class AffineNodeOp : public NaryNodeOp { NodeOps forwardOps() override { using namespace functional; + if(children_.size() == 3) { + ABORT_IF(graph()->getBackend()->getDeviceId().type != DeviceType::gpu, "Snuck passed constructor check? Only supported for GPU backend"); + // We must be using the GPU backend (Would have aborted otherwise) + return { + NodeOp( + Affine(val_, + graph()->allocator(), + child(0)->val(), + child(1)->val(), + child(2)->val(), + transA_, + transB_, + 0.f, + scalar_, + do_relu_)) + }; + } + return { NodeOp( Prod(val_, @@ -272,7 +301,7 @@ class AffineNodeOp : public NaryNodeOp { // beta set to 1.0 in gemm, C = alpha * dot(op(A), op(B)) + beta * C // to sum gradients from different graph parts using namespace functional; - + ABORT_IF(do_relu_, "Snuck passed constructor check? ReLU fusion not supported for training."); if(!transA_ && transB_) return { NodeOp(Prod(child(0)->grad(), @@ -355,95 +384,6 @@ class AffineNodeOp : public NaryNodeOp { const std::string type() override { return "affine"; } - virtual size_t hash() override { - size_t seed = NaryNodeOp::hash(); - util::hash_combine(seed, transA_); - util::hash_combine(seed, transB_); - util::hash_combine(seed, scalar_); - return seed; - } - - virtual bool equal(Expr node) override { - if(!NaryNodeOp::equal(node)) - return false; - auto cnode = std::dynamic_pointer_cast(node); - if(!cnode) - return false; - if(transA_ != cnode->transA_) - return false; - if(transB_ != cnode->transB_) - return false; - if(scalar_ != cnode->scalar_) - return false; - return true; - } - -}; - -class FusedAffineNodeOp : public NaryNodeOp { -private: - friend class SerializationHelpers; - bool transA_; - bool transB_; - float scalar_; - bool do_relu_; - -public: - FusedAffineNodeOp(const std::vector& nodes, - bool transA, - bool transB, - float scalar, - bool do_relu=false) - : NaryNodeOp(nodes, newShape(nodes[0], nodes[1], transA, transB)), - transA_(transA), - transB_(transB), - scalar_(scalar), - do_relu_(do_relu) {} - - Shape newShape(Expr a, Expr b, bool transA, bool transB) { - auto shapeA = a->shape(); - if(transA) { - shapeA.set(shapeA.size() - 2, a->shape()[shapeA.size() - 1]); - shapeA.set(shapeA.size() - 1, a->shape()[shapeA.size() - 2]); - } - - auto shapeB = b->shape(); - if(transB) { - shapeB.set(shapeB.size() - 2, b->shape()[shapeB.size() - 1]); - shapeB.set(shapeB.size() - 1, b->shape()[shapeB.size() - 2]); - } - - Shape outShape = shapeA; - outShape.set(outShape.size() - 1, shapeB[shapeB.size() - 1]); - ABORT_IF(shapeA[shapeA.size() - 1] != shapeB[shapeB.size() - 2], - "Matrix product requires inner dimensions to match in {}{} * {}{}", std::string(shapeA), transA, std::string(shapeB), transB); - return outShape; - } - - NodeOps forwardOps() override { - using namespace functional; - - return { - NodeOp( - Affine(val_, - graph()->allocator(), - child(0)->val(), - child(1)->val(), - child(2)->val(), - transA_, - transB_, - 0.f, - scalar_, - do_relu_)) - }; - } - - NodeOps backwardOps() override { - ABORT("Node only supports inference."); - } - - const std::string type() override { return "fusedAffine"; } - virtual size_t hash() override { size_t seed = NaryNodeOp::hash(); util::hash_combine(seed, transA_); @@ -456,7 +396,7 @@ class FusedAffineNodeOp : public NaryNodeOp { virtual bool equal(Expr node) override { if(!NaryNodeOp::equal(node)) return false; - auto cnode = std::dynamic_pointer_cast(node); + auto cnode = std::dynamic_pointer_cast(node); if(!cnode) return false; if(transA_ != cnode->transA_) From 3e94764abb89d149d1042fb4e742ccef4654a0fd Mon Sep 17 00:00:00 2001 From: Rawn Henry Date: Wed, 9 Dec 2020 11:52:00 -0800 Subject: [PATCH 04/19] Adds basic affine support for cublaslt in cuda 11 --- CMakeLists.txt | 18 ++++- src/tensors/gpu/prod.cpp | 145 ++++++++++++++++++++++++++++++++++++++- 2 files changed, 160 insertions(+), 3 deletions(-) diff --git a/CMakeLists.txt b/CMakeLists.txt index 3c0150587..f6ee826d8 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -1,3 +1,7 @@ +# All or part of this file was contributed by NVIDIA under license: +# Copyright (C) 2020 NVIDIA Corporation +# SPDX-License-Identifier: MIT + cmake_minimum_required(VERSION 3.5.1) set(CMAKE_MODULE_PATH ${CMAKE_CURRENT_SOURCE_DIR}/cmake) @@ -338,8 +342,20 @@ if(CUDA_FOUND) endif() message(STATUS "Found CUDA libraries: ${CUDA_LIBS}") else(USE_STATIC_LIBS) + set(CUDA_LIBS ${CUDA_curand_LIBRARY} ${CUDA_cusparse_LIBRARY} ${CUDA_CUBLAS_LIBRARIES}) + # We actually only need cublasLt here after cuda 11. Marian will work fine without it pre cuda 11. We want to force CMake to use the cublas + # version that ships with CUDA 11 so we force the search to occur inside of the cuda toolkit directory. + set(CUDA_LIBS ${CUDA_curand_LIBRARY} ${CUDA_cusparse_LIBRARY} ${CUDA_CUBLAS_LIBRARIES}) + if ((CUDA_VERSION VERSION_EQUAL "11.0" OR CUDA_VERSION VERSION_GREATER "11.0")) + find_library(CUDA_cublasLt_LIBRARY NAMES cublasLt PATHS ${CUDA_TOOLKIT_ROOT_DIR}/lib64 ${CUDA_TOOLKIT_ROOT_DIR}/lib/x64 NO_DEFAULT_PATH) + if(NOT CUDA_cublasLt_LIBRARY) + message(FATAL_ERROR "cuBLASLt library not found") + endif() + set(EXT_LIBS ${EXT_LIBS} ${CUDA_cublasLt_LIBRARY}) + set(CUDA_LIBS ${CUDA_LIBS} ${CUDA_cublasLt_LIBRARY}) + endif() set(EXT_LIBS ${EXT_LIBS} ${CUDA_curand_LIBRARY} ${CUDA_cusparse_LIBRARY} ${CUDA_CUBLAS_LIBRARIES}) - message(STATUS "Found CUDA libraries: ${CUDA_curand_LIBRARY} ${CUDA_cusparse_LIBRARY} ${CUDA_CUBLAS_LIBRARIES}") + message(STATUS "Found CUDA libraries: ${CUDA_LIBS}") endif(USE_STATIC_LIBS) if(USE_CUDNN) diff --git a/src/tensors/gpu/prod.cpp b/src/tensors/gpu/prod.cpp index d149f36e5..dafabb71c 100755 --- a/src/tensors/gpu/prod.cpp +++ b/src/tensors/gpu/prod.cpp @@ -12,6 +12,10 @@ #include "tensors/gpu/cuda_helpers.h" // clang-format on +#if CUDA_VERSION >= 11000 +#include +#endif + namespace marian { namespace gpu { @@ -539,11 +543,147 @@ void CSRProd(marian::Tensor C, if(St_offsets) allocator->free(St_offsets); } -// NOTE: The allocator isn't used currently but is useful for CUDA >= 11.0.3 with cublasLt so that a workspace can be allocated. -// Earlier versions of cublasLT do not support bias addition for fp32 and fp16. +#if CUDA_VERSION >= 11000 +static cublasStatus_t cublasLtAffineHelper(cublasLtHandle_t ltHandle, cublasOperation_t transA, cublasOperation_t transB, + cudaDataType matrixType, + int m, int n, int k, const void *alpha, const void *A, int lda, const void *B, + int ldb, const void *beta, void *C, int ldc, const void* bias, + void* workspace, size_t workspaceSize, bool do_relu) { + + cublasLtMatmulDesc_t operationDesc = NULL; + cublasLtMatrixLayout_t Adesc = NULL, Bdesc = NULL, Cdesc = NULL; + cublasLtMatmulPreference_t preference = NULL; + + int returnedResults = 0; + cublasLtMatmulHeuristicResult_t heuristicResult = {}; + + cublasLtEpilogue_t epilogue = do_relu? CUBLASLT_EPILOGUE_RELU_BIAS: CUBLASLT_EPILOGUE_BIAS; + cublasComputeType_t computeType = matrixType == CUDA_R_32F? CUBLAS_COMPUTE_32F_FAST_16F: CUBLAS_COMPUTE_16F; + + + CUBLAS_CHECK(cublasLtMatmulDescCreate(&operationDesc, computeType, matrixType)); + CUBLAS_CHECK(cublasLtMatmulDescSetAttribute(operationDesc, CUBLASLT_MATMUL_DESC_TRANSA, &transA, sizeof(transA))); + CUBLAS_CHECK(cublasLtMatmulDescSetAttribute(operationDesc, CUBLASLT_MATMUL_DESC_TRANSB, &transB, sizeof(transB))); + CUBLAS_CHECK(cublasLtMatmulDescSetAttribute(operationDesc, CUBLASLT_MATMUL_DESC_EPILOGUE, &epilogue, sizeof(epilogue))); + CUBLAS_CHECK(cublasLtMatmulDescSetAttribute(operationDesc, CUBLASLT_MATMUL_DESC_BIAS_POINTER, &bias, sizeof(bias))); + + CUBLAS_CHECK(cublasLtMatrixLayoutCreate(&Adesc, matrixType, transA == CUBLAS_OP_N ? m : k, transA == CUBLAS_OP_N ? k : m, lda)); + CUBLAS_CHECK(cublasLtMatrixLayoutCreate(&Bdesc, matrixType, transB == CUBLAS_OP_N ? k : n, transB == CUBLAS_OP_N ? n : k, ldb)); + CUBLAS_CHECK(cublasLtMatrixLayoutCreate(&Cdesc, matrixType, m, n, ldc)); + + CUBLAS_CHECK(cublasLtMatmulPreferenceCreate(&preference)); + CUBLAS_CHECK(cublasLtMatmulPreferenceSetAttribute(preference, CUBLASLT_MATMUL_PREF_MAX_WORKSPACE_BYTES, &workspaceSize, sizeof(workspaceSize))); + CUBLAS_CHECK(cublasLtMatmulAlgoGetHeuristic(ltHandle, operationDesc, Adesc, Bdesc, Cdesc, Cdesc, preference, 1, &heuristicResult, &returnedResults)); + + cublasStatus_t opStatus = cublasLtMatmul(ltHandle, operationDesc, alpha, A, Adesc, B, Bdesc, beta, C, Cdesc, C, Cdesc, &heuristicResult.algo, + workspace, workspaceSize, cudaStreamPerThread); + + if (preference) CUBLAS_CHECK(cublasLtMatmulPreferenceDestroy(preference)); + if (Cdesc) CUBLAS_CHECK(cublasLtMatrixLayoutDestroy(Cdesc)); + if (Bdesc) CUBLAS_CHECK(cublasLtMatrixLayoutDestroy(Bdesc)); + if (Adesc) CUBLAS_CHECK(cublasLtMatrixLayoutDestroy(Adesc)); + if (operationDesc) CUBLAS_CHECK(cublasLtMatmulDescDestroy(operationDesc)); + + return opStatus; +} + +static cublasStatus_t cublasLtAffineTyped(cublasLtHandle_t ltHandle, cublasOperation_t transA, cublasOperation_t transB, + int m, int n, int k, const half *alpha, const half *A, int lda, const half *B, + int ldb, const half *beta, half *C, int ldc, const half* bias, + half* workspace, size_t workspaceSizeBytes, bool do_relu) { + return cublasLtAffineHelper(ltHandle, transA, transB, CUDA_R_16F, m, n, k, alpha, A, lda, B, ldb, beta, C, ldc, bias, + workspace, workspaceSizeBytes, do_relu); +} + +static cublasStatus_t cublasLtAffineTyped(cublasLtHandle_t ltHandle, cublasOperation_t transA, cublasOperation_t transB, + int m, int n, int k, const float *alpha, const float *A, int lda, const float *B, + int ldb, const float *beta, float *C, int ldc, const float* bias, + float* workspace, size_t workspaceSizeBytes,bool do_relu) { + + return cublasLtAffineHelper(ltHandle, transA, transB, CUDA_R_32F, m, n, k, alpha, A, lda, B, ldb, beta, C, ldc, bias, + workspace, workspaceSizeBytes, do_relu); +} + +template +void affineTyped(marian::Tensor C, Ptr allocator, const marian::Tensor& A, const marian::Tensor& B, const marian::Tensor& bias, + bool transA, bool transB, T beta, T scalar, bool do_relu) { + + CUDA_CHECK(cudaSetDevice((int)C->getDeviceId().no)); + T alpha = scalar; + + int m = A->shape().elements() / A->shape().back(); + int k = A->shape().back(); + if(transA) + std::swap(m, k); + + int l = B->shape().elements() / B->shape().back(); + int n = B->shape().back(); + if(transB) + std::swap(l, n); + + int lda = A->shape().back(); + int ldb = B->shape().back(); + int ldc = B->shape().back(); + + if(transB) + ldc = B->shape().elements() / B->shape().back(); + + cublasOperation_t opA = transA ? CUBLAS_OP_T : CUBLAS_OP_N; + cublasOperation_t opB = transB ? CUBLAS_OP_T : CUBLAS_OP_N; + + auto backend = std::static_pointer_cast(C->getBackend()); + auto ltHandle = (cublasLtHandle_t)backend->getCublasHandle(); // A cublas handle encapsulates an lt handle + + size_t numWorkSpaceElts = 8192; // Allows for cublasLt to perform split-K gemms. This is chosen to be at least + // 16 KiB for float16 which is large enough to prevent alloc failed errors + size_t workspaceSizeBytes = numWorkSpaceElts * sizeof(T); + IPtr workspace = allocator->alloc(numWorkSpaceElts); //TODO fix + + CUBLAS_CHECK(cublasLtAffineTyped(ltHandle, + opB, + opA, + n, + m, + k, + &alpha, + B->data(), + ldb, + A->data(), + lda, + &beta, + C->data(), + ldc, + bias->data(), + workspace->data(), + workspaceSizeBytes, + do_relu)); + + allocator->free(workspace); // TODO fix without synchronize (bad for small batch) +} + +// Earlier versions of cublasLT do not support bias addition for fp32 and fp16. + +// This version is needed so that Windows doesn't complain when compiling CUDA < 11. Otherwise, the ifdef could be inside of one +// definition of Affine. void Affine(marian::Tensor C, Ptr allocator, const marian::Tensor& A, const marian::Tensor& B, const marian::Tensor& bias, bool transA, bool transB, float beta, float scalar, bool do_relu) { + if(C->type() == Type::float32) { + affineTyped(C, allocator, A, B, bias, transA, transB, beta, scalar, do_relu); +#if COMPILE_FP16 + } else if(C->type() == Type::float16) { + affineTyped(C, allocator, A, B, bias, transA, transB, __float2half(beta), __float2half(scalar), do_relu); +#endif + } else { + ABORT("Affine not implemented for type {}", C->type()); + } +} + +#else + +void Affine(marian::Tensor C, Ptr /*allocator*/, const marian::Tensor& A, const marian::Tensor& B, const marian::Tensor& bias, + bool transA, bool transB, float beta, float scalar, bool do_relu) { + if(C->type() == Type::float32) { ProdTyped(C, A, B, transA, transB, beta, scalar); #if COMPILE_FP16 @@ -555,6 +695,7 @@ void Affine(marian::Tensor C, Ptr allocator, const marian::Tensor& A, } BiasAdd(C, bias, do_relu); } +#endif } // namespace gpu } // namespace marian From 4b41b86e427e1314b3e407e7fe6ca279d823df71 Mon Sep 17 00:00:00 2001 From: Rawn Henry Date: Wed, 9 Dec 2020 12:06:06 -0800 Subject: [PATCH 05/19] Works around some bugs in CUDA 11 --- src/tensors/gpu/prod.cpp | 74 +++++++++++++++++++++++++++++++++------- 1 file changed, 62 insertions(+), 12 deletions(-) diff --git a/src/tensors/gpu/prod.cpp b/src/tensors/gpu/prod.cpp index dafabb71c..b1c540151 100755 --- a/src/tensors/gpu/prod.cpp +++ b/src/tensors/gpu/prod.cpp @@ -20,6 +20,26 @@ namespace marian { namespace gpu { +// It seems that the bias must be 8 byte aligned for the cublasLt epilogue to work. Therefore, +// if the bias pointer is not 8 byte aligned, we do a normal matmul in cublasLt and invoke a +// custom epilogue kernel. +static constexpr int REQUIRED_BIAS_ALIGNMENT = 8; + +// Used to set preferences for cublasLt to filter out algos if matrices to not meet default 256 byte alignment +int getAlignmentUpTo256(const void *ptr) { + uintptr_t addr = (uintptr_t)ptr; + int trailingZeros = 0; + + for(int shiftAmt = 8, mask = 0xFF; shiftAmt > 0; shiftAmt /= 2, mask >>=shiftAmt) { + if ((addr & mask) == 0) { + trailingZeros += shiftAmt; + addr >>= shiftAmt; + } + } + + return std::min(256, 1 << trailingZeros); +} + // The explicit version of matmult like cublasGemmEx choose their math mode based on the algorithm that // has been passed into the function call and seem to ignore setMathMode. Here we query the used math mode // to choose the algorithm. @@ -543,12 +563,13 @@ void CSRProd(marian::Tensor C, if(St_offsets) allocator->free(St_offsets); } -#if CUDA_VERSION >= 11000 +#if CUDA_VERSION >= 11000 // Earlier versions of cublasLT do not support bias addition for fp32 and fp16. + static cublasStatus_t cublasLtAffineHelper(cublasLtHandle_t ltHandle, cublasOperation_t transA, cublasOperation_t transB, cudaDataType matrixType, int m, int n, int k, const void *alpha, const void *A, int lda, const void *B, int ldb, const void *beta, void *C, int ldc, const void* bias, - void* workspace, size_t workspaceSize, bool do_relu) { + void* workspace, size_t workspaceSize, bool do_relu, cudaStream_t stream) { cublasLtMatmulDesc_t operationDesc = NULL; cublasLtMatrixLayout_t Adesc = NULL, Bdesc = NULL, Cdesc = NULL; @@ -560,6 +581,11 @@ static cublasStatus_t cublasLtAffineHelper(cublasLtHandle_t ltHandle, cublasOper cublasLtEpilogue_t epilogue = do_relu? CUBLASLT_EPILOGUE_RELU_BIAS: CUBLASLT_EPILOGUE_BIAS; cublasComputeType_t computeType = matrixType == CUDA_R_32F? CUBLAS_COMPUTE_32F_FAST_16F: CUBLAS_COMPUTE_16F; + // If the bias is not aligned, just matmul and invoke custom epilogue later. + // cublas fails with a misalignment error if this condition is not true. + if((uintptr_t)bias % REQUIRED_BIAS_ALIGNMENT != 0) { + epilogue = CUBLASLT_EPILOGUE_DEFAULT; + } CUBLAS_CHECK(cublasLtMatmulDescCreate(&operationDesc, computeType, matrixType)); CUBLAS_CHECK(cublasLtMatmulDescSetAttribute(operationDesc, CUBLASLT_MATMUL_DESC_TRANSA, &transA, sizeof(transA))); @@ -571,12 +597,23 @@ static cublasStatus_t cublasLtAffineHelper(cublasLtHandle_t ltHandle, cublasOper CUBLAS_CHECK(cublasLtMatrixLayoutCreate(&Bdesc, matrixType, transB == CUBLAS_OP_N ? k : n, transB == CUBLAS_OP_N ? n : k, ldb)); CUBLAS_CHECK(cublasLtMatrixLayoutCreate(&Cdesc, matrixType, m, n, ldc)); + // I think we need to do this since we can slice matrices... + // The allocator always allocates on 256 byte boundaries but we have no guarantees about the alignment of a matrix slice so we filter out + // algorithms that would not work with matrices not aligned to 256 bytes. + int alignmentA = getAlignmentUpTo256(A); + int alignmentB = getAlignmentUpTo256(B); + int alignmentC = getAlignmentUpTo256(C); + CUBLAS_CHECK(cublasLtMatmulPreferenceCreate(&preference)); CUBLAS_CHECK(cublasLtMatmulPreferenceSetAttribute(preference, CUBLASLT_MATMUL_PREF_MAX_WORKSPACE_BYTES, &workspaceSize, sizeof(workspaceSize))); + CUBLAS_CHECK(cublasLtMatmulPreferenceSetAttribute(preference, CUBLASLT_MATMUL_PREF_MIN_ALIGNMENT_A_BYTES, &alignmentA, sizeof(alignmentA))); + CUBLAS_CHECK(cublasLtMatmulPreferenceSetAttribute(preference, CUBLASLT_MATMUL_PREF_MIN_ALIGNMENT_B_BYTES, &alignmentB, sizeof(alignmentB))); + CUBLAS_CHECK(cublasLtMatmulPreferenceSetAttribute(preference, CUBLASLT_MATMUL_PREF_MIN_ALIGNMENT_C_BYTES, &alignmentC, sizeof(alignmentC))); + CUBLAS_CHECK(cublasLtMatmulPreferenceSetAttribute(preference, CUBLASLT_MATMUL_PREF_MIN_ALIGNMENT_D_BYTES, &alignmentC, sizeof(alignmentC))); CUBLAS_CHECK(cublasLtMatmulAlgoGetHeuristic(ltHandle, operationDesc, Adesc, Bdesc, Cdesc, Cdesc, preference, 1, &heuristicResult, &returnedResults)); - cublasStatus_t opStatus = cublasLtMatmul(ltHandle, operationDesc, alpha, A, Adesc, B, Bdesc, beta, C, Cdesc, C, Cdesc, &heuristicResult.algo, - workspace, workspaceSize, cudaStreamPerThread); + cublasStatus_t opStatus = cublasLtMatmul(ltHandle, operationDesc, alpha, A, Adesc, B, Bdesc, beta, C, Cdesc, C, Cdesc, + &heuristicResult.algo, workspace, workspaceSize, stream); if (preference) CUBLAS_CHECK(cublasLtMatmulPreferenceDestroy(preference)); if (Cdesc) CUBLAS_CHECK(cublasLtMatrixLayoutDestroy(Cdesc)); @@ -590,18 +627,18 @@ static cublasStatus_t cublasLtAffineHelper(cublasLtHandle_t ltHandle, cublasOper static cublasStatus_t cublasLtAffineTyped(cublasLtHandle_t ltHandle, cublasOperation_t transA, cublasOperation_t transB, int m, int n, int k, const half *alpha, const half *A, int lda, const half *B, int ldb, const half *beta, half *C, int ldc, const half* bias, - half* workspace, size_t workspaceSizeBytes, bool do_relu) { + half* workspace, size_t workspaceSizeBytes, bool do_relu, cudaStream_t stream) { return cublasLtAffineHelper(ltHandle, transA, transB, CUDA_R_16F, m, n, k, alpha, A, lda, B, ldb, beta, C, ldc, bias, - workspace, workspaceSizeBytes, do_relu); + workspace, workspaceSizeBytes, do_relu, stream); } static cublasStatus_t cublasLtAffineTyped(cublasLtHandle_t ltHandle, cublasOperation_t transA, cublasOperation_t transB, int m, int n, int k, const float *alpha, const float *A, int lda, const float *B, int ldb, const float *beta, float *C, int ldc, const float* bias, - float* workspace, size_t workspaceSizeBytes,bool do_relu) { + float* workspace, size_t workspaceSizeBytes,bool do_relu, cudaStream_t stream) { return cublasLtAffineHelper(ltHandle, transA, transB, CUDA_R_32F, m, n, k, alpha, A, lda, B, ldb, beta, C, ldc, bias, - workspace, workspaceSizeBytes, do_relu); + workspace, workspaceSizeBytes, do_relu, stream); } template @@ -625,6 +662,9 @@ void affineTyped(marian::Tensor C, Ptr allocator, const marian::Tenso int ldb = B->shape().back(); int ldc = B->shape().back(); + size_t bias_size = bias->shape().elements(); + ABORT_IF(n != bias_size, "The number of elements in the bias must match the number of columns in C"); + if(transB) ldc = B->shape().elements() / B->shape().back(); @@ -632,12 +672,17 @@ void affineTyped(marian::Tensor C, Ptr allocator, const marian::Tenso cublasOperation_t opB = transB ? CUBLAS_OP_T : CUBLAS_OP_N; auto backend = std::static_pointer_cast(C->getBackend()); + auto cublasHandle = backend->getCublasHandle(); auto ltHandle = (cublasLtHandle_t)backend->getCublasHandle(); // A cublas handle encapsulates an lt handle size_t numWorkSpaceElts = 8192; // Allows for cublasLt to perform split-K gemms. This is chosen to be at least // 16 KiB for float16 which is large enough to prevent alloc failed errors size_t workspaceSizeBytes = numWorkSpaceElts * sizeof(T); - IPtr workspace = allocator->alloc(numWorkSpaceElts); //TODO fix + IPtr workspace = allocator->alloc(numWorkSpaceElts); + + cudaStream_t stream = 0; + CUBLAS_CHECK(cublasGetStream(cublasHandle, &stream)); + CUBLAS_CHECK(cublasLtAffineTyped(ltHandle, opB, @@ -661,18 +706,23 @@ void affineTyped(marian::Tensor C, Ptr allocator, const marian::Tenso allocator->free(workspace); // TODO fix without synchronize (bad for small batch) } -// Earlier versions of cublasLT do not support bias addition for fp32 and fp16. - // This version is needed so that Windows doesn't complain when compiling CUDA < 11. Otherwise, the ifdef could be inside of one // definition of Affine. void Affine(marian::Tensor C, Ptr allocator, const marian::Tensor& A, const marian::Tensor& B, const marian::Tensor& bias, bool transA, bool transB, float beta, float scalar, bool do_relu) { - + // There is a bug in CUDA 11 where the bias pointer needs to be 8 byte aligned. This bug will be fix in a subsequent release. For now, + // we launch a custom epilogue if the bias does not meet the alignment requirement. if(C->type() == Type::float32) { affineTyped(C, allocator, A, B, bias, transA, transB, beta, scalar, do_relu); + if((uintptr_t)bias->data() % REQUIRED_BIAS_ALIGNMENT != 0) { + BiasAdd(C, bias, do_relu); + } #if COMPILE_FP16 } else if(C->type() == Type::float16) { affineTyped(C, allocator, A, B, bias, transA, transB, __float2half(beta), __float2half(scalar), do_relu); + if((uintptr_t)bias->data() % REQUIRED_BIAS_ALIGNMENT != 0) { + BiasAdd(C, bias, do_relu); + } #endif } else { ABORT("Affine not implemented for type {}", C->type()); From cc30a32d4a9282dcf240fe753ad9402687d6e24e Mon Sep 17 00:00:00 2001 From: Rawn Henry Date: Mon, 14 Dec 2020 22:42:52 -0800 Subject: [PATCH 06/19] Fixes API and removes comment --- src/tensors/gpu/prod.cpp | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/src/tensors/gpu/prod.cpp b/src/tensors/gpu/prod.cpp index b1c540151..6da2f7b30 100755 --- a/src/tensors/gpu/prod.cpp +++ b/src/tensors/gpu/prod.cpp @@ -701,9 +701,10 @@ void affineTyped(marian::Tensor C, Ptr allocator, const marian::Tenso bias->data(), workspace->data(), workspaceSizeBytes, - do_relu)); + do_relu, + stream)); - allocator->free(workspace); // TODO fix without synchronize (bad for small batch) + allocator->free(workspace); } // This version is needed so that Windows doesn't complain when compiling CUDA < 11. Otherwise, the ifdef could be inside of one From 792669f110438edcefd5c092e72d78b7c13860e7 Mon Sep 17 00:00:00 2001 From: Rawn Henry Date: Tue, 15 Dec 2020 10:25:51 -0800 Subject: [PATCH 07/19] Removes SPDX identifier in operator tests --- src/tests/units/operator_tests.cpp | 4 ---- 1 file changed, 4 deletions(-) diff --git a/src/tests/units/operator_tests.cpp b/src/tests/units/operator_tests.cpp index 43ec72e58..bb6b97058 100644 --- a/src/tests/units/operator_tests.cpp +++ b/src/tests/units/operator_tests.cpp @@ -1,7 +1,3 @@ -/* All or part of this file was contributed by NVIDIA under license: - * Copyright (C) 2020 NVIDIA Corporation - * SPDX-License-Identifier: MIT - */ #include "catch.hpp" #include "graph/expression_graph.h" #include "graph/expression_operators.h" From 0c0b8ac7d5a4d8a5b436200373bc669d6ccd1d88 Mon Sep 17 00:00:00 2001 From: Rawn Henry Date: Tue, 15 Dec 2020 10:36:19 -0800 Subject: [PATCH 08/19] Adds SPDX identifier in operator tests - will remove in a different PR --- src/tests/units/operator_tests.cpp | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/src/tests/units/operator_tests.cpp b/src/tests/units/operator_tests.cpp index bb6b97058..43ec72e58 100644 --- a/src/tests/units/operator_tests.cpp +++ b/src/tests/units/operator_tests.cpp @@ -1,3 +1,7 @@ +/* All or part of this file was contributed by NVIDIA under license: + * Copyright (C) 2020 NVIDIA Corporation + * SPDX-License-Identifier: MIT + */ #include "catch.hpp" #include "graph/expression_graph.h" #include "graph/expression_operators.h" From 6258fd58c3354ee045d2c7ab286b5ddbf8b82168 Mon Sep 17 00:00:00 2001 From: Rawn Henry Date: Tue, 15 Dec 2020 10:50:57 -0800 Subject: [PATCH 09/19] Removes NVIDIA notices --- CMakeLists.txt | 4 ---- 1 file changed, 4 deletions(-) diff --git a/CMakeLists.txt b/CMakeLists.txt index f6ee826d8..bfb26ad86 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -1,7 +1,3 @@ -# All or part of this file was contributed by NVIDIA under license: -# Copyright (C) 2020 NVIDIA Corporation -# SPDX-License-Identifier: MIT - cmake_minimum_required(VERSION 3.5.1) set(CMAKE_MODULE_PATH ${CMAKE_CURRENT_SOURCE_DIR}/cmake) From 764d7b85a15f10094444811c9609478022344926 Mon Sep 17 00:00:00 2001 From: Rawn Henry Date: Tue, 15 Dec 2020 11:05:06 -0800 Subject: [PATCH 10/19] Fixes windows compile errors --- src/tensors/cpu/prod.cpp | 25 ++++++++++++------------- 1 file changed, 12 insertions(+), 13 deletions(-) diff --git a/src/tensors/cpu/prod.cpp b/src/tensors/cpu/prod.cpp index 6112db7ee..2dee32d2c 100755 --- a/src/tensors/cpu/prod.cpp +++ b/src/tensors/cpu/prod.cpp @@ -14,19 +14,18 @@ namespace marian { namespace cpu { -void Affine(marian::Tensor C, - Ptr allocator, - const marian::Tensor& A, - const marian::Tensor& B, - const marian::Tensor& bias, - bool transA, - bool transB, - float beta, - float scalar, - bool relu_postprocess) { - - ABORT("Not supported on CPU"); - } +void Affine(marian::Tensor /*C*/, + Ptr /*allocator*/, + const marian::Tensor& /*A*/, + const marian::Tensor& /*B*/, + const marian::Tensor& /*bias*/, + bool /*transA*/, + bool /*transB*/, + float /*beta*/, + float /*scalar*/, + bool /*relu_postprocess*/) { + ABORT("Not supported on CPU"); +} void Prod(marian::Tensor C, const marian::Tensor& A, From 95fb4277cc1243dee5825a01d777323e78e9a915 Mon Sep 17 00:00:00 2001 From: Rawn Henry Date: Tue, 15 Dec 2020 11:10:29 -0800 Subject: [PATCH 11/19] Format changes --- src/graph/expression_operators.cpp | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/graph/expression_operators.cpp b/src/graph/expression_operators.cpp index d5d6c031a..7e64d38af 100644 --- a/src/graph/expression_operators.cpp +++ b/src/graph/expression_operators.cpp @@ -535,15 +535,15 @@ static Expr affineDefault(Expr a, Expr b, Expr bias, bool transA, bool transB, f std::vector nodes = { clip(a, clipValue), clip(b, clipValue), bias }; // If we are using CPU, we broadcast the ones vector. On GPU, the bias addition can be fused into the GEMM with CUDA >= 11 - if(g->getBackend()->getDeviceId().type == DeviceType::cpu || !g->isInference()) { + if (g->getBackend()->getDeviceId().type == DeviceType::cpu || !g->isInference()) { int rows = a->shape().elements() / a->shape()[-1]; Expr ones = g->ones({ rows, 1 }); nodes.push_back(ones); } - if(do_relu) { + if (do_relu) { // For GPU inference, we can fuse the RELU into the bias addition. - if(g->isInference() && g->getBackend()->getDeviceId().type == DeviceType::gpu) { + if (g->isInference() && g->getBackend()->getDeviceId().type == DeviceType::gpu) { return Expression(nodes, transA, transB, scale, do_relu); } Expr affineOp = Expression(nodes, transA, transB, scale, false); From 4bbf17f480c859575ba0d702510757a856a45967 Mon Sep 17 00:00:00 2001 From: Rawn Henry Date: Tue, 15 Dec 2020 11:30:43 -0800 Subject: [PATCH 12/19] updates changelog --- CHANGELOG.md | 2 ++ 1 file changed, 2 insertions(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index 4182d72b6..bb348f835 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -9,6 +9,8 @@ and this project adheres to [Semantic Versioning](http://semver.org/spec/v2.0.0. ## [Unreleased] ### Added +- Adds custom bias epilogue kernel. +- Adds support for fusing relu and bias addition into gemms when using cuda 11. - Add --train-embedder-rank for fine-tuning any encoder(-decoder) model for multi-lingual similarity via softmax-margin loss - Add --logical-epoch that allows to redefine the displayed epoch counter as a multiple of n data epochs, updates or labels. Also allows to define width of fractional part with second argument. - Add --metrics chrf for computing ChrF according to https://www.aclweb.org/anthology/W15-3049/ and SacreBLEU reference implementation From 256e3978ba67416aaab536656072fab6a0d02af2 Mon Sep 17 00:00:00 2001 From: Rawn Henry Date: Tue, 20 Oct 2020 09:42:38 -0700 Subject: [PATCH 13/19] Adds bigobj flag to windows builds --- CMakeLists.txt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/CMakeLists.txt b/CMakeLists.txt index bfb26ad86..21aa8e062 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -93,7 +93,7 @@ if(MSVC) set(INTRINSICS "/arch:AVX2") # set(INTRINSICS "/arch:AVX512") - set(CMAKE_CXX_FLAGS "/EHsc /DWIN32 /D_WINDOWS /DUNICODE /D_UNICODE /D_CRT_NONSTDC_NO_WARNINGS /D_CRT_SECURE_NO_WARNINGS ${DISABLE_GLOBALLY}") + set(CMAKE_CXX_FLAGS "/EHsc /DWIN32 /D_WINDOWS /DUNICODE /D_UNICODE /bigobj /D_CRT_NONSTDC_NO_WARNINGS /D_CRT_SECURE_NO_WARNINGS ${DISABLE_GLOBALLY}") set(CMAKE_CXX_FLAGS_RELEASE "${CMAKE_CXX_FLAGS} /MT /O2 ${INTRINSICS} /Zi /MP /GL /DNDEBUG") set(CMAKE_CXX_FLAGS_DEBUG "${CMAKE_CXX_FLAGS} /MTd /Od /Ob0 ${INTRINSICS} /RTC1 /Zi /D_DEBUG") From a48489343ae7c63db83f81b76c8ab62ee337428b Mon Sep 17 00:00:00 2001 From: Marcin Junczys-Dowmunt Date: Thu, 25 Mar 2021 05:03:53 +0000 Subject: [PATCH 14/19] refactor --- src/graph/expression_operators.cpp | 62 +++---- src/graph/expression_operators.h | 13 +- src/graph/node_operators_binary.h | 273 +++++++++++++++++++++++++++++ src/layers/generic.h | 38 ++-- src/layers/output.cpp | 2 +- src/models/transformer.h | 33 +--- src/tensors/cpu/prod.cpp | 30 ++-- src/tensors/gpu/prod.cpp | 7 +- 8 files changed, 358 insertions(+), 100 deletions(-) diff --git a/src/graph/expression_operators.cpp b/src/graph/expression_operators.cpp index 84902579d..7f01454e3 100644 --- a/src/graph/expression_operators.cpp +++ b/src/graph/expression_operators.cpp @@ -519,29 +519,13 @@ Expr bdot(Expr a, Expr b, bool transA, bool transB, float scale) { return Expression(a, b, transA, transB, scale); } -static Expr affineDefault(Expr a, Expr b, Expr bias, bool transA, bool transB, float scale, bool do_relu=false) { +Expr affineDefault(Expr a, Expr b, Expr bias, bool transA, bool transB, float scale) { // general version, MKL, CBlas or CUDA - auto g = a->graph(); - std::vector nodes = { a, b, bias }; - - // If we are using CPU, we broadcast the ones vector. On GPU, the bias addition can be fused into the GEMM with CUDA >= 11 - if (g->getBackend()->getDeviceId().type == DeviceType::cpu || !g->isInference()) { - int rows = a->shape().elements() / a->shape()[-1]; - Expr ones = g->ones({ rows, 1 }); - nodes.push_back(ones); - } - - if (do_relu) { - // For GPU inference, we can fuse the RELU into the bias addition. - if (g->isInference() && g->getBackend()->getDeviceId().type == DeviceType::gpu) { - return Expression(nodes, transA, transB, scale, do_relu); - } - Expr affineOp = Expression(nodes, transA, transB, scale, false); - return relu(affineOp); - } - - return Expression(nodes, transA, transB, scale, do_relu); + int rows = a->shape().elements() / a->shape()[-1]; + Expr ones = a->graph()->ones({ rows, 1 }); + std::vector nodes = { a, b, bias, ones }; + return Expression(nodes, transA, transB, scale); } // This operation used to implement auto-tuning. We have removed it for now due to complexity, but plan to revisit it in the future. @@ -549,7 +533,7 @@ static Expr affineDefault(Expr a, Expr b, Expr bias, bool transA, bool transB, f // youki/packed-model-pr-backup1031 // https://machinetranslation.visualstudio.com/Marian/_git/marian-dev?version=GByouki%2Fpacked-model-pr-backup1031 // SHA: 3456a7ed1d1608cfad74cd2c414e7e8fe141aa52 -Expr affine(Expr a, Expr b, Expr bias, bool transA, bool transB, float scale, bool do_relu) { +Expr affine(Expr a, Expr b, Expr bias, bool transA, bool transB, float scale) { auto device = a->graph()->getDeviceId().type; Type aElementType = a->value_type(); @@ -557,12 +541,9 @@ Expr affine(Expr a, Expr b, Expr bias, bool transA, bool transB, float scale, bo if(device == DeviceType::cpu) { if(isFloat(aElementType) && isFloat(bElementType)) { - return affineDefault(a, b, bias, transA, transB, scale, do_relu); + return affineDefault(a, b, bias, transA, transB, scale); } else if(isFloat(aElementType) && isIntgemm(bElementType)) { - Expr affineTransform = cpu::integer::affineOrDot(a, b, bias, transA, transB, scale); - if(do_relu) - affineTransform = relu(affineTransform); - return affineTransform; + return cpu::integer::affineOrDot(a, b, bias, transA, transB, scale); } else if(isFloat(aElementType) && isPacked(bElementType)) { #if USE_FBGEMM // 07/10/2019 - Use packed GEMM only if the cpu architecture supports AVX2 @@ -572,16 +553,13 @@ Expr affine(Expr a, Expr b, Expr bias, bool transA, bool transB, float scale, bo // and this cpu lookup is executed only once and the state is kept in FBGEMM. if(fbgemm::fbgemmHasAvx2Support()) { // This variant of affine product can handle matrix multiplications with packed8 and packed16 weight matrix (B). - Expr affineTransform = cpu::variant::affine(a, - b, - b->shape(), - bias, - transA, - transB, - scale); - if(do_relu) - affineTransform = relu(affineTransform); - return affineTransform; + return cpu::variant::affine(a, + b, + b->shape(), + bias, + transA, + transB, + scale); } else { ABORT("AVX2 is not available. At least, AVX2 is needed to use fbgemm-based packed GEMM"); } @@ -596,10 +574,18 @@ Expr affine(Expr a, Expr b, Expr bias, bool transA, bool transB, float scale, bo ABORT_IF(!isFloat(aElementType) || !isFloat(bElementType), "GPU-based GEMM only supports float types, you have A: {} and B: {}", aElementType, bElementType); - return affineDefault(a, b, bias, transA, transB, scale, do_relu); + return affineDefault(a, b, bias, transA, transB, scale); } } +Expr affineWithRelu(Expr a, Expr b, Expr bias, bool transA, bool transB, float scale) { + auto graph = a->graph(); + if(graph->isInference()) + return Expression(a, b, bias, transA, transB, scale); + else + return relu(affine(a, b, bias, transA, transB, scale)); +} + // @TODO: Not a great place to check this #if CUDA_VERSION < 11000 // multiply a CSR matrix A with a matrix B diff --git a/src/graph/expression_operators.h b/src/graph/expression_operators.h index c9f018dce..d90b063fb 100644 --- a/src/graph/expression_operators.h +++ b/src/graph/expression_operators.h @@ -491,8 +491,17 @@ Expr affine(Expr a, Expr c, bool transA = false, bool transB = false, - float scalar = 1.f, - bool do_relu = false); + float scalar = 1.f); + +/** + * As above, but efficiently applies relu transformation to output. For inference only. + */ +Expr affineWithRelu(Expr a, + Expr b, + Expr c, + bool transA = false, + bool transB = false, + float scalar = 1.f); /** * Computes the dot product of CSR-tensor @p A with @p B. diff --git a/src/graph/node_operators_binary.h b/src/graph/node_operators_binary.h index ff3e05f8f..14efa8f69 100644 --- a/src/graph/node_operators_binary.h +++ b/src/graph/node_operators_binary.h @@ -227,6 +227,7 @@ class DotNodeOp : public NaryNodeOp { const std::string color() override { return "orange"; } }; +#if 0 class AffineNodeOp : public NaryNodeOp { private: friend class SerializationHelpers; @@ -447,6 +448,278 @@ class AffineNodeOp : public NaryNodeOp { } }; +#endif + +class AffineNodeOp : public NaryNodeOp { +private: + friend class SerializationHelpers; + bool transA_; + bool transB_; + float scalar_; + +public: + AffineNodeOp(const std::vector& nodes, + bool transA, + bool transB, + float scalar) + : NaryNodeOp(nodes, newShape(nodes[0], nodes[1], transA, transB)), + transA_(transA), + transB_(transB), + scalar_(scalar) {} + + Shape newShape(Expr a, Expr b, bool transA, bool transB) { + auto shapeA = a->shape(); + if(transA) { + shapeA.set(shapeA.size() - 2, a->shape()[shapeA.size() - 1]); + shapeA.set(shapeA.size() - 1, a->shape()[shapeA.size() - 2]); + } + + auto shapeB = b->shape(); + if(transB) { + shapeB.set(shapeB.size() - 2, b->shape()[shapeB.size() - 1]); + shapeB.set(shapeB.size() - 1, b->shape()[shapeB.size() - 2]); + } + + Shape outShape = shapeA; + outShape.set(outShape.size() - 1, shapeB[shapeB.size() - 1]); + ABORT_IF(shapeA[shapeA.size() - 1] != shapeB[shapeB.size() - 2], + "Matrix product requires inner dimensions to match in {}{} * {}{}", std::string(shapeA), transA, std::string(shapeB), transB); + return outShape; + } + + NodeOps forwardOps() override { + using namespace functional; + + return { + NodeOp(Affine(val_, + graph()->allocator(), + child(0)->val(), + child(1)->val(), + child(2)->val(), + transA_, + transB_, + 0.f, + scalar_, + /*doRelu=*/false)) + }; + } + + NodeOps backwardOps() override { + // D is the adjoint, the matrix of derivatives + // df/dA += alpha * dot(D, op(B).T) + // df/dB += alpha * dot(op(A).T, D) + // beta set to 1.0 in gemm, C = alpha * dot(op(A), op(B)) + beta * C + // to sum gradients from different graph parts + + auto isParameter = [](Expr p) { + return std::dynamic_pointer_cast(p) != nullptr; + }; + + // if child A is not a parameter (i.e. activations) use computeType float32 for accumulation + Type computeTypeA = child(0)->trainable() ? child(0)->grad()->type() : Type::float32; + if(!isParameter(child(0)) && computeTypeA == Type::float16) + computeTypeA = Type::float32; + + // if child B is not a parameter (i.e. activations) use computeType float32 for accumulation + Type computeTypeB = child(1)->trainable() ? child(1)->grad()->type() : Type::float32; + if(!isParameter(child(1)) && computeTypeB == Type::float16) + computeTypeB = Type::float32; + + // if child C (bias) is not a parameter (i.e. activations) use computeType float32 for accumulation + Type computeTypeC = child(2)->trainable() ? child(2)->grad()->type() : Type::float32; + if(!isParameter(child(2)) && computeTypeC == Type::float16) + computeTypeC = Type::float32; + + // We reduce bias gradients with a matrix multiply + if(!transA_ && transB_) + return { + NodeOp(Prod(child(0)->grad(), + adj_, + child(1)->val(), + false, + false, + 1.0, + scalar_, computeTypeA)), + NodeOp(Prod(child(1)->grad(), + adj_, + child(0)->val(), + true, + false, + 1.0, + scalar_, computeTypeB)), + NodeOp(Prod(child(2)->grad(), child(3)->val(), adj_, true, false, 0.f, 1.f, computeTypeC)) + }; + + if(transA_ && !transB_) + return { + NodeOp(Prod(child(0)->grad(), + child(1)->val(), + adj_, + false, + true, + 1.0, + scalar_, computeTypeA)), + NodeOp(Prod(child(1)->grad(), + child(0)->val(), + adj_, + false, + false, + 1.0, + scalar_, computeTypeB)), + NodeOp(Prod(child(2)->grad(), child(3)->val(), adj_, true, false, 0.f, 1.f, computeTypeC)) + }; + + if(transA_ && transB_) + return { + NodeOp(Prod(child(0)->grad(), + child(1)->val(), + adj_, + true, + true, + 1.0, + scalar_, computeTypeA)), + NodeOp(Prod(child(1)->grad(), + adj_, + child(0)->val(), + true, + true, + 1.0, + scalar_, computeTypeB)), + NodeOp(Prod(child(2)->grad(), child(3)->val(), adj_, true, false, 0.f, 1.f, computeTypeC)) + }; + + return { + NodeOp(Prod(child(0)->grad(), + adj_, + child(1)->val(), + false, + true, + 1.0, + scalar_, computeTypeA)), + NodeOp(Prod(child(1)->grad(), + child(0)->val(), + adj_, + true, + false, + 1.0, + scalar_, computeTypeB)), + NodeOp(Prod(child(2)->grad(), child(3)->val(), adj_, true, false, 0.f, 1.f, computeTypeC)) + }; + } + + const std::string type() override { return "affine"; } + + virtual size_t hash() override { + size_t seed = NaryNodeOp::hash(); + util::hash_combine(seed, transA_); + util::hash_combine(seed, transB_); + util::hash_combine(seed, scalar_); + return seed; + } + + virtual bool equal(Expr node) override { + if(!NaryNodeOp::equal(node)) + return false; + auto cnode = std::dynamic_pointer_cast(node); + if(!cnode) + return false; + if(transA_ != cnode->transA_) + return false; + if(transB_ != cnode->transB_) + return false; + if(scalar_ != cnode->scalar_) + return false; + return true; + } + +}; + +class AffineWithReluNodeOp : public NaryNodeOp { +private: + friend class SerializationHelpers; + bool transA_; + bool transB_; + float scalar_; + +public: + AffineWithReluNodeOp(Expr a, + Expr b, + Expr bias, + bool transA, + bool transB, + float scalar) + : NaryNodeOp({a, b, bias}, newShape(a, b, transA, transB)), + transA_(transA), + transB_(transB), + scalar_(scalar) {} + + Shape newShape(Expr a, Expr b, bool transA, bool transB) { + auto shapeA = a->shape(); + if(transA) { + shapeA.set(shapeA.size() - 2, a->shape()[shapeA.size() - 1]); + shapeA.set(shapeA.size() - 1, a->shape()[shapeA.size() - 2]); + } + + auto shapeB = b->shape(); + if(transB) { + shapeB.set(shapeB.size() - 2, b->shape()[shapeB.size() - 1]); + shapeB.set(shapeB.size() - 1, b->shape()[shapeB.size() - 2]); + } + + Shape outShape = shapeA; + outShape.set(outShape.size() - 1, shapeB[shapeB.size() - 1]); + ABORT_IF(shapeA[shapeA.size() - 1] != shapeB[shapeB.size() - 2], + "Matrix product requires inner dimensions to match in {}{} * {}{}", std::string(shapeA), transA, std::string(shapeB), transB); + return outShape; + } + + NodeOps forwardOps() override { + using namespace functional; + + return { + NodeOp(Affine(val_, + graph()->allocator(), + child(0)->val(), + child(1)->val(), + child(2)->val(), + transA_, + transB_, + 0.f, + scalar_, + /*doRelu=*/true)) + }; + } + + NodeOps backwardOps() override { + ABORT("AffineWithReluNodeOp cannot be used for training??"); + return {}; + } + + const std::string type() override { return "affineWithRelu"; } + + virtual size_t hash() override { + size_t seed = NaryNodeOp::hash(); + util::hash_combine(seed, transA_); + util::hash_combine(seed, transB_); + util::hash_combine(seed, scalar_); + return seed; + } + + virtual bool equal(Expr node) override { + if(!NaryNodeOp::equal(node)) + return false; + auto cnode = std::dynamic_pointer_cast(node); + if(!cnode) + return false; + if(transA_ != cnode->transA_) + return false; + if(transB_ != cnode->transB_) + return false; + if(scalar_ != cnode->scalar_) + return false; + return true; + } +}; class DotBatchedNodeOp : public NaryNodeOp { private: diff --git a/src/layers/generic.h b/src/layers/generic.h index fbba8d73e..5eb936151 100644 --- a/src/layers/generic.h +++ b/src/layers/generic.h @@ -170,41 +170,41 @@ class Dense : public LayerBase, public IUnaryLayer { // --- a few layers with built-in parameters created on the fly, without proper object // @TODO: change to a proper layer object +static inline std::function activationByName(const std::string& actName) { + if (actName == "relu") + return (ActivationFunction*)relu; + else if (actName == "swish") + return (ActivationFunction*)swish; + else if (actName == "gelu") + return (ActivationFunction*)gelu; + else if (actName == "") // return identity function if activation name is empty + return [](Expr x) { return x; }; + ABORT("Invalid activation name '{}'", actName); +} + // like affine() but with built-in parameters, activation, and dropout static inline Expr denseInline(Expr x, std::string prefix, std::string suffix, int outDim, Ptr initFn = inits::glorotUniform(), - const std::function& actFn = nullptr, + std::string actName = "", float dropProb = 0.0f) { auto graph = x->graph(); auto W = graph->param(prefix + "_W" + suffix, {x->shape()[-1], outDim}, inits::glorotUniform()); auto b = graph->param(prefix + "_b" + suffix, {1, outDim}, inits::zeros()); - x = affine(x, W, b); - if(actFn) - x = actFn(x); + if(actName == "relu") { + x = affineWithRelu(x, W, b); // speed optimization for inference, @TODO: handle better in future layer framework + } else { + x = affine(x, W, b); + x = activationByName(actName)(x); + } x = dropout(x, dropProb); // @TODO: check for infernce? return x; } -static inline Expr denseInlineRelu(Expr x, - std::string prefix, - std::string suffix, - int outDim, - Ptr initFn = inits::glorotUniform(), - float dropProb = 0.0f) { - auto graph = x->graph(); - - auto W = graph->param(prefix + "_W" + suffix, { x->shape()[-1], outDim }, initFn); - auto b = graph->param(prefix + "_b" + suffix, { 1, outDim }, inits::zeros()); - x = affine(x, W, b, false, false, 1.f, true); - x = dropout(x, dropProb); - return x; -} - static inline Expr layerNorm(Expr x, std::string prefix, std::string suffix = std::string()) { int dimModel = x->shape()[-1]; auto scale = x->graph()->param(prefix + "_ln_scale" + suffix, {1, dimModel}, inits::ones()); diff --git a/src/layers/output.cpp b/src/layers/output.cpp index 1d9c7b4b0..4c34bdcea 100644 --- a/src/layers/output.cpp +++ b/src/layers/output.cpp @@ -170,7 +170,7 @@ Logits Output::applyAsLogits(Expr input) /*override final*/ { /*suffix=*/"1", ffnDim, inits::glorotUniform(), - (ActivationFunction*)relu, + "relu", ffnDropProb); f = denseInline(f, name + "_ffn", /*suffix=*/"2", inputDim); // add & norm diff --git a/src/models/transformer.h b/src/models/transformer.h index 0eeb4743d..79b59000a 100644 --- a/src/models/transformer.h +++ b/src/models/transformer.h @@ -396,18 +396,6 @@ class Transformer : public EncoderOrDecoderBase { opt("transformer-heads"), /*cache=*/false); } - static inline - std::function activationByName(const std::string& actName) - { - if (actName == "relu") - return (ActivationFunction*)relu; - else if (actName == "swish") - return (ActivationFunction*)swish; - else if (actName == "gelu") - return (ActivationFunction*)gelu; - ABORT("Invalid activation name '{}'", actName); - } - Expr LayerFFN(std::string prefix, Expr input) const { int dimModel = input->shape()[-1]; @@ -418,7 +406,6 @@ class Transformer : public EncoderOrDecoderBase { auto actName = opt("transformer-ffn-activation"); int dimFfn = opt("transformer-dim-ffn"); int depthFfn = opt("transformer-ffn-depth"); - auto actFn = activationByName(actName); float ffnDropProb = inference_ ? 0 : opt("transformer-dropout-ffn"); @@ -427,18 +414,12 @@ class Transformer : public EncoderOrDecoderBase { auto initFn = inits::glorotUniform(true, true, depthScaling_ ? 1.f / sqrtf((float)depth_) : 1.f); // the stack of FF layers - for(int i = 1; i < depthFfn; ++i) { - if (actName == "relu") { - output = denseInlineRelu(output, prefix, std::to_string(i), dimFfn, initFn, ffnDropProb); - } else { - output = denseInline(output, prefix, /*suffix=*/std::to_string(i), dimFfn, initFn, actFn, ffnDropProb); - } - } + for(int i = 1; i < depthFfn; ++i) + output = denseInline(output, prefix, /*suffix=*/std::to_string(i), dimFfn, initFn, actName, ffnDropProb); output = denseInline(output, prefix, /*suffix=*/std::to_string(depthFfn), dimModel, initFn); auto opsPost = opt("transformer-postprocess"); - output - = postProcess(prefix + "_ffn", opsPost, output, input, dropProb); + output = postProcess(prefix + "_ffn", opsPost, output, input, dropProb); return output; } @@ -456,21 +437,21 @@ class Transformer : public EncoderOrDecoderBase { // FFN int dimAan = opt("transformer-dim-aan"); int depthAan = opt("transformer-aan-depth"); - auto actFn = activationByName(opt("transformer-aan-activation")); + auto actName = opt("transformer-aan-activation"); float aanDropProb = inference_ ? 0 : opt("transformer-dropout-ffn"); auto initFn = inits::glorotUniform(true, true, depthScaling_ ? 1.f / sqrtf((float)depth_) : 1.f); // the stack of AAN layers for(int i = 1; i < depthAan; ++i) - y = denseInline(y, prefix, /*suffix=*/std::to_string(i), dimAan, initFn, actFn, aanDropProb); + y = denseInline(y, prefix, /*suffix=*/std::to_string(i), dimAan, initFn, actName, aanDropProb); if(y->shape()[-1] != dimModel) // bring it back to the desired dimension if needed y = denseInline(y, prefix, std::to_string(depthAan), dimModel, initFn); bool noGate = opt("transformer-aan-nogate"); if(!noGate) { - auto gi = denseInline(x, prefix, /*suffix=*/"i", dimModel, initFn, (ActivationFunction*)sigmoid); - auto gf = denseInline(y, prefix, /*suffix=*/"f", dimModel, initFn, (ActivationFunction*)sigmoid); + auto gi = denseInline(x, prefix, /*suffix=*/"i", dimModel, initFn, "sigmoid"); + auto gf = denseInline(y, prefix, /*suffix=*/"f", dimModel, initFn, "sigmoid"); y = gi * x + gf * y; } diff --git a/src/tensors/cpu/prod.cpp b/src/tensors/cpu/prod.cpp index 77161989a..6e28bdd23 100755 --- a/src/tensors/cpu/prod.cpp +++ b/src/tensors/cpu/prod.cpp @@ -23,19 +23,6 @@ namespace marian { namespace cpu { -void Affine(marian::Tensor /*C*/, - Ptr /*allocator*/, - const marian::Tensor& /*A*/, - const marian::Tensor& /*B*/, - const marian::Tensor& /*bias*/, - bool /*transA*/, - bool /*transB*/, - float /*beta*/, - float /*scalar*/, - bool /*relu_postprocess*/) { - ABORT("Not supported on CPU"); -} - void Prod(marian::Tensor C, const marian::Tensor& A, const marian::Tensor& B, @@ -225,6 +212,23 @@ void ProdWithBias(marian::Tensor C, cpu::integer::AddBias(C, bias); } +void Affine(marian::Tensor C, + Ptr /*allocator*/, + const marian::Tensor& A, + const marian::Tensor& B, + const marian::Tensor& bias, + bool transA, + bool transB, + float beta, + float scalar, + bool reluPostprocess) { + using namespace functional; + ProdWithBias(C, A, B, bias, transA, transB, beta, scalar); + if(reluPostprocess) + cpu::Element(_1 = ReLU(_1), C); // @TODO: also fuse with AddBias +} + + void CSRProd(marian::Tensor C, Ptr /*allocator*/, const marian::Tensor& S_values, diff --git a/src/tensors/gpu/prod.cpp b/src/tensors/gpu/prod.cpp index 61a832c55..530b7a48d 100755 --- a/src/tensors/gpu/prod.cpp +++ b/src/tensors/gpu/prod.cpp @@ -605,7 +605,12 @@ void Affine(marian::Tensor C, Ptr allocator, const marian::Tensor& A, #else -void Affine(marian::Tensor C, Ptr /*allocator*/, const marian::Tensor& A, const marian::Tensor& B, const marian::Tensor& bias, +void Affine(marian::Tensor C, + Ptr /*allocator*/, + const marian::Tensor& A, + const marian::Tensor& B, + const marian::Tensor& bias, + const marian::Tensor& /*ones*/, bool transA, bool transB, float beta, float scalar, bool do_relu) { if(C->type() == Type::float32) { From 9de84d4ed6e8ea5704782e842db50d06d5d9ee45 Mon Sep 17 00:00:00 2001 From: Marcin Junczys-Dowmunt Date: Thu, 25 Mar 2021 05:06:15 +0000 Subject: [PATCH 15/19] remove previous code --- src/graph/node_operators_binary.h | 223 ------------------------------ 1 file changed, 223 deletions(-) diff --git a/src/graph/node_operators_binary.h b/src/graph/node_operators_binary.h index 14efa8f69..f0a263e7b 100644 --- a/src/graph/node_operators_binary.h +++ b/src/graph/node_operators_binary.h @@ -227,229 +227,6 @@ class DotNodeOp : public NaryNodeOp { const std::string color() override { return "orange"; } }; -#if 0 -class AffineNodeOp : public NaryNodeOp { -private: - friend class SerializationHelpers; - bool transA_; - bool transB_; - float scalar_; - bool do_relu_; - -public: - AffineNodeOp(const std::vector& nodes, - bool transA, - bool transB, - float scalar, - bool do_relu=false) - : NaryNodeOp(nodes, newShape(nodes[0], nodes[1], transA, transB)), - transA_(transA), - transB_(transB), - scalar_(scalar), - do_relu_(do_relu) { - // ReLU fusion checks - ABORT_IF(do_relu && graph()->getBackend()->getDeviceId().type == DeviceType::cpu, "ReLU not fused for CPU backend"); - ABORT_IF(do_relu && !graph()->isInference(), "ReLU fusion is only supported for GPU inference."); - - // bias fusion checks - ABORT_IF(nodes.size() == 3 && graph()->getBackend()->getDeviceId().type == DeviceType::cpu, "Bias addition not specialized for CPU backend"); - // No check for if inference since the forward computation will be the same. - } - - Shape newShape(Expr a, Expr b, bool transA, bool transB) { - auto shapeA = a->shape(); - if(transA) { - shapeA.set(shapeA.size() - 2, a->shape()[shapeA.size() - 1]); - shapeA.set(shapeA.size() - 1, a->shape()[shapeA.size() - 2]); - } - - auto shapeB = b->shape(); - if(transB) { - shapeB.set(shapeB.size() - 2, b->shape()[shapeB.size() - 1]); - shapeB.set(shapeB.size() - 1, b->shape()[shapeB.size() - 2]); - } - - Shape outShape = shapeA; - outShape.set(outShape.size() - 1, shapeB[shapeB.size() - 1]); - ABORT_IF(shapeA[shapeA.size() - 1] != shapeB[shapeB.size() - 2], - "Matrix product requires inner dimensions to match in {}{} * {}{}", std::string(shapeA), transA, std::string(shapeB), transB); - return outShape; - } - - NodeOps forwardOps() override { - using namespace functional; - - if(children_.size() == 3) { - ABORT_IF(graph()->getBackend()->getDeviceId().type != DeviceType::gpu, "Snuck passed constructor check? Only supported for GPU backend"); - // We must be using the GPU backend (Would have aborted otherwise) - return { - NodeOp( - Affine(val_, - graph()->allocator(), - child(0)->val(), - child(1)->val(), - child(2)->val(), - transA_, - transB_, - 0.f, - scalar_, - do_relu_)) - }; - } - - return { - NodeOp( - Prod(val_, - child(0)->val(), - child(1)->val(), - transA_, - transB_, - 0.f, - scalar_); - Prod(val_, child(3)->val(), child(2)->val(), false, false, 1.f, 1.f)) - }; - } - - NodeOps backwardOps() override { - // D is the adjoint, the matrix of derivatives - // df/dA += alpha * dot(D, op(B).T) - // df/dB += alpha * dot(op(A).T, D) - // beta set to 1.0 in gemm, C = alpha * dot(op(A), op(B)) + beta * C - // to sum gradients from different graph parts - using namespace functional; - ABORT_IF(do_relu_, "Snuck passed constructor check? ReLU fusion not supported for training."); - - auto isParameter = [](Expr p) { - return std::dynamic_pointer_cast(p) != nullptr; - }; - - // if child A is not a parameter (i.e. activations) use computeType float32 for accumulation - Type computeTypeA = child(0)->trainable() ? child(0)->grad()->type() : Type::float32; - if(!isParameter(child(0)) && computeTypeA == Type::float16) - computeTypeA = Type::float32; - - // if child B is not a parameter (i.e. activations) use computeType float32 for accumulation - Type computeTypeB = child(1)->trainable() ? child(1)->grad()->type() : Type::float32; - if(!isParameter(child(1)) && computeTypeB == Type::float16) - computeTypeB = Type::float32; - - // if child C (bias) is not a parameter (i.e. activations) use computeType float32 for accumulation - Type computeTypeC = child(2)->trainable() ? child(2)->grad()->type() : Type::float32; - if(!isParameter(child(2)) && computeTypeC == Type::float16) - computeTypeC = Type::float32; - - // We reduce bias gradients with a matrix multiply - if(!transA_ && transB_) - return { - NodeOp(Prod(child(0)->grad(), - adj_, - child(1)->val(), - false, - false, - 1.0, - scalar_, computeTypeA)), - NodeOp(Prod(child(1)->grad(), - adj_, - child(0)->val(), - true, - false, - 1.0, - scalar_, computeTypeB)), - NodeOp(Prod( - child(2)->grad(), child(3)->val(), adj_, true, false, 0.f, 1.f, computeTypeC)) - }; - - if(transA_ && !transB_) - return { - NodeOp(Prod(child(0)->grad(), - child(1)->val(), - adj_, - false, - true, - 1.0, - scalar_, computeTypeA)), - NodeOp(Prod(child(1)->grad(), - child(0)->val(), - adj_, - false, - false, - 1.0, - scalar_, computeTypeB)), - NodeOp(Prod( - child(2)->grad(), child(3)->val(), adj_, true, false, 0.f, 1.f, computeTypeC)) - }; - - if(transA_ && transB_) - return { - NodeOp(Prod(child(0)->grad(), - child(1)->val(), - adj_, - true, - true, - 1.0, - scalar_, computeTypeA)), - NodeOp(Prod(child(1)->grad(), - adj_, - child(0)->val(), - true, - true, - 1.0, - scalar_, computeTypeB)), - NodeOp(Prod( - child(2)->grad(), child(3)->val(), adj_, true, false, 0.f, 1.f, computeTypeC)) - }; - - return { - NodeOp(Prod(child(0)->grad(), - adj_, - child(1)->val(), - false, - true, - 1.0, - scalar_, computeTypeA)), - NodeOp(Prod(child(1)->grad(), - child(0)->val(), - adj_, - true, - false, - 1.0, - scalar_, computeTypeB)), - NodeOp(Prod( - child(2)->grad(), child(3)->val(), adj_, true, false, 0.f, 1.f, computeTypeC)) - }; - } - - const std::string type() override { return "affine"; } - - virtual size_t hash() override { - size_t seed = NaryNodeOp::hash(); - util::hash_combine(seed, transA_); - util::hash_combine(seed, transB_); - util::hash_combine(seed, scalar_); - util::hash_combine(seed, do_relu_); - return seed; - } - - virtual bool equal(Expr node) override { - if(!NaryNodeOp::equal(node)) - return false; - auto cnode = std::dynamic_pointer_cast(node); - if(!cnode) - return false; - if(transA_ != cnode->transA_) - return false; - if(transB_ != cnode->transB_) - return false; - if(scalar_ != cnode->scalar_) - return false; - if(do_relu_ != cnode->do_relu_) - return false; - return true; - } - -}; -#endif - class AffineNodeOp : public NaryNodeOp { private: friend class SerializationHelpers; From 43c54ce50c730d4805bb7c3625da810d50bc2aac Mon Sep 17 00:00:00 2001 From: Marcin Junczys-Dowmunt Date: Thu, 25 Mar 2021 05:33:34 +0000 Subject: [PATCH 16/19] add unit tests --- src/graph/expression_operators.h | 4 ++-- src/tests/units/operator_tests.cpp | 21 ++++++++++++++++----- 2 files changed, 18 insertions(+), 7 deletions(-) diff --git a/src/graph/expression_operators.h b/src/graph/expression_operators.h index d90b063fb..81b0f5ea2 100644 --- a/src/graph/expression_operators.h +++ b/src/graph/expression_operators.h @@ -488,7 +488,7 @@ Expr bdot(Expr a, */ Expr affine(Expr a, Expr b, - Expr c, + Expr bias, bool transA = false, bool transB = false, float scalar = 1.f); @@ -498,7 +498,7 @@ Expr affine(Expr a, */ Expr affineWithRelu(Expr a, Expr b, - Expr c, + Expr bias, bool transA = false, bool transB = false, float scalar = 1.f); diff --git a/src/tests/units/operator_tests.cpp b/src/tests/units/operator_tests.cpp index 27ccf1396..c3fd4a9e7 100644 --- a/src/tests/units/operator_tests.cpp +++ b/src/tests/units/operator_tests.cpp @@ -32,6 +32,8 @@ void tests(DeviceType device, Type floatType = Type::float32) { Config::seed = 1234; auto graph = New(); + + graph->setInference(true); graph->setDefaultElementType(floatType); graph->setDevice({0, device}); graph->reserveWorkspaceMB(16); @@ -539,15 +541,19 @@ void tests(DeviceType device, Type floatType = Type::float32) { values.clear(); std::vector vA({1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12}); - std::vector vB({1, 2, 3, 4, 5, 6}); - std::vector vAff({24, 30, 51, 66, 78, 102, 105, 138}); + std::vector vB({1, -2, 3, 4, -5, 6}); + std::vector vAff({-6, 26, -9, 50, -12, 74, -15, 98}); + std::vector vAffRelu({0, 26, 0, 50, 0, 74, 0, 98}); auto A = graph->param("A", {4, 3}, inits::fromVector(vA)); auto B = graph->param("B", {3, 2}, inits::fromVector(vB)); - auto C = graph->param("C", {4, 2}, inits::fromValue(2)); + auto bias = graph->param("C", {1, 2}, inits::fromValue(2)); + + auto aff1 = affine(A, B, bias); + auto aff2 = dot(A, B) + bias; - auto aff1 = affine(A, B, C); - auto aff2 = dot(A, B) + C; + auto affRelu1 = affineWithRelu(A, B, bias); + auto affRelu2 = relu(dot(A, B) + bias); graph->forward(); @@ -559,6 +565,11 @@ void tests(DeviceType device, Type floatType = Type::float32) { CHECK(aff2->shape() == aff1->shape()); aff2->val()->get(values2); CHECK(values2 == values); + + affRelu1->val()->get(values); + affRelu2->val()->get(values2); + CHECK(values2 == vAffRelu); + CHECK(values2 == values); } SECTION("repeat") { From 04b0f95f46f69c3882d43365a83511ea00e4ed2b Mon Sep 17 00:00:00 2001 From: Marcin Junczys-Dowmunt Date: Thu, 25 Mar 2021 05:44:43 +0000 Subject: [PATCH 17/19] fix incorrect function signature for CUDA 10 and smaller --- src/tensors/gpu/prod.cpp | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/src/tensors/gpu/prod.cpp b/src/tensors/gpu/prod.cpp index 530b7a48d..8cfa78cab 100755 --- a/src/tensors/gpu/prod.cpp +++ b/src/tensors/gpu/prod.cpp @@ -582,7 +582,11 @@ void affineTyped(marian::Tensor C, Ptr allocator, const marian::Tenso // This version is needed so that Windows doesn't complain when compiling CUDA < 11. Otherwise, the ifdef could be inside of one // definition of Affine. -void Affine(marian::Tensor C, Ptr allocator, const marian::Tensor& A, const marian::Tensor& B, const marian::Tensor& bias, +void Affine(marian::Tensor C, + Ptr allocator, + const marian::Tensor& A, + const marian::Tensor& B, + const marian::Tensor& bias, bool transA, bool transB, float beta, float scalar, bool do_relu) { // There is a bug in CUDA 11 where the bias pointer needs to be 8 byte aligned. This bug will be fix in a subsequent release. For now, // we launch a custom epilogue if the bias does not meet the alignment requirement. @@ -610,7 +614,6 @@ void Affine(marian::Tensor C, const marian::Tensor& A, const marian::Tensor& B, const marian::Tensor& bias, - const marian::Tensor& /*ones*/, bool transA, bool transB, float beta, float scalar, bool do_relu) { if(C->type() == Type::float32) { From 5c7b6eec494943536b0bf134bd286308a7861d65 Mon Sep 17 00:00:00 2001 From: Marcin Junczys-Dowmunt Date: Thu, 25 Mar 2021 06:51:41 +0000 Subject: [PATCH 18/19] only use AffineWithReluNodeOp for float types --- src/graph/expression_operators.cpp | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/src/graph/expression_operators.cpp b/src/graph/expression_operators.cpp index 7f01454e3..63cedb59d 100644 --- a/src/graph/expression_operators.cpp +++ b/src/graph/expression_operators.cpp @@ -580,7 +580,10 @@ Expr affine(Expr a, Expr b, Expr bias, bool transA, bool transB, float scale) { Expr affineWithRelu(Expr a, Expr b, Expr bias, bool transA, bool transB, float scale) { auto graph = a->graph(); - if(graph->isInference()) + Type aElementType = a->value_type(); + Type bElementType = b->value_type(); + + if(graph->isInference() && isFloat(aElementType) && isFloat(bElementType)) return Expression(a, b, bias, transA, transB, scale); else return relu(affine(a, b, bias, transA, transB, scale)); From c027c6d3d60f17560083ec6122ab703e2ed82853 Mon Sep 17 00:00:00 2001 From: Marcin Junczys-Dowmunt Date: Thu, 25 Mar 2021 14:26:42 +0000 Subject: [PATCH 19/19] switch back to using affineWithRelu on gpu only for now --- src/graph/expression_operators.cpp | 4 +--- src/graph/node_operators_binary.h | 8 ++++++-- 2 files changed, 7 insertions(+), 5 deletions(-) diff --git a/src/graph/expression_operators.cpp b/src/graph/expression_operators.cpp index 63cedb59d..048c74789 100644 --- a/src/graph/expression_operators.cpp +++ b/src/graph/expression_operators.cpp @@ -580,10 +580,8 @@ Expr affine(Expr a, Expr b, Expr bias, bool transA, bool transB, float scale) { Expr affineWithRelu(Expr a, Expr b, Expr bias, bool transA, bool transB, float scale) { auto graph = a->graph(); - Type aElementType = a->value_type(); - Type bElementType = b->value_type(); - if(graph->isInference() && isFloat(aElementType) && isFloat(bElementType)) + if(graph->isInference() && graph->getDeviceId().type == DeviceType::gpu) return Expression(a, b, bias, transA, transB, scale); else return relu(affine(a, b, bias, transA, transB, scale)); diff --git a/src/graph/node_operators_binary.h b/src/graph/node_operators_binary.h index f0a263e7b..55f105a96 100644 --- a/src/graph/node_operators_binary.h +++ b/src/graph/node_operators_binary.h @@ -428,7 +428,10 @@ class AffineWithReluNodeOp : public NaryNodeOp { : NaryNodeOp({a, b, bias}, newShape(a, b, transA, transB)), transA_(transA), transB_(transB), - scalar_(scalar) {} + scalar_(scalar) { + ABORT_IF(!graph()->isInference() || graph()->getDeviceId().type != DeviceType::gpu, + "AffineWithReluNodeOp currently only supported for inference on GPU"); + } Shape newShape(Expr a, Expr b, bool transA, bool transB) { auto shapeA = a->shape(); @@ -451,7 +454,8 @@ class AffineWithReluNodeOp : public NaryNodeOp { } NodeOps forwardOps() override { - using namespace functional; + ABORT_IF(!graph()->isInference() || graph()->getDeviceId().type != DeviceType::gpu, + "AffineWithReluNodeOp currently only supported for inference on GPU"); return { NodeOp(Affine(val_,