From 1100596cabc570a10286ba0c0ca822449b7640c8 Mon Sep 17 00:00:00 2001 From: Rawn Henry Date: Thu, 15 Oct 2020 12:40:18 -0700 Subject: [PATCH 01/17] Introduces a special operator to handle getting factor logits for inference when processing the lemmas --- src/data/factored_vocab.h | 6 ++ src/graph/expression_graph.h | 28 ++++++ src/graph/expression_operators.cpp | 18 ++++ src/graph/expression_operators.h | 7 ++ src/graph/node_operators_binary.h | 71 +++++++++++++ src/layers/generic.cpp | 64 ++++++++++-- src/layers/generic.h | 6 ++ src/tensors/cpu/tensor_operators.cpp | 15 +++ src/tensors/gpu/tensor_operators.cu | 145 +++++++++++++++++++++++++++ src/tensors/tensor_operators.h | 7 ++ src/translator/beam_search.cpp | 8 +- 11 files changed, 368 insertions(+), 7 deletions(-) diff --git a/src/data/factored_vocab.h b/src/data/factored_vocab.h index 215e92f09..39b893b63 100755 --- a/src/data/factored_vocab.h +++ b/src/data/factored_vocab.h @@ -3,6 +3,11 @@ // and via dynamic_cast to FactoredVocab for factored-specific things used by // the Embedding and Output layers. +/* Part of this file was contributed by NVIDIA under license: + * Copyright (C) 2020 NVIDIA Corporation + * SPDX-License-Identifier: MIT + */ + #pragma once #include "common/definitions.h" @@ -69,6 +74,7 @@ class FactoredVocab : public IVocab { const std::string& getFactorGroupPrefix(size_t groupIndex) const { return groupPrefixes_[groupIndex]; } // for diagnostics only const std::string& getFactorName(size_t groupIndex, size_t factorIndex) const { return factorVocab_[(WordIndex)(factorIndex + groupRanges_[groupIndex].first)]; } std::string decodeForDiagnostics(const Words& sentence) const; + const std::vector>& getLemmaHasFactorGroupVector() const {return lemmaHasFactorGroup_;}; static constexpr size_t FACTOR_NOT_APPLICABLE = (SIZE_MAX - 1); static constexpr size_t FACTOR_NOT_SPECIFIED = (SIZE_MAX - 2); diff --git a/src/graph/expression_graph.h b/src/graph/expression_graph.h index b4f0c1e29..b5231dca8 100644 --- a/src/graph/expression_graph.h +++ b/src/graph/expression_graph.h @@ -1,3 +1,8 @@ +/* Part of this file was contributed by NVIDIA under license: + * Copyright (C) 2020 NVIDIA Corporation + * SPDX-License-Identifier: MIT + */ + #pragma once #include "common/config.h" @@ -27,6 +32,8 @@ class Tensors { typedef std::unordered_map> WeakMemory; typedef std::unordered_map> Memory; + std::map memoizationMap_; + Ptr shortterm_; Ptr longterm_; @@ -100,6 +107,19 @@ class Tensors { return nullptr; } + void rememberByName(const std::string& name, Expr e) { + ABORT_IF(e == nullptr, "Expression must be non-null"); + ABORT_IF(e->type() == "param", "Not intended for graph parameters"); + memoizationMap_[name] = e; + findOrRemember(e); + } + + Expr findByName(const std::string&name) { + if(memoizationMap_.count(name)) + return findOrRemember(memoizationMap_[name]); + return nullptr; + } + void clear() { tensors_->clear(); shortterm_->clear(); @@ -471,6 +491,14 @@ class ExpressionGraph : public std::enable_shared_from_this { // Returns the tensor allocator of the graph workspace, different from above as proper tensor objects are allocated Ptr getTensorAllocator() { return tensors_->getTensorAllocator(); } + void rememberByName(const std::string& name, Expr e) { + tensors_->rememberByName(name, e); + } + + Expr findByName(const std::string&name) { + return tensors_->findByName(name); + } + void clear() { // clear everything apart from parameters and memoized nodes count_ = 0; diff --git a/src/graph/expression_operators.cpp b/src/graph/expression_operators.cpp index f571f9f8b..54f438228 100644 --- a/src/graph/expression_operators.cpp +++ b/src/graph/expression_operators.cpp @@ -1,3 +1,7 @@ +/* Part of this file was contributed by NVIDIA under license: + * Copyright (C) 2020 NVIDIA Corporation + * SPDX-License-Identifier: MIT + */ #include "graph/expression_operators.h" #include "layers/constructors.h" @@ -679,6 +683,20 @@ Expr unlikelihood(Expr logits, Expr indices) { return -log(gather(1.f - softmax(logits), /*axis=*/-1, indicesWithLayout)); } +Expr addFactorMaxes(Expr lemmaHasFactorGroup, std::vector groupLosses, Expr hypIndices, size_t groupStart, size_t numLemmas) { + if(groupLosses.size() == 1) { + return groupLosses[0]; + } + std::vector nodes({lemmaHasFactorGroup}); + if (hypIndices) { + nodes.push_back(hypIndices); + } + nodes.insert(nodes.end(), groupLosses.begin(), groupLosses.end()); + bool hasShortList = hypIndices != nullptr; + return Expression(nodes, hasShortList, groupStart, numLemmas); +} + + Expr plus(const std::vector& nodes) { ABORT_IF(nodes.size() > 1, "Not implemented"); return nodes[0]; diff --git a/src/graph/expression_operators.h b/src/graph/expression_operators.h index a4a2eeee4..72a6aba99 100755 --- a/src/graph/expression_operators.h +++ b/src/graph/expression_operators.h @@ -1,3 +1,8 @@ +/* Part of this file was contributed by NVIDIA under license: + * Copyright (C) 2020 NVIDIA Corporation + * SPDX-License-Identifier: MIT + */ + #pragma once #include "graph/expression_graph.h" #include "graph/node_initializers.h" @@ -172,6 +177,8 @@ Expr atleast_3d(Expr a); Expr atleast_4d(Expr a); Expr atleast_nd(Expr a, size_t dims); +Expr addFactorMaxes(Expr lemmaHasFactorGroup, std::vector groupLosses, Expr hypIndices, size_t groupStart, size_t numLemmas); + // create a constant of shape a->shape() and initialize with init // @TODO: add a && version, to avoid a ref count. NodeInitializers are typically temps. // @TODO: and/or make this a template on init diff --git a/src/graph/node_operators_binary.h b/src/graph/node_operators_binary.h index d596619b2..af3cd9765 100644 --- a/src/graph/node_operators_binary.h +++ b/src/graph/node_operators_binary.h @@ -1,3 +1,7 @@ +/* Part of this file was contributed by NVIDIA under license: + * Copyright (C) 2020 NVIDIA Corporation + * SPDX-License-Identifier: MIT + */ #pragma once #include @@ -1245,6 +1249,73 @@ struct LayerNormalizationOp : public NaryNodeOp { float eps_; }; +struct AddFactorMaxesOp : public NaryNodeOp { +size_t groupStart_; +size_t numLemmas_; +bool hasShortlist_; +public: + AddFactorMaxesOp(const std::vector& nodes, bool hasShortlist, size_t groupStart, size_t numLemmas) + : NaryNodeOp(nodes, getShape(nodes, hasShortlist), nodes[hasShortlist? 2 : 1]->value_type()) { + groupStart_ = groupStart; + numLemmas_ = hasShortlist? nodes[1]->shape().size(): numLemmas; + hasShortlist_ = hasShortlist; + } + + Shape getShape(const std::vector& nodes, bool hasShortlist) { + ABORT_IF(nodes.empty(), "No child nodes given"); + int start = hasShortlist? 2 : 1; + return nodes[start]->shape(); + } + + NodeOps forwardOps() override { + int start = hasShortlist_? 2 : 1; + std::vector losses; + for(int i = start; i < children().size(); ++i) { + losses.push_back(child(i)->val()); + } + return {NodeOp( + AddFactorMaxes(val_, + graph()->allocator(), + child(0)->val(), // lemmaHasFactorGroupTensor + hasShortlist_? child(1)->val() : nullptr, // indices + losses, + groupStart_, numLemmas_))}; + } + + NodeOps backwardOps() override { + ABORT("Not Implemented for Training"); + } + + const std::string type() override { return "AddFactorMaxesOp"; } + + virtual size_t hash() override { + size_t seed = NaryNodeOp::hash(); + util::hash_combine(seed, hasShortlist_); + util::hash_combine(seed, groupStart_); + util::hash_combine(seed, numLemmas_); + 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(hasShortlist_ != cnode->hasShortlist_) + return false; + if(groupStart_ != cnode->groupStart_) + return false; + if(numLemmas_ != cnode->numLemmas_) + return false; + return true; + } + +private: + friend class SerializationHelpers; // @TODO: use the same name for this as SqrtNodeOp +}; + + struct HighwayNodeOp : public NaryNodeOp { HighwayNodeOp(const std::vector& nodes) : NaryNodeOp(nodes) {} diff --git a/src/layers/generic.cpp b/src/layers/generic.cpp index 6c760aace..84f0350fe 100755 --- a/src/layers/generic.cpp +++ b/src/layers/generic.cpp @@ -1,3 +1,8 @@ +/* Part of this file was contributed by NVIDIA under license: + * Copyright (C) 2020 NVIDIA Corporation + * SPDX-License-Identifier: MIT + */ + #include "marian.h" #include "layers/generic.h" @@ -77,15 +82,20 @@ namespace marian { // - lemma: add all maxes of applicable factors if (groupIndex > 0) { sel = sel - max(sel, -1); - } - else { + } else { auto numGroups = getNumFactorGroups(); - for (size_t g = 1; g < numGroups; g++) { - auto factorMaxima = max(logits_[g]->loss(), -1); - auto factorMasks = constant(getFactorMasks(g, shortlist ? shortlist->indices() : std::vector())); - sel = sel + factorMaxima * factorMasks; // those lemmas that don't have a factor get multiplied with 0 + if(numGroups > 1 && graph()->isInference() && graph()->getBackend()->getDeviceId().type == DeviceType::gpu) { + Expr shortlistExpr = shortlist? constant(shortlist->indices()) : nullptr; + sel = addFactorMaxesHelper(shortlistExpr); + } else { + for (size_t g = 1; g < numGroups; g++) { + auto factorMaxima = max(logits_[g]->loss(), -1); + Expr factorMasks = constant(getFactorMasks(g, shortlist ? shortlist->indices() : std::vector())); + sel = sel + factorMaxima * factorMasks; // those lemmas that don't have a factor get multiplied with 0 + } } } + // if selIdx are given, then we must reshuffle accordingly if (!hypIndices.empty()) // use the same function that shuffles decoder state @@ -186,6 +196,48 @@ namespace marian { return res; } + Expr Logits::addFactorMaxesHelper(Expr indices) const { + auto g = graph(); + const std::string name = "lemmaHasFactorGroup"; + auto lemmaHasFactorGroupTensor = g->findByName(name); + if(!lemmaHasFactorGroupTensor) { + // We want to make this a graph param so the GPU can use this to implement lemmaHasFactorGroup to avoid unneeded memcpys. + const auto lemmaHasFactorGroup = factoredVocab_->getLemmaHasFactorGroupVector(); + int dimLemma = (int)lemmaHasFactorGroup.size(); + int dimFactorGroup = (int)lemmaHasFactorGroup[0].size(); + + for(const auto& lemma : lemmaHasFactorGroup) { + // Paranoid check - Instead of aborting I think we can pad here and copy the array + ABORT_IF(lemma.size() != dimFactorGroup, "All groups must be the same size"); + } + + auto initFunc = [lemmaHasFactorGroup, dimLemma, dimFactorGroup] (Tensor t) { + std::vector flattened(dimLemma * dimFactorGroup); + int row = 0; + for(const auto& v : lemmaHasFactorGroup) { + int offset = row * dimFactorGroup; + for(int i = 0; i < v.size(); ++i) { + flattened[offset + i] = static_cast(v[i]); + } + ++row; + } + t->set(flattened); + }; + + lemmaHasFactorGroupTensor = graph()->constant({dimLemma, dimFactorGroup}, inits::fromLambda(initFunc), Type::int8); + lemmaHasFactorGroupTensor->setMemoize(true); + g->rememberByName(name, lemmaHasFactorGroupTensor); + } + size_t n = indices ? indices->shape().elements() : (factoredVocab_->getGroupRange(0).second - factoredVocab_->getGroupRange(0).first); + + std::vector groupLosses(getNumFactorGroups()); + for(int g = 0; g < getNumFactorGroups(); ++g) { + groupLosses[g] = logits_[g]->loss(); + } + + return addFactorMaxes(lemmaHasFactorGroupTensor, groupLosses, indices, factoredVocab_->getGroupRange(0).first, n); + } + Logits Logits::applyUnaryFunction(const std::function& f) const { // clone this but apply f to all loss values std::vector> newLogits; for (const auto& l : logits_) diff --git a/src/layers/generic.h b/src/layers/generic.h index e83663357..8bb4962f5 100755 --- a/src/layers/generic.h +++ b/src/layers/generic.h @@ -1,3 +1,8 @@ +/* Part of this file was contributed by NVIDIA under license: + * Copyright (C) 2020 NVIDIA Corporation + * SPDX-License-Identifier: MIT + */ + #pragma once #include "marian.h" @@ -142,6 +147,7 @@ class Logits { template Expr constant(const std::vector& data) const { return constant(Shape{(int)data.size()}, data); } // same as constant() but assuming vector Expr indices(const std::vector& data) const { return graph()->indices(data); } // actually the same as constant(data) for this data type std::vector getFactorMasks(size_t factorGroup, const std::vector& indices) const; + Expr addFactorMaxesHelper(Expr indices=nullptr) const; private: // members // @TODO: we don't use the RationalLoss component anymore, can be removed again, and replaced just by the Expr diff --git a/src/tensors/cpu/tensor_operators.cpp b/src/tensors/cpu/tensor_operators.cpp index 211283d58..bd88289cb 100755 --- a/src/tensors/cpu/tensor_operators.cpp +++ b/src/tensors/cpu/tensor_operators.cpp @@ -3,6 +3,11 @@ * SPDX-License-Identifier: MIT */ +/* Part of this file was contributed by NVIDIA under license: + * Copyright (C) 2020 NVIDIA Corporation + * SPDX-License-Identifier: MIT + */ + #include "tensors/tensor_operators.h" #include "tensors/cpu/backend.h" #include "tensors/allocator.h" @@ -24,6 +29,16 @@ namespace cpu { ABORT("Not implemented"); } +void AddFactorMaxes(Tensor /*out*/, + Ptr /*allocator*/, + const Tensor /*lemmaHasFactorGroupTensor*/, + const Tensor /*indices*/, + const std::vector& /*groupLosses*/, + size_t /*groupStart*/, + size_t /*numLemmas*/) { + ABORT("AddFactorMaxes not implemented on CPU"); +} + template void CopyCastTo(To* out, const From* in, int length) { for(int i = 0; i < length; ++i) diff --git a/src/tensors/gpu/tensor_operators.cu b/src/tensors/gpu/tensor_operators.cu index 2552b7c7e..92a0ab86f 100644 --- a/src/tensors/gpu/tensor_operators.cu +++ b/src/tensors/gpu/tensor_operators.cu @@ -1,3 +1,10 @@ +/* Part of this file was contributed by NVIDIA under license: + * Copyright (C) 2020 NVIDIA Corporation + * SPDX-License-Identifier: MIT + */ + +#include +#include "common/logging.h" #include "common/types.h" #include "tensors/tensor_operators.h" @@ -2930,5 +2937,143 @@ void PoolingWithMaskingBackward(Tensor adj, width, lastWidth); } + +// Finds max within warp. Assumes val in non-participating warps set to T::min() +// or a valid non-garbage value +template +__device__ __forceinline__ T dWarpReduceMax(T val) { + constexpr unsigned fullMask = 0xffffffff; + #pragma unroll + for (int offset = (warpSize / 2); offset > 0; offset /= 2) { + val = max(val, __shfl_down_sync(fullMask, val, offset)); + } + return val; +} + +// Elts to reduce must be <= blockDim.x. Needed since factorDims can be varied lengths. +template +__device__ __forceinline__ void dBlockReduceMax(T val, int eltsToReduce, volatile T* shared, volatile T* writeLoc) { + const int lane = threadIdx.x % warpSize; + const int wid = threadIdx.x / warpSize; + const int warpsNeeded = (eltsToReduce + warpSize - 1) / warpSize; + + if (wid < warpsNeeded) val = dWarpReduceMax(val); + if (lane==0) shared[wid]=val; + __syncthreads(); // Needed to ensure all threads write to shared before reading. + + //read from shared memory lane only if that warp existed. Otherwise just get first lane. + val = (threadIdx.x < warpsNeeded) ? shared[lane] : shared[0]; + if(wid==0) { + T temp = dWarpReduceMax(val); + if(lane == 0) writeLoc[0] = temp; + } +} + +template +struct ptrInnerDimPair { + T* ptr; + int innerDim; +}; + +template +__global__ void gAddFactorMaxes(T* out, const int8_t* const lemmaHasFactorGroup, const IndexType* const indices, + ptrInnerDimPair* lossAndLastDimSize, size_t numGroups, size_t sizeWithoutInnerDim, + size_t groupStart, int lemmaHasFactorGroupWidth, size_t numLemmas, T minimal) { + + extern __shared__ T _sharedMem[]; + T* sel = _sharedMem; + T* factorMaximasInBlock = _sharedMem + blockDim.x; + T shared[32]; + + // First, each block computes the max across the row for each group and stores in in factorMaximasInBlock[g] + for(int g = 1; g < (int)numGroups; ++g) { + T* groupLosses = lossAndLastDimSize[g].ptr; + int groupLossesInnerDim = lossAndLastDimSize[g].innerDim; + int groupLossSize = groupLossesInnerDim * sizeWithoutInnerDim; + + // Now, we get the max for the factors. If the size of the inner dim is greater than the blocksize, + // first reduce maxes so they fit within a block. + T factorMaxima = minimal; + int blockRowStart = blockIdx.x * groupLossesInnerDim; + for(int lossCol = threadIdx.x; lossCol < groupLossesInnerDim; lossCol += blockDim.x) { + int offset = blockRowStart + lossCol; + if(offset < groupLossSize) { + factorMaxima = max(factorMaxima, groupLosses[offset]); + } + } + dBlockReduceMax(factorMaxima, min(blockDim.x, groupLossesInnerDim), shared, factorMaximasInBlock + g); + __syncthreads(); + } + + // Each block has the max for each factor, so we iterate over the groups again and accumulate the + // output in shared mem one block at a time before writing the results for the row to out + T* selGlobal = lossAndLastDimSize[0].ptr; + const int outTensorInnerDim = lossAndLastDimSize[0].innerDim; + const int outputSize = sizeWithoutInnerDim * outTensorInnerDim; + const int outRowOffset = blockIdx.x * outTensorInnerDim; + + for(int offset = threadIdx.x; offset < outTensorInnerDim; offset += blockDim.x) { + const int outOffset = outRowOffset + offset; + if(outOffset < outputSize) sel[threadIdx.x] = selGlobal[outOffset]; + + // Accumulate the row's portion in shared memory + for(int g = 1; g < (int)numGroups; ++g) { + int lemma = indices? indices[offset] - groupStart: offset; + T factorMask = static_cast(lemmaHasFactorGroup[lemma * lemmaHasFactorGroupWidth + g]); + T factorMaxima = factorMaximasInBlock[g]; + + if(outOffset < outputSize) { + sel[threadIdx.x] += (factorMask * factorMaxima); + } + } + if(outOffset < outputSize) out[outOffset] = sel[threadIdx.x]; + } +} + +void AddFactorMaxes(Tensor out, + Ptr allocator, + const Tensor lemmaHasFactorGroupTensor, + const Tensor indices, + const std::vector& groupLosses, + size_t groupStart, + size_t numLemmas) { + + cudaSetDevice(out->getDeviceId().no); + + ABORT_IF(lemmaHasFactorGroupTensor->type() != Type::int8, "lemmaHasFactor group tensor wrong type"); + ABORT_IF(out->shape()[-1] != (int) numLemmas, "Output shape {} or numLemmas {} incorrect", out->shape(), numLemmas); + + const int threads = std::min(MAX_THREADS, std::max(32, out->shape()[-1])); + const int sizeWithoutInnerDim = out->shape().elements() / out->shape()[-1]; + + if(out->type() == Type::float32) { + IPtr mp_ptrs = allocator->alloc>(groupLosses.size()); + ptrInnerDimPair* dest = mp_ptrs->data>(); + + std::vector> lossPtrs; + for(const auto& t : groupLosses) { + ptrInnerDimPair pair = {t->data(), t->shape()[-1]}; + lossPtrs.push_back(pair); + } + + // Async just so that call is issued in per thread default stream + CUDA_CHECK(cudaMemcpyAsync(dest, lossPtrs.data(), lossPtrs.size() * sizeof(ptrInnerDimPair), cudaMemcpyHostToDevice, 0)); + const int blocks = sizeWithoutInnerDim; + const int sharedBytes = sizeof(float) * (groupLosses.size() + threads); + gAddFactorMaxes<<>>(out->data(), + lemmaHasFactorGroupTensor->data(), + indices? indices->data(): nullptr, + dest, + groupLosses.size(), + sizeWithoutInnerDim, + groupStart, lemmaHasFactorGroupTensor->shape()[1], numLemmas, + std::numeric_limits::lowest()); + CUDA_CHECK(cudaStreamSynchronize(0)); + allocator->free(mp_ptrs); + } else { + ABORT("AddFactorMaxes not implemented for type {}", out->type()); + } +} + } // namespace gpu } // namespace marian diff --git a/src/tensors/tensor_operators.h b/src/tensors/tensor_operators.h index e075244f5..5d3df5dc9 100644 --- a/src/tensors/tensor_operators.h +++ b/src/tensors/tensor_operators.h @@ -1,3 +1,8 @@ +/* Part of this file was contributed by NVIDIA under license: + * Copyright (C) 2020 NVIDIA Corporation + * SPDX-License-Identifier: MIT + */ + #pragma once #include "common/definitions.h" @@ -120,6 +125,8 @@ DISPATCH4(ShiftGrad, marian::Tensor, marian::Tensor, marian::Shape, bool) DISPATCH3(Concatenate, marian::Tensor, const std::vector&, int) +DISPATCH7(AddFactorMaxes, marian::Tensor, Ptr, const marian::Tensor, const marian::Tensor, const std::vector&, size_t, size_t) + // clang-format on // Bernoulli(tensor, 0.5f, 2.f, -1.f) generates a tensor composed of 50% of 1 and 50% of -1. diff --git a/src/translator/beam_search.cpp b/src/translator/beam_search.cpp index 9335c55bd..3c39b2cf7 100755 --- a/src/translator/beam_search.cpp +++ b/src/translator/beam_search.cpp @@ -1,3 +1,8 @@ +/* Part of this file was contributed by NVIDIA under license: + * Copyright (C) 2020 NVIDIA Corporation + * SPDX-License-Identifier: MIT + */ + #include "translator/beam_search.h" #include "data/factored_vocab.h" @@ -396,7 +401,8 @@ Histories BeamSearch::search(Ptr graph, Ptr } if(factorGroup == 0) currentDimBatch = (IndexType) batchIndices.size(); // keep batch size constant for all factor groups in a time step - prevPathScores = graph->constant({(int)maxBeamSize, 1, (int)currentDimBatch, 1}, inits::fromVector(prevScores)); + // Avoid unnecessary memcpy on GPU + if(anyCanExpand) prevPathScores = graph->constant({(int)maxBeamSize, 1, (int)currentDimBatch, 1}, inits::fromVector(prevScores)); } if (!anyCanExpand) // all words cannot expand this factor: skip continue; From 8694fd93150b87a065693488fd5b8f92aa539577 Mon Sep 17 00:00:00 2001 From: Rawn Henry Date: Fri, 30 Oct 2020 12:36:46 -0700 Subject: [PATCH 02/17] Adds fp16 support for AddFactorMaxes --- src/graph/node_operators_binary.h | 2 +- src/tensors/gpu/tensor_operators.cu | 39 ++++++++++++++++++++++++++--- 2 files changed, 37 insertions(+), 4 deletions(-) diff --git a/src/graph/node_operators_binary.h b/src/graph/node_operators_binary.h index af3cd9765..e928628f7 100644 --- a/src/graph/node_operators_binary.h +++ b/src/graph/node_operators_binary.h @@ -1255,7 +1255,7 @@ size_t numLemmas_; bool hasShortlist_; public: AddFactorMaxesOp(const std::vector& nodes, bool hasShortlist, size_t groupStart, size_t numLemmas) - : NaryNodeOp(nodes, getShape(nodes, hasShortlist), nodes[hasShortlist? 2 : 1]->value_type()) { + : NaryNodeOp(nodes, getShape(nodes, hasShortlist), commonType(std::vector(nodes.begin() + 1, nodes.end())) ) { groupStart_ = groupStart; numLemmas_ = hasShortlist? nodes[1]->shape().size(): numLemmas; hasShortlist_ = hasShortlist; diff --git a/src/tensors/gpu/tensor_operators.cu b/src/tensors/gpu/tensor_operators.cu index 92a0ab86f..c530eda20 100644 --- a/src/tensors/gpu/tensor_operators.cu +++ b/src/tensors/gpu/tensor_operators.cu @@ -16,6 +16,13 @@ #include "tensors/gpu/add_all.h" +#if COMPILE_FP16 +#include +__device__ __forceinline__ half max(const half a, const half b) { + return a > b ? a : b; +} +#endif + namespace marian { namespace gpu { @@ -2980,9 +2987,9 @@ __global__ void gAddFactorMaxes(T* out, const int8_t* const lemmaHasFactorGroup, ptrInnerDimPair* lossAndLastDimSize, size_t numGroups, size_t sizeWithoutInnerDim, size_t groupStart, int lemmaHasFactorGroupWidth, size_t numLemmas, T minimal) { - extern __shared__ T _sharedMem[]; - T* sel = _sharedMem; - T* factorMaximasInBlock = _sharedMem + blockDim.x; + extern __shared__ uint8_t _sharedBytes[]; + T* sel = (T*)_sharedBytes; + T* factorMaximasInBlock = sel + blockDim.x; T shared[32]; // First, each block computes the max across the row for each group and stores in in factorMaximasInBlock[g] @@ -3070,6 +3077,32 @@ void AddFactorMaxes(Tensor out, std::numeric_limits::lowest()); CUDA_CHECK(cudaStreamSynchronize(0)); allocator->free(mp_ptrs); + #if COMPILE_FP16 + } else if(out->type() == Type::float16) { + IPtr mp_ptrs = allocator->alloc>(groupLosses.size()); + ptrInnerDimPair* dest = mp_ptrs->data>(); + + std::vector> lossPtrs; + for(const auto& t : groupLosses) { + ptrInnerDimPair pair = {t->data(), t->shape()[-1]}; + lossPtrs.push_back(pair); + } + + // Async just so that call is issued in per thread default stream + CUDA_CHECK(cudaMemcpyAsync(dest, lossPtrs.data(), lossPtrs.size() * sizeof(ptrInnerDimPair), cudaMemcpyHostToDevice, 0)); + const int blocks = sizeWithoutInnerDim; + const int sharedBytes = sizeof(half) * (groupLosses.size() + threads); + gAddFactorMaxes<<>>(out->data(), + lemmaHasFactorGroupTensor->data(), + indices? indices->data(): nullptr, + dest, + groupLosses.size(), + sizeWithoutInnerDim, + groupStart, lemmaHasFactorGroupTensor->shape()[1], numLemmas, + __float2half(std::numeric_limits::lowest()) ); + CUDA_CHECK(cudaStreamSynchronize(0)); + allocator->free(mp_ptrs); + #endif } else { ABORT("AddFactorMaxes not implemented for type {}", out->type()); } From 0fd9aa1f2bb13a30b8365d4b74cb44b433e3873c Mon Sep 17 00:00:00 2001 From: Rawn Henry Date: Sat, 17 Oct 2020 14:53:30 -0700 Subject: [PATCH 03/17] Adds cub as a submodule --- .gitmodules | 3 +++ src/3rd_party/cub | 1 + 2 files changed, 4 insertions(+) create mode 160000 src/3rd_party/cub diff --git a/.gitmodules b/.gitmodules index 6cb63fc0b..e3de79a86 100644 --- a/.gitmodules +++ b/.gitmodules @@ -17,3 +17,6 @@ [submodule "src/3rd_party/simple-websocket-server"] path = src/3rd_party/simple-websocket-server url = https://github.com/marian-nmt/Simple-WebSocket-Server +[submodule "src/3rd_party/cub"] + path = src/3rd_party/cub + url = https://github.com/NVIDIA/cub diff --git a/src/3rd_party/cub b/src/3rd_party/cub new file mode 160000 index 000000000..52d58a889 --- /dev/null +++ b/src/3rd_party/cub @@ -0,0 +1 @@ +Subproject commit 52d58a88904da39c374e44a6a8ae0e4dcca5b71a From d700727a037f80f5450a6febe6ad82ede1d083c6 Mon Sep 17 00:00:00 2001 From: Rawn Henry Date: Fri, 11 Dec 2020 14:39:09 -0800 Subject: [PATCH 04/17] Adds definitions for cub compilation --- src/CMakeLists.txt | 2 ++ 1 file changed, 2 insertions(+) diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index 6dcf7fd89..db97da7d3 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -1,3 +1,5 @@ +add_definitions(-DCUB_IGNORE_DEPRECATED_CPP_DIALECT=1) +add_definitions(-DTHRUST_IGNORE_DEPRECATED_CPP_DIALECT=1) add_subdirectory(3rd_party) include_directories(.) From f7aba6d6247c47388e4d0e5bbd9ef9d880b6bdc5 Mon Sep 17 00:00:00 2001 From: Rawn Henry Date: Fri, 11 Dec 2020 14:39:38 -0800 Subject: [PATCH 05/17] Bug fix in add factor maxes --- src/tensors/gpu/tensor_operators.cu | 55 +++++++++++++++++++++-------- 1 file changed, 40 insertions(+), 15 deletions(-) diff --git a/src/tensors/gpu/tensor_operators.cu b/src/tensors/gpu/tensor_operators.cu index c530eda20..c5d04caab 100644 --- a/src/tensors/gpu/tensor_operators.cu +++ b/src/tensors/gpu/tensor_operators.cu @@ -23,6 +23,13 @@ __device__ __forceinline__ half max(const half a, const half b) { } #endif + +#if CUDA_VERSION >= 11000 +#include +#else +#include "cub/cub/cub.cuh" +#endif + namespace marian { namespace gpu { @@ -2990,28 +2997,46 @@ __global__ void gAddFactorMaxes(T* out, const int8_t* const lemmaHasFactorGroup, extern __shared__ uint8_t _sharedBytes[]; T* sel = (T*)_sharedBytes; T* factorMaximasInBlock = sel + blockDim.x; - T shared[32]; - - // First, each block computes the max across the row for each group and stores in in factorMaximasInBlock[g] - for(int g = 1; g < (int)numGroups; ++g) { - T* groupLosses = lossAndLastDimSize[g].ptr; - int groupLossesInnerDim = lossAndLastDimSize[g].innerDim; + typedef cub::WarpReduce WarpReduce; + __shared__ typename WarpReduce::TempStorage temp_storage[32]; // Max thread size (1024) divided by warp size (32) + + + // We exploit the fact that the factor groups (except for the lemmas) tends to be small. Therefore, we perform + // all of the reductions across factor groups in parallel by splitting the work across warps. + const int lane = threadIdx.x % warpSize; + const int wid = threadIdx.x / warpSize; + const int numWarps = blockDim.x / warpSize; + + for(int warpFactorGroup = wid + 1; warpFactorGroup < numGroups; warpFactorGroup+=numWarps) { + T* groupLosses = lossAndLastDimSize[warpFactorGroup].ptr; + int groupLossesInnerDim = lossAndLastDimSize[warpFactorGroup].innerDim; int groupLossSize = groupLossesInnerDim * sizeWithoutInnerDim; - // Now, we get the max for the factors. If the size of the inner dim is greater than the blocksize, - // first reduce maxes so they fit within a block. T factorMaxima = minimal; - int blockRowStart = blockIdx.x * groupLossesInnerDim; - for(int lossCol = threadIdx.x; lossCol < groupLossesInnerDim; lossCol += blockDim.x) { - int offset = blockRowStart + lossCol; - if(offset < groupLossSize) { - factorMaxima = max(factorMaxima, groupLosses[offset]); + // Each block computes the factor maxes within a given row. + // We start by each warp computing offset to the row its block is responsible for in the loss + + // This is per warp since groupLossesInnerDim is warp specific and the groupLossesPtr is also warp specific. + const int blockRowStart = blockIdx.x * groupLossesInnerDim; + for(int lossCol = lane; lossCol < groupLossesInnerDim; lossCol += warpSize) { + int warpOffset = blockRowStart + lossCol; + if(warpOffset < groupLossSize) { + factorMaxima = max(factorMaxima, groupLosses[warpOffset]); } } - dBlockReduceMax(factorMaxima, min(blockDim.x, groupLossesInnerDim), shared, factorMaximasInBlock + g); - __syncthreads(); + + // Each warp reduces the accumulated values. Warps can do this in parallel. Then each each writes its final value to shared mem. + factorMaxima = WarpReduce(temp_storage[wid]).Reduce(factorMaxima, cub::Max()); + if(lane == 0) { + factorMaximasInBlock[warpFactorGroup] = factorMaxima; + } } + // Wait for all warps to write out to shared mem so the second phase of this kernel can proceed. + __syncthreads(); + + // First, each block computes the max across the row for each group and stores in in factorMaximasInBlock[g] + // Each block has the max for each factor, so we iterate over the groups again and accumulate the // output in shared mem one block at a time before writing the results for the row to out T* selGlobal = lossAndLastDimSize[0].ptr; From 7aba3b40538f3666fd6ff54267896e9dac06947b Mon Sep 17 00:00:00 2001 From: Rawn Henry Date: Tue, 10 Nov 2020 16:52:48 -0800 Subject: [PATCH 06/17] Removes sync before free in addFactorMaxes. I think it is not needed since the allocator has a memory pool that it manages for it won't get released by a cuda free. Additionally, two kernels may get the same pointer but they cannot execute concurrently since a single thread does not launch concurrent kernels. Since there is an allocator per thread, this means that no two kernels can ever race on the same pointer (I think). I have not seen any issues after removing this sync --- src/tensors/gpu/tensor_operators.cu | 2 -- 1 file changed, 2 deletions(-) diff --git a/src/tensors/gpu/tensor_operators.cu b/src/tensors/gpu/tensor_operators.cu index c5d04caab..5dc5da28f 100644 --- a/src/tensors/gpu/tensor_operators.cu +++ b/src/tensors/gpu/tensor_operators.cu @@ -3100,7 +3100,6 @@ void AddFactorMaxes(Tensor out, sizeWithoutInnerDim, groupStart, lemmaHasFactorGroupTensor->shape()[1], numLemmas, std::numeric_limits::lowest()); - CUDA_CHECK(cudaStreamSynchronize(0)); allocator->free(mp_ptrs); #if COMPILE_FP16 } else if(out->type() == Type::float16) { @@ -3125,7 +3124,6 @@ void AddFactorMaxes(Tensor out, sizeWithoutInnerDim, groupStart, lemmaHasFactorGroupTensor->shape()[1], numLemmas, __float2half(std::numeric_limits::lowest()) ); - CUDA_CHECK(cudaStreamSynchronize(0)); allocator->free(mp_ptrs); #endif } else { From baab83f8a01271986131c501b9bdaaf1c9d50fdb Mon Sep 17 00:00:00 2001 From: Rawn Henry Date: Mon, 30 Nov 2020 21:41:36 -0800 Subject: [PATCH 07/17] WIP - Rework addFactorMaxes. Starts splitting it into two kernels to expose more parallelism when adding into the lemmas --- src/tensors/gpu/tensor_operators.cu | 155 ++++++++++++++++++---------- 1 file changed, 102 insertions(+), 53 deletions(-) diff --git a/src/tensors/gpu/tensor_operators.cu b/src/tensors/gpu/tensor_operators.cu index 5dc5da28f..371957005 100644 --- a/src/tensors/gpu/tensor_operators.cu +++ b/src/tensors/gpu/tensor_operators.cu @@ -786,7 +786,7 @@ void LogSoftmax(Tensor out, Tensor in) { size_t m = out->shape().elements() / out->shape().back(); size_t k = out->shape().back(); - int blocks = std::min(MAX_BLOCKS, (int)m); + int blocks = (int)m; int threads = std::min(MAX_THREADS, (int)k); int shared = sizeof(float) * threads; // use float32 as accumulation type @@ -2990,13 +2990,11 @@ struct ptrInnerDimPair { }; template -__global__ void gAddFactorMaxes(T* out, const int8_t* const lemmaHasFactorGroup, const IndexType* const indices, - ptrInnerDimPair* lossAndLastDimSize, size_t numGroups, size_t sizeWithoutInnerDim, - size_t groupStart, int lemmaHasFactorGroupWidth, size_t numLemmas, T minimal) { +__global__ void gAddFactorMaxesPhase1(T* maxes, // [#blocks, numGroups - 1] + ptrInnerDimPair* lossAndLastDimSize, + size_t numGroups, size_t sizeWithoutInnerDim, + T minimal) { - extern __shared__ uint8_t _sharedBytes[]; - T* sel = (T*)_sharedBytes; - T* factorMaximasInBlock = sel + blockDim.x; typedef cub::WarpReduce WarpReduce; __shared__ typename WarpReduce::TempStorage temp_storage[32]; // Max thread size (1024) divided by warp size (32) @@ -3013,7 +3011,7 @@ __global__ void gAddFactorMaxes(T* out, const int8_t* const lemmaHasFactorGroup, int groupLossSize = groupLossesInnerDim * sizeWithoutInnerDim; T factorMaxima = minimal; - // Each block computes the factor maxes within a given row. + // Each block computes the maxes for each secondary factor within a given row. // We start by each warp computing offset to the row its block is responsible for in the loss // This is per warp since groupLossesInnerDim is warp specific and the groupLossesPtr is also warp specific. @@ -3025,40 +3023,68 @@ __global__ void gAddFactorMaxes(T* out, const int8_t* const lemmaHasFactorGroup, } } - // Each warp reduces the accumulated values. Warps can do this in parallel. Then each each writes its final value to shared mem. + // Each warp reduces the accumulated values. Then each each writes its final value to global for phase 2 of this op factorMaxima = WarpReduce(temp_storage[wid]).Reduce(factorMaxima, cub::Max()); if(lane == 0) { - factorMaximasInBlock[warpFactorGroup] = factorMaxima; + T* factorMaximasInBlock = maxes + blockIdx.x * (numGroups - 1); + factorMaximasInBlock[warpFactorGroup - 1] = factorMaxima; } } +} - // Wait for all warps to write out to shared mem so the second phase of this kernel can proceed. - __syncthreads(); - - // First, each block computes the max across the row for each group and stores in in factorMaximasInBlock[g] + + template +__global__ void gAddFactorMaxesPhase2(T* out, const int8_t* const lemmaHasFactorGroup, + const IndexType* const indices, + const T* lemmaLosses, + size_t numGroups, + const int sizeWithoutInnerDim, + const T* maxes, + size_t groupStart, + int lemmaHasFactorGroupWidth, + size_t numLemmas, T minimal) { + + extern __shared__ uint8_t _sharedBytes[]; + T* sums = (T*)_sharedBytes; + T* factorMaxesInBlockShared = sums + blockDim.x; - // Each block has the max for each factor, so we iterate over the groups again and accumulate the - // output in shared mem one block at a time before writing the results for the row to out - T* selGlobal = lossAndLastDimSize[0].ptr; - const int outTensorInnerDim = lossAndLastDimSize[0].innerDim; - const int outputSize = sizeWithoutInnerDim * outTensorInnerDim; - const int outRowOffset = blockIdx.x * outTensorInnerDim; - - for(int offset = threadIdx.x; offset < outTensorInnerDim; offset += blockDim.x) { - const int outOffset = outRowOffset + offset; - if(outOffset < outputSize) sel[threadIdx.x] = selGlobal[outOffset]; - - // Accumulate the row's portion in shared memory + // A group of blocks is assigned a row to process. We compute some the row index for each block in the group + // and the row each group of blocks is responsible for processing. + const int blockStartIndexInRow = blockIdx.x % BLOCKS_PER_ROW; + const int blockRow = blockIdx.x / BLOCKS_PER_ROW; + + // We need to know the total number of chunks in a row to safely iterate over the input and output + const int totalBlockChunksInRow = (numLemmas + blockDim.x - 1) / blockDim.x; + + // Calculate the row pointer for the global memory which holds for max for each factor. This is the same for + // all blocks in a given group + const T* factorMaxesInBlock = maxes + blockRow * (numGroups - 1); + T* outputRowForBlockGroup = out + blockRow * numLemmas; + const T* inputRowForBlockGroup = lemmaLosses + blockRow * numLemmas; + if(threadIdx.x < (numGroups - 1)) factorMaxesInBlockShared[threadIdx.x] = factorMaxesInBlock[threadIdx.x]; + __syncthreads(); + + for(int rowChunk = blockStartIndexInRow; rowChunk < totalBlockChunksInRow; rowChunk += BLOCKS_PER_ROW) { + // Compute the start of the chunk each block needs to read from in the input and output tensors + const T* lemmaLossesChunk = inputRowForBlockGroup + rowChunk * blockDim.x; + T* blockOutputChunk = outputRowForBlockGroup + rowChunk * blockDim.x; + + // Calculate the column being read to/written from to use as a read/write guard + const int columnIndex = threadIdx.x * rowChunk; + if(columnIndex < numLemmas) sums[threadIdx.x] = lemmaLossesChunk[threadIdx.x]; + + // Each block accumulates its output in shared memory then writes the final result to global. We use __ldg since the + // data being read does not change throughout the life of the kernel so the constant cache can be used. for(int g = 1; g < (int)numGroups; ++g) { - int lemma = indices? indices[offset] - groupStart: offset; - T factorMask = static_cast(lemmaHasFactorGroup[lemma * lemmaHasFactorGroupWidth + g]); - T factorMaxima = factorMaximasInBlock[g]; + int lemma = indices? indices[columnIndex] - groupStart: columnIndex; + T factorMask = static_cast(__ldg(&lemmaHasFactorGroup[lemma * lemmaHasFactorGroupWidth + g])); + T factorMaxima = factorMaxesInBlockShared[g-1]; - if(outOffset < outputSize) { - sel[threadIdx.x] += (factorMask * factorMaxima); + if(columnIndex < numLemmas) { + sums[threadIdx.x] += (factorMask * factorMaxima); } } - if(outOffset < outputSize) out[outOffset] = sel[threadIdx.x]; + if(columnIndex < numLemmas) blockOutputChunk[threadIdx.x] = sums[threadIdx.x]; } } @@ -3075,11 +3101,18 @@ void AddFactorMaxes(Tensor out, ABORT_IF(lemmaHasFactorGroupTensor->type() != Type::int8, "lemmaHasFactor group tensor wrong type"); ABORT_IF(out->shape()[-1] != (int) numLemmas, "Output shape {} or numLemmas {} incorrect", out->shape(), numLemmas); - const int threads = std::min(MAX_THREADS, std::max(32, out->shape()[-1])); const int sizeWithoutInnerDim = out->shape().elements() / out->shape()[-1]; + constexpr int warpSize = 32; + const int blocksPhase1 = sizeWithoutInnerDim; + const int threadsPhase1 = std::min(MAX_THREADS, (int)(warpSize * (groupLosses.size() - 1)) ); + + constexpr int blocksPerRow = 8; + const int threadsPhase2 = std::min(MAX_THREADS, std::max(32, out->shape()[-1])); + const int blocksPhase2 = sizeWithoutInnerDim * blocksPerRow; if(out->type() == Type::float32) { IPtr mp_ptrs = allocator->alloc>(groupLosses.size()); + IPtr factorMaxes = allocator->alloc(sizeWithoutInnerDim * (groupLosses.size() - 1)); ptrInnerDimPair* dest = mp_ptrs->data>(); std::vector> lossPtrs; @@ -3090,20 +3123,29 @@ void AddFactorMaxes(Tensor out, // Async just so that call is issued in per thread default stream CUDA_CHECK(cudaMemcpyAsync(dest, lossPtrs.data(), lossPtrs.size() * sizeof(ptrInnerDimPair), cudaMemcpyHostToDevice, 0)); - const int blocks = sizeWithoutInnerDim; - const int sharedBytes = sizeof(float) * (groupLosses.size() + threads); - gAddFactorMaxes<<>>(out->data(), - lemmaHasFactorGroupTensor->data(), - indices? indices->data(): nullptr, - dest, - groupLosses.size(), - sizeWithoutInnerDim, - groupStart, lemmaHasFactorGroupTensor->shape()[1], numLemmas, - std::numeric_limits::lowest()); + gAddFactorMaxesPhase1<<>>(factorMaxes->data(), // [#rows, numGroups - 1] + dest, + groupLosses.size(), + sizeWithoutInnerDim, + std::numeric_limits::lowest()); + + const int sharedBytes = sizeof(float) * (groupLosses.size() + threadsPhase2); + gAddFactorMaxesPhase2<<>>(out->data(), + lemmaHasFactorGroupTensor->data(), + indices? indices->data(): nullptr, + groupLosses[0]->data(), + groupLosses.size(), sizeWithoutInnerDim, + factorMaxes->data(), groupStart, + lemmaHasFactorGroupTensor->shape()[1], + numLemmas, + std::numeric_limits::lowest()); + allocator->free(mp_ptrs); + allocator->free(factorMaxes); #if COMPILE_FP16 } else if(out->type() == Type::float16) { IPtr mp_ptrs = allocator->alloc>(groupLosses.size()); + IPtr factorMaxes = allocator->alloc(sizeWithoutInnerDim * (groupLosses.size() - 1)); ptrInnerDimPair* dest = mp_ptrs->data>(); std::vector> lossPtrs; @@ -3114,17 +3156,24 @@ void AddFactorMaxes(Tensor out, // Async just so that call is issued in per thread default stream CUDA_CHECK(cudaMemcpyAsync(dest, lossPtrs.data(), lossPtrs.size() * sizeof(ptrInnerDimPair), cudaMemcpyHostToDevice, 0)); - const int blocks = sizeWithoutInnerDim; - const int sharedBytes = sizeof(half) * (groupLosses.size() + threads); - gAddFactorMaxes<<>>(out->data(), - lemmaHasFactorGroupTensor->data(), - indices? indices->data(): nullptr, - dest, - groupLosses.size(), - sizeWithoutInnerDim, - groupStart, lemmaHasFactorGroupTensor->shape()[1], numLemmas, - __float2half(std::numeric_limits::lowest()) ); + gAddFactorMaxesPhase1<<>>(factorMaxes->data(), // [#rows, numGroups - 1] + dest, + groupLosses.size(), + sizeWithoutInnerDim, + __float2half(std::numeric_limits::lowest()) ); + + const int sharedBytes = sizeof(half) * (groupLosses.size() + threadsPhase2); + gAddFactorMaxesPhase2<<>>(out->data(), + lemmaHasFactorGroupTensor->data(), + indices? indices->data(): nullptr, + groupLosses[0]->data(), + groupLosses.size(), sizeWithoutInnerDim, + factorMaxes->data(), groupStart, + lemmaHasFactorGroupTensor->shape()[1], + numLemmas, + __float2half(std::numeric_limits::lowest()) ); allocator->free(mp_ptrs); + allocator->free(factorMaxes); #endif } else { ABORT("AddFactorMaxes not implemented for type {}", out->type()); From f2d77e0377aa25788864599d18d1fa88dacff6d2 Mon Sep 17 00:00:00 2001 From: Rawn Henry Date: Fri, 11 Dec 2020 14:52:30 -0800 Subject: [PATCH 08/17] Stylistic changes to addFactorMaxes --- src/tensors/gpu/tensor_operators.cu | 20 ++++++++++++-------- 1 file changed, 12 insertions(+), 8 deletions(-) diff --git a/src/tensors/gpu/tensor_operators.cu b/src/tensors/gpu/tensor_operators.cu index 371957005..2deba8b16 100644 --- a/src/tensors/gpu/tensor_operators.cu +++ b/src/tensors/gpu/tensor_operators.cu @@ -3018,14 +3018,14 @@ __global__ void gAddFactorMaxesPhase1(T* maxes, // [#blocks, numGroups - 1] const int blockRowStart = blockIdx.x * groupLossesInnerDim; for(int lossCol = lane; lossCol < groupLossesInnerDim; lossCol += warpSize) { int warpOffset = blockRowStart + lossCol; - if(warpOffset < groupLossSize) { + if (warpOffset < groupLossSize) { factorMaxima = max(factorMaxima, groupLosses[warpOffset]); } } // Each warp reduces the accumulated values. Then each each writes its final value to global for phase 2 of this op factorMaxima = WarpReduce(temp_storage[wid]).Reduce(factorMaxima, cub::Max()); - if(lane == 0) { + if (lane == 0) { T* factorMaximasInBlock = maxes + blockIdx.x * (numGroups - 1); factorMaximasInBlock[warpFactorGroup - 1] = factorMaxima; } @@ -3061,7 +3061,10 @@ __global__ void gAddFactorMaxesPhase2(T* out, const int8_t* const lemmaHasFactor const T* factorMaxesInBlock = maxes + blockRow * (numGroups - 1); T* outputRowForBlockGroup = out + blockRow * numLemmas; const T* inputRowForBlockGroup = lemmaLosses + blockRow * numLemmas; - if(threadIdx.x < (numGroups - 1)) factorMaxesInBlockShared[threadIdx.x] = factorMaxesInBlock[threadIdx.x]; + + if (threadIdx.x < (numGroups - 1)) + factorMaxesInBlockShared[threadIdx.x] = factorMaxesInBlock[threadIdx.x]; + __syncthreads(); for(int rowChunk = blockStartIndexInRow; rowChunk < totalBlockChunksInRow; rowChunk += BLOCKS_PER_ROW) { @@ -3071,20 +3074,21 @@ __global__ void gAddFactorMaxesPhase2(T* out, const int8_t* const lemmaHasFactor // Calculate the column being read to/written from to use as a read/write guard const int columnIndex = threadIdx.x * rowChunk; - if(columnIndex < numLemmas) sums[threadIdx.x] = lemmaLossesChunk[threadIdx.x]; + if (columnIndex < numLemmas) + sums[threadIdx.x] = lemmaLossesChunk[threadIdx.x]; // Each block accumulates its output in shared memory then writes the final result to global. We use __ldg since the // data being read does not change throughout the life of the kernel so the constant cache can be used. for(int g = 1; g < (int)numGroups; ++g) { - int lemma = indices? indices[columnIndex] - groupStart: columnIndex; + int lemma = indices? indices[columnIndex] - groupStart : columnIndex; T factorMask = static_cast(__ldg(&lemmaHasFactorGroup[lemma * lemmaHasFactorGroupWidth + g])); T factorMaxima = factorMaxesInBlockShared[g-1]; - if(columnIndex < numLemmas) { + if (columnIndex < numLemmas) sums[threadIdx.x] += (factorMask * factorMaxima); - } } - if(columnIndex < numLemmas) blockOutputChunk[threadIdx.x] = sums[threadIdx.x]; + if (columnIndex < numLemmas) + blockOutputChunk[threadIdx.x] = sums[threadIdx.x]; } } From 706bbe7ff0dd4ce3acac09c7be5133cdd52ff113 Mon Sep 17 00:00:00 2001 From: Rawn Henry Date: Fri, 11 Dec 2020 15:01:41 -0800 Subject: [PATCH 09/17] Initializes reduction variable to negative inf instead of float::lowest --- src/tensors/gpu/tensor_operators.cu | 16 +++++++--------- 1 file changed, 7 insertions(+), 9 deletions(-) diff --git a/src/tensors/gpu/tensor_operators.cu b/src/tensors/gpu/tensor_operators.cu index 2deba8b16..f08b7ab93 100644 --- a/src/tensors/gpu/tensor_operators.cu +++ b/src/tensors/gpu/tensor_operators.cu @@ -3042,7 +3042,7 @@ __global__ void gAddFactorMaxesPhase2(T* out, const int8_t* const lemmaHasFactor const T* maxes, size_t groupStart, int lemmaHasFactorGroupWidth, - size_t numLemmas, T minimal) { + size_t numLemmas) { extern __shared__ uint8_t _sharedBytes[]; T* sums = (T*)_sharedBytes; @@ -3114,7 +3114,7 @@ void AddFactorMaxes(Tensor out, const int threadsPhase2 = std::min(MAX_THREADS, std::max(32, out->shape()[-1])); const int blocksPhase2 = sizeWithoutInnerDim * blocksPerRow; - if(out->type() == Type::float32) { + if (out->type() == Type::float32) { IPtr mp_ptrs = allocator->alloc>(groupLosses.size()); IPtr factorMaxes = allocator->alloc(sizeWithoutInnerDim * (groupLosses.size() - 1)); ptrInnerDimPair* dest = mp_ptrs->data>(); @@ -3131,7 +3131,7 @@ void AddFactorMaxes(Tensor out, dest, groupLosses.size(), sizeWithoutInnerDim, - std::numeric_limits::lowest()); + -std::numeric_limits::infinity()); const int sharedBytes = sizeof(float) * (groupLosses.size() + threadsPhase2); gAddFactorMaxesPhase2<<>>(out->data(), @@ -3141,13 +3141,12 @@ void AddFactorMaxes(Tensor out, groupLosses.size(), sizeWithoutInnerDim, factorMaxes->data(), groupStart, lemmaHasFactorGroupTensor->shape()[1], - numLemmas, - std::numeric_limits::lowest()); + numLemmas); allocator->free(mp_ptrs); allocator->free(factorMaxes); #if COMPILE_FP16 - } else if(out->type() == Type::float16) { + } else if (out->type() == Type::float16) { IPtr mp_ptrs = allocator->alloc>(groupLosses.size()); IPtr factorMaxes = allocator->alloc(sizeWithoutInnerDim * (groupLosses.size() - 1)); ptrInnerDimPair* dest = mp_ptrs->data>(); @@ -3164,7 +3163,7 @@ void AddFactorMaxes(Tensor out, dest, groupLosses.size(), sizeWithoutInnerDim, - __float2half(std::numeric_limits::lowest()) ); + __float2half(-std::numeric_limits::infinity()) ); const int sharedBytes = sizeof(half) * (groupLosses.size() + threadsPhase2); gAddFactorMaxesPhase2<<>>(out->data(), @@ -3174,8 +3173,7 @@ void AddFactorMaxes(Tensor out, groupLosses.size(), sizeWithoutInnerDim, factorMaxes->data(), groupStart, lemmaHasFactorGroupTensor->shape()[1], - numLemmas, - __float2half(std::numeric_limits::lowest()) ); + numLemmas); allocator->free(mp_ptrs); allocator->free(factorMaxes); #endif From edebe59a942abfd137ea66206be39ca8f20a82ad Mon Sep 17 00:00:00 2001 From: Rawn Henry Date: Fri, 11 Dec 2020 15:23:18 -0800 Subject: [PATCH 10/17] Small refactoring --- src/tensors/gpu/tensor_operators.cu | 26 +++++++++++++++++--------- 1 file changed, 17 insertions(+), 9 deletions(-) diff --git a/src/tensors/gpu/tensor_operators.cu b/src/tensors/gpu/tensor_operators.cu index f08b7ab93..9b0e499c9 100644 --- a/src/tensors/gpu/tensor_operators.cu +++ b/src/tensors/gpu/tensor_operators.cu @@ -3005,13 +3005,16 @@ __global__ void gAddFactorMaxesPhase1(T* maxes, // [#blocks, numGroups - 1] const int wid = threadIdx.x / warpSize; const int numWarps = blockDim.x / warpSize; + // Each block computes the maxes for each secondary factor within a given row. For example, if we have 6 factor groups, + // (so lemma[fg0] + 5 secondary groups) block 0 will compute the max in row 0 for fg1 to fg6. These is currently performed + // in parallel by assigning fg1 to warp 0. fg2 to warp 1 etc. If there are not enough warps to do all the factor groups then + // multiple batched iterations are needed to calculate the factor maxes. for(int warpFactorGroup = wid + 1; warpFactorGroup < numGroups; warpFactorGroup+=numWarps) { T* groupLosses = lossAndLastDimSize[warpFactorGroup].ptr; int groupLossesInnerDim = lossAndLastDimSize[warpFactorGroup].innerDim; int groupLossSize = groupLossesInnerDim * sizeWithoutInnerDim; T factorMaxima = minimal; - // Each block computes the maxes for each secondary factor within a given row. // We start by each warp computing offset to the row its block is responsible for in the loss // This is per warp since groupLossesInnerDim is warp specific and the groupLossesPtr is also warp specific. @@ -3062,9 +3065,12 @@ __global__ void gAddFactorMaxesPhase2(T* out, const int8_t* const lemmaHasFactor T* outputRowForBlockGroup = out + blockRow * numLemmas; const T* inputRowForBlockGroup = lemmaLosses + blockRow * numLemmas; + // We keep the maxes in shared memory since they are read repeatedly in the loop below. if (threadIdx.x < (numGroups - 1)) factorMaxesInBlockShared[threadIdx.x] = factorMaxesInBlock[threadIdx.x]; - + + // Sync here to ensure all factor maxes are in shared memory before other threads attempt to read the + // shared memory array. __syncthreads(); for(int rowChunk = blockStartIndexInRow; rowChunk < totalBlockChunksInRow; rowChunk += BLOCKS_PER_ROW) { @@ -3114,9 +3120,12 @@ void AddFactorMaxes(Tensor out, const int threadsPhase2 = std::min(MAX_THREADS, std::max(32, out->shape()[-1])); const int blocksPhase2 = sizeWithoutInnerDim * blocksPerRow; + IPtr mp_ptrs; + IPtr factorMaxes; + if (out->type() == Type::float32) { - IPtr mp_ptrs = allocator->alloc>(groupLosses.size()); - IPtr factorMaxes = allocator->alloc(sizeWithoutInnerDim * (groupLosses.size() - 1)); + mp_ptrs = allocator->alloc>(groupLosses.size()); + allocator->alloc(sizeWithoutInnerDim * (groupLosses.size() - 1)); ptrInnerDimPair* dest = mp_ptrs->data>(); std::vector> lossPtrs; @@ -3142,13 +3151,10 @@ void AddFactorMaxes(Tensor out, factorMaxes->data(), groupStart, lemmaHasFactorGroupTensor->shape()[1], numLemmas); - - allocator->free(mp_ptrs); - allocator->free(factorMaxes); #if COMPILE_FP16 } else if (out->type() == Type::float16) { - IPtr mp_ptrs = allocator->alloc>(groupLosses.size()); - IPtr factorMaxes = allocator->alloc(sizeWithoutInnerDim * (groupLosses.size() - 1)); + mp_ptrs = allocator->alloc>(groupLosses.size()); + factorMaxes = allocator->alloc(sizeWithoutInnerDim * (groupLosses.size() - 1)); ptrInnerDimPair* dest = mp_ptrs->data>(); std::vector> lossPtrs; @@ -3180,6 +3186,8 @@ void AddFactorMaxes(Tensor out, } else { ABORT("AddFactorMaxes not implemented for type {}", out->type()); } + allocator->free(mp_ptrs); + allocator->free(factorMaxes); } } // namespace gpu From 07de32dd27577fa97d38d7c641df116aa77a5524 Mon Sep 17 00:00:00 2001 From: Rawn Henry Date: Fri, 11 Dec 2020 17:03:13 -0800 Subject: [PATCH 11/17] Changes kernel parallelization strategy for addFactorMaxesKernel2 --- src/tensors/gpu/tensor_operators.cu | 161 +++++++++++----------------- 1 file changed, 61 insertions(+), 100 deletions(-) diff --git a/src/tensors/gpu/tensor_operators.cu b/src/tensors/gpu/tensor_operators.cu index 9b0e499c9..6e4c93697 100644 --- a/src/tensors/gpu/tensor_operators.cu +++ b/src/tensors/gpu/tensor_operators.cu @@ -2952,37 +2952,6 @@ void PoolingWithMaskingBackward(Tensor adj, lastWidth); } -// Finds max within warp. Assumes val in non-participating warps set to T::min() -// or a valid non-garbage value -template -__device__ __forceinline__ T dWarpReduceMax(T val) { - constexpr unsigned fullMask = 0xffffffff; - #pragma unroll - for (int offset = (warpSize / 2); offset > 0; offset /= 2) { - val = max(val, __shfl_down_sync(fullMask, val, offset)); - } - return val; -} - -// Elts to reduce must be <= blockDim.x. Needed since factorDims can be varied lengths. -template -__device__ __forceinline__ void dBlockReduceMax(T val, int eltsToReduce, volatile T* shared, volatile T* writeLoc) { - const int lane = threadIdx.x % warpSize; - const int wid = threadIdx.x / warpSize; - const int warpsNeeded = (eltsToReduce + warpSize - 1) / warpSize; - - if (wid < warpsNeeded) val = dWarpReduceMax(val); - if (lane==0) shared[wid]=val; - __syncthreads(); // Needed to ensure all threads write to shared before reading. - - //read from shared memory lane only if that warp existed. Otherwise just get first lane. - val = (threadIdx.x < warpsNeeded) ? shared[lane] : shared[0]; - if(wid==0) { - T temp = dWarpReduceMax(val); - if(lane == 0) writeLoc[0] = temp; - } -} - template struct ptrInnerDimPair { T* ptr; @@ -3036,7 +3005,7 @@ __global__ void gAddFactorMaxesPhase1(T* maxes, // [#blocks, numGroups - 1] } - template + template __global__ void gAddFactorMaxesPhase2(T* out, const int8_t* const lemmaHasFactorGroup, const IndexType* const indices, const T* lemmaLosses, @@ -3048,54 +3017,49 @@ __global__ void gAddFactorMaxesPhase2(T* out, const int8_t* const lemmaHasFactor size_t numLemmas) { extern __shared__ uint8_t _sharedBytes[]; - T* sums = (T*)_sharedBytes; - T* factorMaxesInBlockShared = sums + blockDim.x; - - // A group of blocks is assigned a row to process. We compute some the row index for each block in the group - // and the row each group of blocks is responsible for processing. - const int blockStartIndexInRow = blockIdx.x % BLOCKS_PER_ROW; - const int blockRow = blockIdx.x / BLOCKS_PER_ROW; - - // We need to know the total number of chunks in a row to safely iterate over the input and output - const int totalBlockChunksInRow = (numLemmas + blockDim.x - 1) / blockDim.x; - - // Calculate the row pointer for the global memory which holds for max for each factor. This is the same for - // all blocks in a given group - const T* factorMaxesInBlock = maxes + blockRow * (numGroups - 1); - T* outputRowForBlockGroup = out + blockRow * numLemmas; - const T* inputRowForBlockGroup = lemmaLosses + blockRow * numLemmas; - - // We keep the maxes in shared memory since they are read repeatedly in the loop below. - if (threadIdx.x < (numGroups - 1)) - factorMaxesInBlockShared[threadIdx.x] = factorMaxesInBlock[threadIdx.x]; - + T* factorMaxesInBlockShared = (T*)_sharedBytes; + + // Each thread is given an element in the output to construct. First we compute every thread's offset from the + // base pointer along with the row and column indicies for the output cell the thread is responsible for constructing. + const int offsetFromBasePtr = blockIdx.x * blockDim.x + threadIdx.x; + const int rowForThread = offsetFromBasePtr / numLemmas; + const int colForThread = offsetFromBasePtr % numLemmas; + + // We keep the maxes in shared memory since they are read repeatedly by each thread in a block + const T* globalFactorMaxesRow = maxes + rowForThread * (numGroups - 1); + for(int factorGroup = threadIdx.x; factorGroup < (numGroups - 1); factorGroup += blockDim.x) { + factorMaxesInBlockShared[factorGroup] = globalFactorMaxesRow[factorGroup]; + } + // Sync here to ensure all factor maxes are in shared memory before other threads attempt to read the // shared memory array. __syncthreads(); - for(int rowChunk = blockStartIndexInRow; rowChunk < totalBlockChunksInRow; rowChunk += BLOCKS_PER_ROW) { - // Compute the start of the chunk each block needs to read from in the input and output tensors - const T* lemmaLossesChunk = inputRowForBlockGroup + rowChunk * blockDim.x; - T* blockOutputChunk = outputRowForBlockGroup + rowChunk * blockDim.x; - - // Calculate the column being read to/written from to use as a read/write guard - const int columnIndex = threadIdx.x * rowChunk; - if (columnIndex < numLemmas) - sums[threadIdx.x] = lemmaLossesChunk[threadIdx.x]; - - // Each block accumulates its output in shared memory then writes the final result to global. We use __ldg since the - // data being read does not change throughout the life of the kernel so the constant cache can be used. - for(int g = 1; g < (int)numGroups; ++g) { - int lemma = indices? indices[columnIndex] - groupStart : columnIndex; - T factorMask = static_cast(__ldg(&lemmaHasFactorGroup[lemma * lemmaHasFactorGroupWidth + g])); - T factorMaxima = factorMaxesInBlockShared[g-1]; - - if (columnIndex < numLemmas) - sums[threadIdx.x] += (factorMask * factorMaxima); - } - if (columnIndex < numLemmas) - blockOutputChunk[threadIdx.x] = sums[threadIdx.x]; + // Compute a write guard for the output array + const int outputSize = numLemmas * sizeWithoutInnerDim; + + // Now we start accumulating the thread's sum by initializing the the lemma value. The out tensor and lemmaLosses + // tensor are the same size. + T threadSum = 0; + if (offsetFromBasePtr < outputSize) + threadSum = lemmaLosses[offsetFromBasePtr]; + + // Each block accumulates its output locally then writes the final result to global. We use __ldg since the + // data being read does not change throughout the life of the kernel so the constant cache can be used. + // This is more of a compiler hint than a necessity. + for(int g = 1; g < (int)numGroups; ++g) { + int lemma = indices? indices[colForThread] - groupStart : colForThread; + T factorMask = static_cast(__ldg(&lemmaHasFactorGroup[lemma * lemmaHasFactorGroupWidth + g])); + T factorMaxima = factorMaxesInBlockShared[g-1]; + + if (offsetFromBasePtr < outputSize) + threadSum += (factorMask * factorMaxima); } + + // Finally, the accumulated sum is written back out to global memory. + if (offsetFromBasePtr < outputSize) + out[offsetFromBasePtr] = threadSum; + } void AddFactorMaxes(Tensor out, @@ -3116,16 +3080,15 @@ void AddFactorMaxes(Tensor out, const int blocksPhase1 = sizeWithoutInnerDim; const int threadsPhase1 = std::min(MAX_THREADS, (int)(warpSize * (groupLosses.size() - 1)) ); - constexpr int blocksPerRow = 8; - const int threadsPhase2 = std::min(MAX_THREADS, std::max(32, out->shape()[-1])); - const int blocksPhase2 = sizeWithoutInnerDim * blocksPerRow; + const int totalPhase2Threads = out->shape().elements(); + const int threadsPhase2 = std::min(MAX_THREADS, out->shape()[-1]); + const int blocksPhase2 = (totalPhase2Threads + threadsPhase2 - 1) / threadsPhase2; IPtr mp_ptrs; IPtr factorMaxes; - if (out->type() == Type::float32) { mp_ptrs = allocator->alloc>(groupLosses.size()); - allocator->alloc(sizeWithoutInnerDim * (groupLosses.size() - 1)); + factorMaxes = allocator->alloc(sizeWithoutInnerDim * (groupLosses.size() - 1)); ptrInnerDimPair* dest = mp_ptrs->data>(); std::vector> lossPtrs; @@ -3142,15 +3105,15 @@ void AddFactorMaxes(Tensor out, sizeWithoutInnerDim, -std::numeric_limits::infinity()); - const int sharedBytes = sizeof(float) * (groupLosses.size() + threadsPhase2); - gAddFactorMaxesPhase2<<>>(out->data(), - lemmaHasFactorGroupTensor->data(), - indices? indices->data(): nullptr, - groupLosses[0]->data(), - groupLosses.size(), sizeWithoutInnerDim, - factorMaxes->data(), groupStart, - lemmaHasFactorGroupTensor->shape()[1], - numLemmas); + const int sharedBytes = sizeof(float) * (groupLosses.size()); + gAddFactorMaxesPhase2<<>>(out->data(), + lemmaHasFactorGroupTensor->data(), + indices? indices->data(): nullptr, + groupLosses[0]->data(), + groupLosses.size(), sizeWithoutInnerDim, + factorMaxes->data(), groupStart, + lemmaHasFactorGroupTensor->shape()[1], + numLemmas); #if COMPILE_FP16 } else if (out->type() == Type::float16) { mp_ptrs = allocator->alloc>(groupLosses.size()); @@ -3171,17 +3134,15 @@ void AddFactorMaxes(Tensor out, sizeWithoutInnerDim, __float2half(-std::numeric_limits::infinity()) ); - const int sharedBytes = sizeof(half) * (groupLosses.size() + threadsPhase2); - gAddFactorMaxesPhase2<<>>(out->data(), - lemmaHasFactorGroupTensor->data(), - indices? indices->data(): nullptr, - groupLosses[0]->data(), - groupLosses.size(), sizeWithoutInnerDim, - factorMaxes->data(), groupStart, - lemmaHasFactorGroupTensor->shape()[1], - numLemmas); - allocator->free(mp_ptrs); - allocator->free(factorMaxes); + const int sharedBytes = sizeof(half) * (groupLosses.size()); + gAddFactorMaxesPhase2<<>>(out->data(), + lemmaHasFactorGroupTensor->data(), + indices? indices->data(): nullptr, + groupLosses[0]->data(), + groupLosses.size(), sizeWithoutInnerDim, + factorMaxes->data(), groupStart, + lemmaHasFactorGroupTensor->shape()[1], + numLemmas); #endif } else { ABORT("AddFactorMaxes not implemented for type {}", out->type()); From eb49d6c3c2355e5c38a493d973af441c59938932 Mon Sep 17 00:00:00 2001 From: Rawn Henry Date: Fri, 11 Dec 2020 17:44:46 -0800 Subject: [PATCH 12/17] More refactorings --- src/tensors/gpu/tensor_operators.cu | 34 ++++++++++++++--------------- 1 file changed, 16 insertions(+), 18 deletions(-) diff --git a/src/tensors/gpu/tensor_operators.cu b/src/tensors/gpu/tensor_operators.cu index 6e4c93697..b10baceb5 100644 --- a/src/tensors/gpu/tensor_operators.cu +++ b/src/tensors/gpu/tensor_operators.cu @@ -3038,27 +3038,25 @@ __global__ void gAddFactorMaxesPhase2(T* out, const int8_t* const lemmaHasFactor // Compute a write guard for the output array const int outputSize = numLemmas * sizeWithoutInnerDim; - // Now we start accumulating the thread's sum by initializing the the lemma value. The out tensor and lemmaLosses - // tensor are the same size. - T threadSum = 0; - if (offsetFromBasePtr < outputSize) - threadSum = lemmaLosses[offsetFromBasePtr]; - - // Each block accumulates its output locally then writes the final result to global. We use __ldg since the - // data being read does not change throughout the life of the kernel so the constant cache can be used. - // This is more of a compiler hint than a necessity. - for(int g = 1; g < (int)numGroups; ++g) { - int lemma = indices? indices[colForThread] - groupStart : colForThread; - T factorMask = static_cast(__ldg(&lemmaHasFactorGroup[lemma * lemmaHasFactorGroupWidth + g])); - T factorMaxima = factorMaxesInBlockShared[g-1]; - - if (offsetFromBasePtr < outputSize) + // Now we start accumulating the thread's sum by initializing to the lemma value. The out tensor and lemmaLosses + // tensor are the same size so the output guard is used to prevent invalid reads. + if (offsetFromBasePtr < outputSize) { + T threadSum = lemmaLosses[offsetFromBasePtr]; + + // Each block accumulates its output locally then writes the final result to global. We use __ldg since the + // data being read does not change throughout the life of the kernel so the constant cache can be used. + // This is more of a compiler hint than a necessity. + for(int g = 1; g < (int)numGroups; ++g) { + int lemma = indices? indices[colForThread] - groupStart : colForThread; + T factorMask = static_cast(__ldg(&lemmaHasFactorGroup[lemma * lemmaHasFactorGroupWidth + g])); + T factorMaxima = factorMaxesInBlockShared[g-1]; + threadSum += (factorMask * factorMaxima); - } + } - // Finally, the accumulated sum is written back out to global memory. - if (offsetFromBasePtr < outputSize) + // Finally, the accumulated sum is written back out to global memory. out[offsetFromBasePtr] = threadSum; + } } From 975682b41080cbd6e25ab7b51ae5940c08f33c2b Mon Sep 17 00:00:00 2001 From: Rawn Henry Date: Fri, 11 Dec 2020 19:18:55 -0800 Subject: [PATCH 13/17] Fixes bug in factor maxes --- src/tensors/gpu/tensor_operators.cu | 47 +++++++++++------------------ 1 file changed, 17 insertions(+), 30 deletions(-) diff --git a/src/tensors/gpu/tensor_operators.cu b/src/tensors/gpu/tensor_operators.cu index b10baceb5..76c2e9438 100644 --- a/src/tensors/gpu/tensor_operators.cu +++ b/src/tensors/gpu/tensor_operators.cu @@ -3016,9 +3016,6 @@ __global__ void gAddFactorMaxesPhase2(T* out, const int8_t* const lemmaHasFactor int lemmaHasFactorGroupWidth, size_t numLemmas) { - extern __shared__ uint8_t _sharedBytes[]; - T* factorMaxesInBlockShared = (T*)_sharedBytes; - // Each thread is given an element in the output to construct. First we compute every thread's offset from the // base pointer along with the row and column indicies for the output cell the thread is responsible for constructing. const int offsetFromBasePtr = blockIdx.x * blockDim.x + threadIdx.x; @@ -3027,13 +3024,6 @@ __global__ void gAddFactorMaxesPhase2(T* out, const int8_t* const lemmaHasFactor // We keep the maxes in shared memory since they are read repeatedly by each thread in a block const T* globalFactorMaxesRow = maxes + rowForThread * (numGroups - 1); - for(int factorGroup = threadIdx.x; factorGroup < (numGroups - 1); factorGroup += blockDim.x) { - factorMaxesInBlockShared[factorGroup] = globalFactorMaxesRow[factorGroup]; - } - - // Sync here to ensure all factor maxes are in shared memory before other threads attempt to read the - // shared memory array. - __syncthreads(); // Compute a write guard for the output array const int outputSize = numLemmas * sizeWithoutInnerDim; @@ -3049,8 +3039,7 @@ __global__ void gAddFactorMaxesPhase2(T* out, const int8_t* const lemmaHasFactor for(int g = 1; g < (int)numGroups; ++g) { int lemma = indices? indices[colForThread] - groupStart : colForThread; T factorMask = static_cast(__ldg(&lemmaHasFactorGroup[lemma * lemmaHasFactorGroupWidth + g])); - T factorMaxima = factorMaxesInBlockShared[g-1]; - + T factorMaxima = globalFactorMaxesRow[g-1]; threadSum += (factorMask * factorMaxima); } @@ -3103,15 +3092,14 @@ void AddFactorMaxes(Tensor out, sizeWithoutInnerDim, -std::numeric_limits::infinity()); - const int sharedBytes = sizeof(float) * (groupLosses.size()); - gAddFactorMaxesPhase2<<>>(out->data(), - lemmaHasFactorGroupTensor->data(), - indices? indices->data(): nullptr, - groupLosses[0]->data(), - groupLosses.size(), sizeWithoutInnerDim, - factorMaxes->data(), groupStart, - lemmaHasFactorGroupTensor->shape()[1], - numLemmas); + gAddFactorMaxesPhase2<<>>(out->data(), + lemmaHasFactorGroupTensor->data(), + indices? indices->data(): nullptr, + groupLosses[0]->data(), + groupLosses.size(), sizeWithoutInnerDim, + factorMaxes->data(), groupStart, + lemmaHasFactorGroupTensor->shape()[1], + numLemmas); #if COMPILE_FP16 } else if (out->type() == Type::float16) { mp_ptrs = allocator->alloc>(groupLosses.size()); @@ -3132,15 +3120,14 @@ void AddFactorMaxes(Tensor out, sizeWithoutInnerDim, __float2half(-std::numeric_limits::infinity()) ); - const int sharedBytes = sizeof(half) * (groupLosses.size()); - gAddFactorMaxesPhase2<<>>(out->data(), - lemmaHasFactorGroupTensor->data(), - indices? indices->data(): nullptr, - groupLosses[0]->data(), - groupLosses.size(), sizeWithoutInnerDim, - factorMaxes->data(), groupStart, - lemmaHasFactorGroupTensor->shape()[1], - numLemmas); + gAddFactorMaxesPhase2<<>>(out->data(), + lemmaHasFactorGroupTensor->data(), + indices? indices->data(): nullptr, + groupLosses[0]->data(), + groupLosses.size(), sizeWithoutInnerDim, + factorMaxes->data(), groupStart, + lemmaHasFactorGroupTensor->shape()[1], + numLemmas); #endif } else { ABORT("AddFactorMaxes not implemented for type {}", out->type()); From f0ffee14e78f9fa1719643dd978e5fcd7a2295b2 Mon Sep 17 00:00:00 2001 From: Rawn Henry Date: Sun, 13 Dec 2020 13:33:06 -0800 Subject: [PATCH 14/17] Refactor of addFactorMaxes entrance --- src/graph/expression_operators.cpp | 8 +++++--- src/graph/expression_operators.h | 2 +- src/layers/generic.cpp | 31 +++++++++++++++--------------- src/layers/generic.h | 1 + 4 files changed, 23 insertions(+), 19 deletions(-) diff --git a/src/graph/expression_operators.cpp b/src/graph/expression_operators.cpp index 54f438228..f6e639aad 100644 --- a/src/graph/expression_operators.cpp +++ b/src/graph/expression_operators.cpp @@ -683,20 +683,22 @@ Expr unlikelihood(Expr logits, Expr indices) { return -log(gather(1.f - softmax(logits), /*axis=*/-1, indicesWithLayout)); } -Expr addFactorMaxes(Expr lemmaHasFactorGroup, std::vector groupLosses, Expr hypIndices, size_t groupStart, size_t numLemmas) { +Expr addFactorMaxes(Expr lemmaHasFactorGroup, std::vector groupLosses, Expr hypIndices, size_t group0Start) { if(groupLosses.size() == 1) { return groupLosses[0]; } + + int numLemmas = groupLosses[0]->shape()[-1]; + std::vector nodes({lemmaHasFactorGroup}); if (hypIndices) { nodes.push_back(hypIndices); } nodes.insert(nodes.end(), groupLosses.begin(), groupLosses.end()); bool hasShortList = hypIndices != nullptr; - return Expression(nodes, hasShortList, groupStart, numLemmas); + return Expression(nodes, hasShortList, group0Start, numLemmas); } - Expr plus(const std::vector& nodes) { ABORT_IF(nodes.size() > 1, "Not implemented"); return nodes[0]; diff --git a/src/graph/expression_operators.h b/src/graph/expression_operators.h index 72a6aba99..945ad671c 100755 --- a/src/graph/expression_operators.h +++ b/src/graph/expression_operators.h @@ -177,7 +177,7 @@ Expr atleast_3d(Expr a); Expr atleast_4d(Expr a); Expr atleast_nd(Expr a, size_t dims); -Expr addFactorMaxes(Expr lemmaHasFactorGroup, std::vector groupLosses, Expr hypIndices, size_t groupStart, size_t numLemmas); +Expr addFactorMaxes(Expr lemmaHasFactorGroup, std::vector groupLosses, Expr hypIndices, size_t group0Start); // create a constant of shape a->shape() and initialize with init // @TODO: add a && version, to avoid a ref count. NodeInitializers are typically temps. diff --git a/src/layers/generic.cpp b/src/layers/generic.cpp index 84f0350fe..3820d89e5 100755 --- a/src/layers/generic.cpp +++ b/src/layers/generic.cpp @@ -3,6 +3,7 @@ * SPDX-License-Identifier: MIT */ +#include #include "marian.h" #include "layers/generic.h" @@ -85,8 +86,13 @@ namespace marian { } else { auto numGroups = getNumFactorGroups(); if(numGroups > 1 && graph()->isInference() && graph()->getBackend()->getDeviceId().type == DeviceType::gpu) { - Expr shortlistExpr = shortlist? constant(shortlist->indices()) : nullptr; - sel = addFactorMaxesHelper(shortlistExpr); + Expr shortlistIndices = shortlist? constant(shortlist->indices()) : nullptr; + Expr lemmaHasFactorGroupTensor = getLemmaHasFactorGroupTensor(); + std::vector groupLosses(logits_.size()); + std::transform(logits_.begin(), logits_.end(), groupLosses.begin(), [](const Ptr& loss) -> Expr {return loss->loss();}); + + sel = addFactorMaxes(lemmaHasFactorGroupTensor, groupLosses, + shortlistIndices, factoredVocab_->getGroupRange(0).first); } else { for (size_t g = 1; g < numGroups; g++) { auto factorMaxima = max(logits_[g]->loss(), -1); @@ -196,11 +202,12 @@ namespace marian { return res; } - Expr Logits::addFactorMaxesHelper(Expr indices) const { - auto g = graph(); - const std::string name = "lemmaHasFactorGroup"; - auto lemmaHasFactorGroupTensor = g->findByName(name); - if(!lemmaHasFactorGroupTensor) { + Expr Logits::getLemmaHasFactorGroupTensor() const { + auto g = graph(); + const std::string name = "lemmaHasFactorGroup"; + auto lemmaHasFactorGroupTensor = g->findByName(name); + + if(!lemmaHasFactorGroupTensor) { // We want to make this a graph param so the GPU can use this to implement lemmaHasFactorGroup to avoid unneeded memcpys. const auto lemmaHasFactorGroup = factoredVocab_->getLemmaHasFactorGroupVector(); int dimLemma = (int)lemmaHasFactorGroup.size(); @@ -228,14 +235,8 @@ namespace marian { lemmaHasFactorGroupTensor->setMemoize(true); g->rememberByName(name, lemmaHasFactorGroupTensor); } - size_t n = indices ? indices->shape().elements() : (factoredVocab_->getGroupRange(0).second - factoredVocab_->getGroupRange(0).first); - - std::vector groupLosses(getNumFactorGroups()); - for(int g = 0; g < getNumFactorGroups(); ++g) { - groupLosses[g] = logits_[g]->loss(); - } - - return addFactorMaxes(lemmaHasFactorGroupTensor, groupLosses, indices, factoredVocab_->getGroupRange(0).first, n); + ABORT_IF(!lemmaHasFactorGroupTensor, "Lemma has factor group tensor undefined?"); + return lemmaHasFactorGroupTensor; } Logits Logits::applyUnaryFunction(const std::function& f) const { // clone this but apply f to all loss values diff --git a/src/layers/generic.h b/src/layers/generic.h index 8bb4962f5..2d344a3c7 100755 --- a/src/layers/generic.h +++ b/src/layers/generic.h @@ -147,6 +147,7 @@ class Logits { template Expr constant(const std::vector& data) const { return constant(Shape{(int)data.size()}, data); } // same as constant() but assuming vector Expr indices(const std::vector& data) const { return graph()->indices(data); } // actually the same as constant(data) for this data type std::vector getFactorMasks(size_t factorGroup, const std::vector& indices) const; + Expr getLemmaHasFactorGroupTensor() const; Expr addFactorMaxesHelper(Expr indices=nullptr) const; private: // members From 4f57fdb4594e5cf86ba6fc7f10882f0400716fb0 Mon Sep 17 00:00:00 2001 From: Rawn Henry Date: Mon, 14 Dec 2020 12:08:10 -0800 Subject: [PATCH 15/17] Updates changelog --- CHANGELOG.md | 1 + 1 file changed, 1 insertion(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index 4182d72b6..182f190ff 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -9,6 +9,7 @@ and this project adheres to [Semantic Versioning](http://semver.org/spec/v2.0.0. ## [Unreleased] ### Added +- Add new operator to compute path score for lemmas for a factored vocabulary - 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 de5449bd049521ca29788cd276f91633656637aa Mon Sep 17 00:00:00 2001 From: Rawn Henry Date: Tue, 20 Oct 2020 09:42:38 -0700 Subject: [PATCH 16/17] 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 3c0150587..36192aa7a 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 7be492aa446e22da5409c6a2c857fa79765ecf92 Mon Sep 17 00:00:00 2001 From: Rawn Henry Date: Thu, 8 Jul 2021 17:20:48 -0700 Subject: [PATCH 17/17] Fix factor map op for shortlist --- src/graph/node_operators_binary.h | 4 ++-- src/layers/generic.cpp | 2 +- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/src/graph/node_operators_binary.h b/src/graph/node_operators_binary.h index 0d79986b1..474ead783 100644 --- a/src/graph/node_operators_binary.h +++ b/src/graph/node_operators_binary.h @@ -1291,9 +1291,9 @@ size_t numLemmas_; bool hasShortlist_; public: AddFactorMaxesOp(const std::vector& nodes, bool hasShortlist, size_t groupStart, size_t numLemmas) - : NaryNodeOp(nodes, getShape(nodes, hasShortlist), commonType(std::vector(nodes.begin() + 1, nodes.end())) ) { + : NaryNodeOp(nodes, getShape(nodes, hasShortlist), commonType(std::vector(nodes.begin() + 1 + (int)hasShortlist, nodes.end())) ) { groupStart_ = groupStart; - numLemmas_ = hasShortlist? nodes[1]->shape().size(): numLemmas; + numLemmas_ = numLemmas; hasShortlist_ = hasShortlist; } diff --git a/src/layers/generic.cpp b/src/layers/generic.cpp index c807dc060..5c9a9ccdb 100755 --- a/src/layers/generic.cpp +++ b/src/layers/generic.cpp @@ -89,7 +89,7 @@ namespace marian { } else { auto numGroups = getNumFactorGroups(); if(numGroups > 1 && graph()->isInference() && graph()->getBackend()->getDeviceId().type == DeviceType::gpu) { - Expr shortlistIndices = shortlist? constant(shortlist->indices()) : nullptr; + Expr shortlistIndices = shortlist? indices(shortlist->indices()) : nullptr; Expr lemmaHasFactorGroupTensor = getLemmaHasFactorGroupTensor(); std::vector groupLosses(logits_.size()); std::transform(logits_.begin(), logits_.end(), groupLosses.begin(), [](const Ptr& loss) -> Expr {return loss->loss();});