From df03c9fd97b8cec2671769a94ea9c9a01dfc1274 Mon Sep 17 00:00:00 2001 From: Jordan Fix Date: Wed, 10 Feb 2021 11:42:14 -0800 Subject: [PATCH] Improvements for FX tracer Differential Revision: D26355583 fbshipit-source-id: 90bf0b0bf3aedc2b06006979d238c0fabd1ee036 --- lib/Backends/NNPI/FXIRImporter.cpp | 333 +++++++++++-------- lib/Backends/NNPI/FXIRImporter.h | 23 +- lib/Backends/NNPI/FXNNPICompiledFunction.cpp | 2 +- 3 files changed, 216 insertions(+), 142 deletions(-) diff --git a/lib/Backends/NNPI/FXIRImporter.cpp b/lib/Backends/NNPI/FXIRImporter.cpp index f69ffe1b13..687ca75abf 100644 --- a/lib/Backends/NNPI/FXIRImporter.cpp +++ b/lib/Backends/NNPI/FXIRImporter.cpp @@ -9,12 +9,24 @@ using namespace utils; namespace { /// \returns true if \p opCode is not either "placeholder" or "output" (which -/// indicate the node is an operater). -bool isOps(const std::string &opCode) { +/// indicate the node is an operator). +static bool isOps(const std::string &opCode) { return opCode != "placeholder" && opCode != "output"; } -// Node Importers +/// \returns the name of some input \p node. Fatals if \p node is not specified +/// as is_node. If \p optional then \returns an empty string if \p node is null. +static const std::string &getInputNodeName(const folly::dynamic &node, + bool optional = false) { + if (optional && node.isNull()) { + static const std::string empty; + return empty; + } + CHECK(node["is_node"].asBool()) << "Expected is_node"; + return node["name"].getString(); +} + +/// Node Importers template class BinaryEltwiseNodeImporter : public INNPIFXNodeImporter { public: @@ -28,9 +40,9 @@ class BinaryEltwiseNodeImporter : public INNPIFXNodeImporter { // TODO: broadcast inputs if input is not a node. std::array inputNames; snprintf(inputNames[0], NNPI_MAX_STRING_LEN, "%s", - inputs[0]["name"].getString().c_str()); + getInputNodeName(inputs[0]).c_str()); snprintf(inputNames[1], NNPI_MAX_STRING_LEN, "%s", - inputs[1]["name"].getString().c_str()); + getInputNodeName(inputs[1]).c_str()); importer.setUsedTensors({inputNames[0], inputNames[1]}, {name}); return nnpiNetworkAddElementwiseOp(importer.getNetwork(), name.c_str(), @@ -41,48 +53,55 @@ class BinaryEltwiseNodeImporter : public INNPIFXNodeImporter { class LinearNodeImporter : public INNPIFXNodeImporter { public: - NNPIErrorCode importNode(const folly::dynamic &node, - const std::function &getQualName, - FXNNPIImporter &importer) override { - const auto &inputs = node["args"]; + NNPIErrorCode + importNode(const folly::dynamic &node, + const std::function & /* getQualName */, + FXNNPIImporter &importer) override { const auto &name = node["name"].getString(); - const auto &weightName = getQualName("weight"); - const auto &biasName = getQualName("bias"); - - importer.setUsedTensors( - {inputs[0]["name"].getString(), weightName, biasName}, {name}); - return nnpiNetworkAddFullyConnectedOp( - importer.getNetwork(), name.c_str(), - inputs[0]["name"].getString().c_str(), name.c_str(), weightName.c_str(), - importer.getConstant(biasName) ? biasName.c_str() : nullptr); + const auto &kwargs = node["kwargs"]; + const auto &inputName = getInputNodeName(kwargs["input"]); + const auto &weightName = getInputNodeName(kwargs["weight"]); + const auto &biasName = + getInputNodeName(kwargs["bias"], /* optional */ true); + + std::unordered_set readTensors = {inputName, weightName}; + if (!biasName.empty()) { + readTensors.insert(biasName); + } + importer.setUsedTensors(readTensors, {name}); + return nnpiNetworkAddFullyConnectedOp(importer.getNetwork(), name.c_str(), + inputName.c_str(), name.c_str(), + importer.getConstantName(weightName), + importer.getConstantName(biasName)); } }; template class ConvolutionNodeImporter : public INNPIFXNodeImporter { public: - NNPIErrorCode importNode(const folly::dynamic &node, - const std::function &getQualName, - FXNNPIImporter &importer) override { + NNPIErrorCode + importNode(const folly::dynamic &node, + const std::function & /* getQualName */, + FXNNPIImporter &importer) override { const auto &inputs = node["args"]; const auto &name = node["name"].getString(); - const auto &filterName = getQualName("weight"); - const auto &biasName = getQualName("bias"); + const auto &inputName = getInputNodeName(inputs[0]); + const auto &filterName = getInputNodeName(inputs[1]); + const auto &biasName = getInputNodeName(inputs[2], /* optional */ true); + + // Kernel size is implicit in the shape of the filter. + const auto &filterDesc = importer.getTensorDesc(filterName); + LOG_AND_RETURN_IF_NOT(ERROR, filterDesc.numDims == convDims + 2, + "Unexpected filter dims", NNPI_INVALID_PARAM); + std::vector kernelSize; + for (size_t i = 0; i < convDims; i++) { + kernelSize.push_back(filterDesc.dims[i + 2]); + } - // Get parameters. - auto kernelSize = - toIntegerArray(node["parameters"]["kernel_size"].getString(), - /* length */ convDims); - auto stride = toIntegerArray( - node["parameters"]["stride"].getString(), /* length */ convDims); - auto padding = - toIntegerArray(node["parameters"]["padding"].getString(), - /* length */ convDims); - auto dilation = - toIntegerArray(node["parameters"]["dilation"].getString(), - /* length */ convDims); - const auto &groups = node["parameters"]["groups"].asInt(); - const auto &paddingMode = node["parameters"]["padding_mode"].getString(); + auto stride = toIntegerArray(inputs[3], /* length */ convDims); + auto padding = toIntegerArray(inputs[4], /* length */ convDims); + auto dilation = toIntegerArray(inputs[5], /* length */ convDims); + const auto &groups = inputs[6].asInt(); LOG_AND_RETURN_IF_NOT(ERROR, std::all_of(dilation.cbegin(), dilation.cend(), @@ -90,18 +109,17 @@ class ConvolutionNodeImporter : public INNPIFXNodeImporter { "Dilation is not supported", NNPI_INVALID_PARAM); LOG_AND_RETURN_IF_NOT(ERROR, groups == 1, "Group is not supported", NNPI_INVALID_PARAM); - LOG_AND_RETURN_IF_NOT(ERROR, paddingMode == "zeros", - "Only support zeros padding mode", - NNPI_INVALID_PARAM); - importer.setUsedTensors( - {inputs[0]["name"].getString(), filterName, biasName}, {name}); + std::unordered_set readTensors = {inputName, filterName}; + if (!biasName.empty()) { + readTensors.insert(biasName); + } + importer.setUsedTensors(readTensors, {name}); return nnpiNetworkAddConvolutionOp( - importer.getNetwork(), name.c_str(), - inputs[0]["name"].getString().c_str(), name.c_str(), filterName.c_str(), - importer.getConstant(biasName) ? biasName.c_str() : nullptr, - kernelSize.data(), padding.data(), padding.data(), stride.data(), - dilation.data(), convDims, groups); + importer.getNetwork(), name.c_str(), inputName.c_str(), name.c_str(), + importer.getConstantName(filterName), + importer.getConstantName(biasName), kernelSize.data(), padding.data(), + padding.data(), stride.data(), dilation.data(), convDims, groups); } }; @@ -110,8 +128,9 @@ class BatchNormalizationNodeImporter : public INNPIFXNodeImporter { NNPIErrorCode importNode(const folly::dynamic &node, const std::function &getQualName, FXNNPIImporter &importer) override { - const auto &inputs = node["args"]; const auto &name = node["name"].getString(); + const auto &kwargs = node["kwargs"]; + const auto &inputName = getInputNodeName(kwargs["input"]); const auto &weightName = getQualName("weight"); const auto &biasName = getQualName("bias"); const auto &meanName = getQualName("running_mean"); @@ -120,13 +139,12 @@ class BatchNormalizationNodeImporter : public INNPIFXNodeImporter { // Get parameters. const auto &eps = node["parameters"]["eps"].asDouble(); - importer.setUsedTensors({inputs[0]["name"].getString(), weightName, - biasName, meanName, varName}, - {name}); - return nnpiNetworkAddBatchNormOp( - importer.getNetwork(), name.c_str(), - inputs[0]["name"].getString().c_str(), name.c_str(), meanName.c_str(), - varName.c_str(), weightName.c_str(), biasName.c_str(), eps); + importer.setUsedTensors( + {inputName, weightName, biasName, meanName, varName}, {name}); + return nnpiNetworkAddBatchNormOp(importer.getNetwork(), name.c_str(), + inputName.c_str(), name.c_str(), + meanName.c_str(), varName.c_str(), + weightName.c_str(), biasName.c_str(), eps); } }; @@ -136,18 +154,18 @@ class ReluNodeImporter : public INNPIFXNodeImporter { importNode(const folly::dynamic &node, const std::function & /* getQualName */, FXNNPIImporter &importer) override { - const auto &inputs = node["args"]; const auto &name = node["name"].getString(); + const auto &kwargs = node["kwargs"]; + const auto &inputName = getInputNodeName(kwargs["input"]); // Get parameters. - const auto &inplace_ = node["parameters"]["inplace"].asBool(); + const auto &inplace = kwargs["inplace"].asBool(); // TODO: replace users of ReLU input after ReLU with ReLU output. - LOG_IF_NOT(WARNING, !inplace_) << "Inplace ReLU is not supported"; + LOG_IF_NOT(WARNING, !inplace) << "Inplace ReLU is not supported"; - importer.setUsedTensors({inputs[0]["name"].getString()}, {name}); + importer.setUsedTensors({inputName}, {name}); return nnpiNetworkAddReluOp(importer.getNetwork(), name.c_str(), - inputs[0]["name"].getString().c_str(), - name.c_str()); + inputName.c_str(), name.c_str()); } }; @@ -158,13 +176,14 @@ class AdaptivePoolNodeImporter : public INNPIFXNodeImporter { importNode(const folly::dynamic &node, const std::function & /* getQualName */, FXNNPIImporter &importer) override { - const auto &inputs = node["args"]; const auto &name = node["name"].getString(); + const auto &kwargs = node["kwargs"]; + const auto &inputName = getInputNodeName(kwargs["input"]); - importer.setUsedTensors({inputs[0]["name"].getString()}, {name}); - return nnpiNetworkAddAdaptivePoolingOp( - importer.getNetwork(), name.c_str(), - inputs[0]["name"].getString().c_str(), name.c_str(), poolType); + importer.setUsedTensors({inputName}, {name}); + return nnpiNetworkAddAdaptivePoolingOp(importer.getNetwork(), name.c_str(), + inputName.c_str(), name.c_str(), + poolType); } }; @@ -175,29 +194,27 @@ class PoolNodeImporter : public INNPIFXNodeImporter { importNode(const folly::dynamic &node, const std::function & /* getQualName */, FXNNPIImporter &importer) override { - const auto &inputs = node["args"]; const auto &name = node["name"].getString(); + const auto &kwargs = node["kwargs"]; + const auto &inputName = getInputNodeName(kwargs["input"]); // Get parameters - auto kernelSize = - toIntegerArray(node["parameters"]["kernel_size"].getString(), - /* length */ poolDims); - auto stride = toIntegerArray( - node["parameters"]["stride"].getString(), /* length */ poolDims); - auto padding = - toIntegerArray(node["parameters"]["padding"].getString(), - /* length */ poolDims); - const auto &ceilMode = node["parameters"]["ceil_mode"].asBool(); + auto kernelSize = toIntegerArray(kwargs["kernel_size"], + /* length */ poolDims); + auto stride = toIntegerArray(kwargs["stride"], + /* length */ poolDims); + auto padding = toIntegerArray(kwargs["padding"], + /* length */ poolDims); + const auto &ceilMode = kwargs["ceil_mode"].asBool(); std::vector dilation(poolDims, 1); auto returnIndices = false; bool countIncludePads = true; if (poolType == NNPI_POOL_AVG) { - countIncludePads = node["parameters"]["count_include_pad"].asBool(); + countIncludePads = kwargs["count_include_pad"].asBool(); } if (poolType == NNPI_POOL_MAX) { - returnIndices = node["parameters"]["return_indices"].asBool(); - dilation = - toIntegerArray(node["parameters"]["dilation"].getString()); + returnIndices = kwargs["return_indices"].asBool(); + dilation = toIntegerArray(kwargs["dilation"]); } LOG_AND_RETURN_IF_NOT(ERROR, @@ -208,10 +225,9 @@ class PoolNodeImporter : public INNPIFXNodeImporter { "Return_indices is not supported", NNPI_INVALID_PARAM); - importer.setUsedTensors({inputs[0]["name"].getString()}, {name}); + importer.setUsedTensors({inputName}, {name}); return nnpiNetworkAddPoolingOp( - importer.getNetwork(), name.c_str(), - inputs[0]["name"].getString().c_str(), name.c_str(), + importer.getNetwork(), name.c_str(), inputName.c_str(), name.c_str(), /* return indices tensor */ nullptr, kernelSize.data(), padding.data(), padding.data(), stride.data(), poolDims, poolType, !countIncludePads, ceilMode); @@ -226,41 +242,41 @@ class ReshapeNodeImporter : public INNPIFXNodeImporter { FXNNPIImporter &importer) override { const auto &inputs = node["args"]; const auto &name = node["name"].getString(); + const auto &inputName = getInputNodeName(inputs[0]); NNPITensorDesc desc; importer.updateDescDimsFromFX( toIntegerArray(node["shape"].getString()), desc); - importer.setUsedTensors({inputs[0]["name"].getString()}, {name}); + importer.setUsedTensors({inputName}, {name}); return nnpiNetworkAddReshapeOp(importer.getNetwork(), name.c_str(), - inputs[0]["name"].getString().c_str(), - name.c_str(), &desc); + inputName.c_str(), name.c_str(), &desc); } }; -std::unordered_map< - std::string, - std::unique_ptr>::value_type FXImporterInit[] = { - // _operator - {"_operator.add", - std::make_unique>()}, - - // torch - {"torch.flatten", std::make_unique()}, - - // torch.nn.modules - {"torch.nn.modules.linear.Linear", std::make_unique()}, - {"torch.nn.modules.conv.Conv2d", - std::make_unique>()}, - {"torch.nn.modules.batchnorm.BatchNorm2d", - std::make_unique()}, - {"torch.nn.modules.activation.ReLU", std::make_unique()}, - {"torch.nn.modules.pooling.AdaptiveAvgPool2d", - std::make_unique>()}, - {"torch.nn.modules.pooling.MaxPool2d", - std::make_unique>()}, +static std::unordered_map>::value_type + FXImporterInit[] = { + // _operator + {"_operator.add", + std::make_unique>()}, + + // torch + {"torch.flatten", std::make_unique()}, + + // torch.nn.modules + {"torch.nn.functional.linear", std::make_unique()}, + {"torch.conv2d", std::make_unique>()}, + {"torch.nn.modules.batchnorm.BatchNorm2d", + std::make_unique()}, + {"torch.nn.functional.relu", std::make_unique()}, + {"torch.nn.functional.adaptive_avg_pool2d", + std::make_unique>()}, + {"torch.nn.functional.max_pool2d", + std::make_unique>()}, }; -const std::unordered_map> +static const std::unordered_map> FXNodeImporters = {std::make_move_iterator(std::begin(FXImporterInit)), std::make_move_iterator(std::end(FXImporterInit))}; @@ -281,8 +297,37 @@ FXNNPIImporter::~FXNNPIImporter() { } } -const void *FXNNPIImporter::getConstant(const std::string &name) const { - return constants_->count(name) ? constants_->find(name)->second : nullptr; +const char *FXNNPIImporter::getConstantName(llvm::StringRef name) const { + if (name.empty()) { + return nullptr; + } + // If this is a Constant already, return the name back untouched. + if (constants_->count(name)) { + return name.data(); + } + // Else return the name the getattr maps to if it exists. + auto it = getattrs_.find(name); + return it != getattrs_.end() ? it->second.c_str() : nullptr; +} + +const void *FXNNPIImporter::getConstant(llvm::StringRef name) const { + const char *baseName = getConstantName(name); + if (!baseName) { + return nullptr; + } + // There must be a constant with name baseName, so return it. + auto it = constants_->find(baseName); + CHECK(it != constants_->end()) + << "Should have found constant with name " << baseName; + return it->second; +} + +const NNPITensorDesc & +FXNNPIImporter::getTensorDesc(llvm::StringRef name) const { + auto it = tensorDescs_.find(name); + CHECK(it != tensorDescs_.end()) + << "Did not find NNPITensorDesc for " << name.str(); + return it->second; } void FXNNPIImporter::updateDescQuantFromFX(const DTYPE &dtype, @@ -391,6 +436,12 @@ FXNNPIImporter::addTensor(const std::string &name, const string &dtypeStr, } } + // Outputs have already had their descs added by the named node itself. + if (!output) { + bool inserted = tensorDescs_.try_emplace(name, desc).second; + CHECK(inserted) << "Already found tensor desc with name " << name; + } + return nnpiNetworkAddTensor(network_, name.c_str(), &desc, pRawData); } @@ -423,32 +474,39 @@ FXNNPIImporter::importFunction(const folly::dynamic &FXIR, // Add ops node. for (const auto &node : mod["nodes"]) { const auto &opCode = node["op_code"].getString(); - if (isOps(opCode)) { - DBG("Importing Node: " << node["name"].getString()); - // Add node outputs. - LOG_NNPI_IF_ERROR_RETURN_INVALID_HANDLE( - addTensor(node["name"].getString(), node["dtype"].getString(), - toIntegerArray(node["shape"].getString())), - "Failed to add intermediate"); - - if (opCode != "get_attr") { - const auto &targetName = node["target"].getString(); - const auto &functionName = opCode != "call_module" - ? targetName - : node["parameters"]["name"].getString(); - // Import node. - LOG_NNPI_IF_ERROR_RETURN_INVALID_HANDLE( - FXNodeImporters.at(functionName) - ->importNode( - node, - [&prefix, &targetName](const std::string &name) { - return folly::to(prefix, targetName, ".", - name); - }, - *this), - "Failed to import node"); - } + const auto &nodeName = node["name"].getString(); + if (!isOps(opCode)) { + continue; } + DBG("Importing Node: " << nodeName); + // Add node outputs. + LOG_NNPI_IF_ERROR_RETURN_INVALID_HANDLE( + addTensor(nodeName, node["dtype"].getString(), + toIntegerArray(node["shape"].getString())), + "Failed to add intermediate"); + + // Track what Constant each get_attr points to. + if (opCode == "get_attr") { + bool inserted = + getattrs_.try_emplace(nodeName, node["target"].getString()).second; + CHECK(inserted) << "Already mapped a getattr by name " << nodeName + << " to its underlying Constant"; + continue; + } + const auto &targetName = node["target"].getString(); + const auto &functionName = opCode != "call_module" + ? targetName + : node["parameters"]["name"].getString(); + // Import node. + LOG_NNPI_IF_ERROR_RETURN_INVALID_HANDLE( + FXNodeImporters.at(functionName) + ->importNode( + node, + [&prefix, &targetName](const std::string &name) { + return folly::to(prefix, targetName, ".", name); + }, + *this), + "Failed to import node"); } // Add placeholder. @@ -456,9 +514,8 @@ FXNNPIImporter::importFunction(const folly::dynamic &FXIR, const auto &opCode = node["op_code"].getString(); if (!isOps(opCode)) { - const auto &name = opCode != "output" - ? node["name"].getString() - : node["args"][0]["name"].getString(); + const auto &name = opCode != "output" ? node["name"].getString() + : getInputNodeName(node["args"][0]); bool inputVar(readTensors_.count(name) && !writeTensors_.count(name)); bool outputVar(!readTensors_.count(name) && writeTensors_.count(name)); DBG("Add placeholder: " << name); diff --git a/lib/Backends/NNPI/FXIRImporter.h b/lib/Backends/NNPI/FXIRImporter.h index 8e63dab2d8..41f9cef6b9 100644 --- a/lib/Backends/NNPI/FXIRImporter.h +++ b/lib/Backends/NNPI/FXIRImporter.h @@ -16,12 +16,13 @@ #ifndef FX_NNPI_IMPORTER_H #define FX_NNPI_IMPORTER_H -#include "glow/fb/fx_nnpi_importer/Utils.h" +#include "glow/fb/fx/nnpi_importer/Utils.h" #include "glow/lib/Backends/NNPI/NNPIOptions.h" #include "nnpi_network_builder.h" #include "nnpi_network_builder_EXPERIMENTAL.h" #include "nnpi_transformer.h" #include "llvm/ADT/ArrayRef.h" +#include "llvm/ADT/StringRef.h" #include "llvm/Support/raw_ostream.h" #include #include @@ -72,14 +73,30 @@ class FXNNPIImporter { const std::string &offsetTensor = {}, bool forceSymlowp = false); - /// \returns the constants - const void *getConstant(const std::string &name) const; + /// \returns the Constant known by \p name. If \p name is from a getattr the + /// constant will be looked up by the underlying name found by + /// \ref getConstantName. + const void *getConstant(llvm::StringRef name) const; + + /// \returns the underlying name of a Constant given provided \p name. If \p + /// name is already the name of a Constant it is returned, else looks for + /// getattr aliases to return the name of the actual underlying Constant. + const char *getConstantName(llvm::StringRef name) const; + + /// \returns the NNPITensorDesc for node by \p name. + const NNPITensorDesc &getTensorDesc(llvm::StringRef name) const; private: /// Mapping from constant name to a void pointer points to where the constant /// is actually stored. const llvm::StringMap *constants_; + /// Mapping from getattrs to the underlying name of the constants they alias. + llvm::StringMap getattrs_; + + /// Mapping from output node names to the tensor descriptor for that node. + llvm::StringMap tensorDescs_; + /// Set of tensors written to by the function. std::unordered_set writeTensors_; diff --git a/lib/Backends/NNPI/FXNNPICompiledFunction.cpp b/lib/Backends/NNPI/FXNNPICompiledFunction.cpp index 72974d1153..681dab2372 100644 --- a/lib/Backends/NNPI/FXNNPICompiledFunction.cpp +++ b/lib/Backends/NNPI/FXNNPICompiledFunction.cpp @@ -1,7 +1,7 @@ // Copyright 2004-present Facebook. All Rights Reserved. #include "glow/Flags/Flags.h" -#include "glow/fb/fx_nnpi_importer/Utils.h" +#include "glow/fb/fx/nnpi_importer/Utils.h" #include "glow/lib/Backends/NNPI/DebugMacros.h" #include "glow/lib/Backends/NNPI/FXIRImporter.h" #include "glow/lib/Backends/NNPI/NNPICompiledFunction.h"