llvm
/
torch-mlir
mirror of https://github.com/llvm/torch-mlir
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity
torch-mlir/frontends/pytorch/csrc/builder/acap_dispatch.cpp

724 lines
27 KiB
C++
Raw Normal View History

Add boilerplate to do device capture (pytorch 1.6). * Uses the new dispatcher API. * Just prints to the console for the moment when an op is captured. * Executes the op through the existing implementation.
2020-09-26 09:13:16 +08:00
//===- acap_dispatch.cpp --------------------------------------------------===//
//
// This file is licensed under a pytorch-style license
// See frontends/pytorch/LICENSE for license information.
//
//===----------------------------------------------------------------------===//
#include "acap_dispatch.h"
Support optional args/returns and other odds and ends. * None's out Device? args. * Emits bool tensors if needed. * Adds some stderr tracing to better see what is going on. * Test case that exercises NLLLoss. * This test case emits something for backward calculations but there are some issues still to be worked out, so that part is left out of the test case. * Progress on #97
2020-10-30 08:41:15 +08:00
#include "debug.h"
Add boilerplate to do device capture (pytorch 1.6). * Uses the new dispatcher API. * Just prints to the console for the moment when an op is captured. * Executes the op through the existing implementation.
2020-09-26 09:13:16 +08:00
Implement torch.kernel_call capture. * Had to stop short of modifying the function return signature because of a missing C-API upstream. * Committing here is good enough for a test and will resolve the various TODOs about upstream APIs next.
2020-10-06 14:21:21 +08:00
#include "mlir-c/StandardAttributes.h"
More progress on PyTorch acap device capture. * Now gets far enough to capture batch_norm. * Has some issues still with in-place ops. * Can materialize constants. * Includes an upgrade to PyTorch nightly, which has important bug fixes for fallback and boxed kernel dispatch. * Fixes #78, #79, #80. * Will do more testing in a follow-up once further bugs are fixed that facilitate getting at the other features.
2020-10-16 09:28:30 +08:00
#include "mlir-c/StandardTypes.h"
#include "npcomp-c/Types.h"
Add boilerplate to do device capture (pytorch 1.6). * Uses the new dispatcher API. * Just prints to the console for the moment when an op is captured. * Executes the op through the existing implementation.
2020-09-26 09:13:16 +08:00
#include "npcomp/Python/PybindUtils.h"
Implement torch.kernel_call capture. * Had to stop short of modifying the function return signature because of a missing C-API upstream. * Committing here is good enough for a test and will resolve the various TODOs about upstream APIs next.
2020-10-06 14:21:21 +08:00
#include <ATen/core/function_schema.h>
#include <ATen/core/ivalue.h>
#include <ATen/core/stack.h>
Add boilerplate to do device capture (pytorch 1.6). * Uses the new dispatcher API. * Just prints to the console for the moment when an op is captured. * Executes the op through the existing implementation.
2020-09-26 09:13:16 +08:00
#include <c10/core/DispatchKey.h>
#include <torch/library.h>
using namespace torch_mlir;
namespace py = pybind11;
Implement torch.kernel_call capture. * Had to stop short of modifying the function return signature because of a missing C-API upstream. * Committing here is good enough for a test and will resolve the various TODOs about upstream APIs next.
2020-10-06 14:21:21 +08:00
using c10::FunctionSchema;
Work around various PyTorch issues in support of convolution. * Enables the conv2d fwd test and ResA (which are both small). * Deletes resnet18 and vgg, which both run but generate output that crashes FileCheck and lit (or at least makes them take an eternity).
2020-10-19 12:32:29 +08:00
using c10::IValue;
Implement torch.kernel_call capture. * Had to stop short of modifying the function return signature because of a missing C-API upstream. * Committing here is good enough for a test and will resolve the various TODOs about upstream APIs next.
2020-10-06 14:21:21 +08:00
using c10::OperatorHandle;
using c10::Stack;
Add boilerplate to do device capture (pytorch 1.6). * Uses the new dispatcher API. * Just prints to the console for the moment when an op is captured. * Executes the op through the existing implementation.
2020-09-26 09:13:16 +08:00
// TODO: Private use dispatch keys are not made for real uses. Allocate a proper
// dispatch key in upstream PyTorch (DispatchKey.h) prior to maturity. Note
// that the TORCH_LIBRARY_* macros expand this by name and other APIs use its
// enum value, so we define both. We can get rid of both once we have our
// own key.
Add remaining pieces to capture full example models. * Adds Basicpy List, Tuple, Dict types and plumbs through C API. * Started debugging the issues around aten::conv2d capture, but a PyTorch bug is suspected. * Was able to manually verify that the basic conv2d forward test captures correctly with a workaround. * Need to resolve some printing issues upstream and move these tests to an integration test target (they take ~seconds to run).
2020-10-17 08:38:07 +08:00
// TODO: Ask the PT devs why conv is special and only shows up if dispatching
// through the autograd keys.
// https://github.com/llvm/mlir-npcomp/issues/86
Work around various PyTorch issues in support of convolution. * Enables the conv2d fwd test and ResA (which are both small). * Deletes resnet18 and vgg, which both run but generate output that crashes FileCheck and lit (or at least makes them take an eternity).
2020-10-19 12:32:29 +08:00
#define ACAP_DISPATCH_KEY PrivateUse2
#define ACAP_GRAD_DISPATCH_KEY AutogradPrivateUse2
Add boilerplate to do device capture (pytorch 1.6). * Uses the new dispatcher API. * Just prints to the console for the moment when an op is captured. * Executes the op through the existing implementation.
2020-09-26 09:13:16 +08:00
static c10::DispatchKey kAcapDispatchKey = c10::DispatchKey::ACAP_DISPATCH_KEY;
Work around various PyTorch issues in support of convolution. * Enables the conv2d fwd test and ResA (which are both small). * Deletes resnet18 and vgg, which both run but generate output that crashes FileCheck and lit (or at least makes them take an eternity).
2020-10-19 12:32:29 +08:00
static c10::DispatchKey kAcapGradDispatchKey =
c10::DispatchKey::ACAP_GRAD_DISPATCH_KEY;
Expose signature metadata to ops and implement ATenRecognizeKernelsPass pass. * Two op interfaces, one for querying instance metadata and one for getting static data needed to construct an op from a generic form. * For torch.generic_kernel ops, metadata is splatted in during capture from Torch (it comes from the op registry, which will work for either device capture or graph import). * Moved the 'add' out of the generated set so I can experiment on it. It implements the TorchBuildableKernelOpInterface interface which provides its metadata. * The ATenRecognizeKernelsPass pass generically lowers from a torch.generic_kernel to recognized ops that implement the TorchBuildableKernelOpInterface, handling the various types of transformations that we allow at this stage.
2020-10-23 14:31:34 +08:00
AcapController::KernelCallBuilder::KernelCallBuilder(
AcapController &parent, MlirContext context, MlirLocation loc,
Capture backward conv and copy_ kernels. * This is sufficient to capture the forward and backward pass and gradients of a convolutional model with an nllloss. * As with the forward conv, the backward conv is a special case wrapped in an enigma on the PyTorch side. There aren't many like it, so special casing is just what we do. * When I traced this, I found that the copy_ op is not yet boxing compatible so I had to map it manually. If there are many more like this, I'll probably do something a bit more clever to reduce duplication. * This exposes new signature patterns that will need to be handled by the ATen lowering. Will take care of that next: It will be nice to have an e2e of a non-trivial case with full gradients. * Fixes #97.
2020-10-31 13:52:46 +08:00
const c10::OperatorHandle &opHandle,
llvm::Optional<std::string> overrideKernelName)
Expose signature metadata to ops and implement ATenRecognizeKernelsPass pass. * Two op interfaces, one for querying instance metadata and one for getting static data needed to construct an op from a generic form. * For torch.generic_kernel ops, metadata is splatted in during capture from Torch (it comes from the op registry, which will work for either device capture or graph import). * Moved the 'add' out of the generated set so I can experiment on it. It implements the TorchBuildableKernelOpInterface interface which provides its metadata. * The ATenRecognizeKernelsPass pass generically lowers from a torch.generic_kernel to recognized ops that implement the TorchBuildableKernelOpInterface, handling the various types of transformations that we allow at this stage.
2020-10-23 14:31:34 +08:00
: parent(parent), context(context), loc(loc), opHandle(opHandle),
Work around various PyTorch issues in support of convolution. * Enables the conv2d fwd test and ResA (which are both small). * Deletes resnet18 and vgg, which both run but generate output that crashes FileCheck and lit (or at least makes them take an eternity).
2020-10-19 12:32:29 +08:00
state("torch.kernel_call", loc) {
(void)this->context; // Preserve for future.
Capture backward conv and copy_ kernels. * This is sufficient to capture the forward and backward pass and gradients of a convolutional model with an nllloss. * As with the forward conv, the backward conv is a special case wrapped in an enigma on the PyTorch side. There aren't many like it, so special casing is just what we do. * When I traced this, I found that the copy_ op is not yet boxing compatible so I had to map it manually. If there are many more like this, I'll probably do something a bit more clever to reduce duplication. * This exposes new signature patterns that will need to be handled by the ATen lowering. Will take care of that next: It will be nice to have an e2e of a non-trivial case with full gradients. * Fixes #97.
2020-10-31 13:52:46 +08:00
const std::string &kernelName =
overrideKernelName ? *overrideKernelName : opHandle.operator_name().name;
Work around various PyTorch issues in support of convolution. * Enables the conv2d fwd test and ResA (which are both small). * Deletes resnet18 and vgg, which both run but generate output that crashes FileCheck and lit (or at least makes them take an eternity).
2020-10-19 12:32:29 +08:00
MlirNamedAttribute kernelNameAttr = mlirNamedAttributeGet(
Expose signature metadata to ops and implement ATenRecognizeKernelsPass pass. * Two op interfaces, one for querying instance metadata and one for getting static data needed to construct an op from a generic form. * For torch.generic_kernel ops, metadata is splatted in during capture from Torch (it comes from the op registry, which will work for either device capture or graph import). * Moved the 'add' out of the generated set so I can experiment on it. It implements the TorchBuildableKernelOpInterface interface which provides its metadata. * The ATenRecognizeKernelsPass pass generically lowers from a torch.generic_kernel to recognized ops that implement the TorchBuildableKernelOpInterface, handling the various types of transformations that we allow at this stage.
2020-10-23 14:31:34 +08:00
"kernelName",
Work around various PyTorch issues in support of convolution. * Enables the conv2d fwd test and ResA (which are both small). * Deletes resnet18 and vgg, which both run but generate output that crashes FileCheck and lit (or at least makes them take an eternity).
2020-10-19 12:32:29 +08:00
mlirStringAttrGet(context, kernelName.size(), kernelName.data()));
mlirOperationStateAddAttributes(state, 1, &kernelNameAttr);
Expose signature metadata to ops and implement ATenRecognizeKernelsPass pass. * Two op interfaces, one for querying instance metadata and one for getting static data needed to construct an op from a generic form. * For torch.generic_kernel ops, metadata is splatted in during capture from Torch (it comes from the op registry, which will work for either device capture or graph import). * Moved the 'add' out of the generated set so I can experiment on it. It implements the TorchBuildableKernelOpInterface interface which provides its metadata. * The ATenRecognizeKernelsPass pass generically lowers from a torch.generic_kernel to recognized ops that implement the TorchBuildableKernelOpInterface, handling the various types of transformations that we allow at this stage.
2020-10-23 14:31:34 +08:00
addSchemaAttrs();
}
void AcapController::KernelCallBuilder::addSchemaAttrs() {
// Map the op schema to the kernel_call attributes:
// sigArgTypes
// sigRetTypes
// sigIsVararg
// sigIsVarret
// sigIsMutable
const c10::FunctionSchema &schema = opHandle.schema();
llvm::SmallVector<MlirNamedAttribute, 8> attrs;
attrs.push_back(mlirNamedAttributeGet(
"sigIsMutable", mlirBoolAttrGet(context, schema.is_mutable())));
attrs.push_back(mlirNamedAttributeGet(
"sigIsVararg", mlirBoolAttrGet(context, schema.is_vararg())));
attrs.push_back(mlirNamedAttributeGet(
"sigIsVarret", mlirBoolAttrGet(context, schema.is_varret())));
// Arg types.
llvm::SmallVector<MlirAttribute, 4> args;
for (auto &arg : schema.arguments()) {
const std::string &typeStr = arg.type()->str();
args.push_back(mlirStringAttrGet(context, typeStr.size(), typeStr.data()));
}
attrs.push_back(mlirNamedAttributeGet(
"sigArgTypes", mlirArrayAttrGet(context, args.size(), args.data())));
// Return types.
llvm::SmallVector<MlirAttribute, 4> returns;
for (auto &ret : schema.returns()) {
const std::string &typeStr = ret.type()->str();
returns.push_back(
mlirStringAttrGet(context, typeStr.size(), typeStr.data()));
}
attrs.push_back(mlirNamedAttributeGet(
"sigRetTypes",
mlirArrayAttrGet(context, returns.size(), returns.data())));
// Add attrs to op.
mlirOperationStateAddAttributes(state, attrs.size(), attrs.data());
Work around various PyTorch issues in support of convolution. * Enables the conv2d fwd test and ResA (which are both small). * Deletes resnet18 and vgg, which both run but generate output that crashes FileCheck and lit (or at least makes them take an eternity).
2020-10-19 12:32:29 +08:00
}
void AcapController::KernelCallBuilder::addOperand(const IValue &value) {
MlirValue mlirValue = parent.mapIValueToMlirValue(loc, value);
if (mlirValueIsNull(mlirValue)) {
std::stringstream out;
Expose signature metadata to ops and implement ATenRecognizeKernelsPass pass. * Two op interfaces, one for querying instance metadata and one for getting static data needed to construct an op from a generic form. * For torch.generic_kernel ops, metadata is splatted in during capture from Torch (it comes from the op registry, which will work for either device capture or graph import). * Moved the 'add' out of the generated set so I can experiment on it. It implements the TorchBuildableKernelOpInterface interface which provides its metadata. * The ATenRecognizeKernelsPass pass generically lowers from a torch.generic_kernel to recognized ops that implement the TorchBuildableKernelOpInterface, handling the various types of transformations that we allow at this stage.
2020-10-23 14:31:34 +08:00
const std::string &kernelName = opHandle.operator_name().name;
Support optional args/returns and other odds and ends. * None's out Device? args. * Emits bool tensors if needed. * Adds some stderr tracing to better see what is going on. * Test case that exercises NLLLoss. * This test case emits something for backward calculations but there are some issues still to be worked out, so that part is left out of the test case. * Progress on #97
2020-10-30 08:41:15 +08:00
out << "Unsupported capture value passed to kernel '" << kernelName << "' ("
<< value.tagKind() << "): " << value;
Work around various PyTorch issues in support of convolution. * Enables the conv2d fwd test and ResA (which are both small). * Deletes resnet18 and vgg, which both run but generate output that crashes FileCheck and lit (or at least makes them take an eternity).
2020-10-19 12:32:29 +08:00
throw std::invalid_argument(out.str());
}
mlirOperationStateAddOperands(state, 1, &mlirValue);
}
void AcapController::KernelCallBuilder::addResult(const IValue &value) {
MlirType resultType = parent.mapIValueToMlirType(loc, value);
if (mlirTypeIsNull(resultType)) {
std::stringstream out;
Expose signature metadata to ops and implement ATenRecognizeKernelsPass pass. * Two op interfaces, one for querying instance metadata and one for getting static data needed to construct an op from a generic form. * For torch.generic_kernel ops, metadata is splatted in during capture from Torch (it comes from the op registry, which will work for either device capture or graph import). * Moved the 'add' out of the generated set so I can experiment on it. It implements the TorchBuildableKernelOpInterface interface which provides its metadata. * The ATenRecognizeKernelsPass pass generically lowers from a torch.generic_kernel to recognized ops that implement the TorchBuildableKernelOpInterface, handling the various types of transformations that we allow at this stage.
2020-10-23 14:31:34 +08:00
const std::string &kernelName = opHandle.operator_name().name;
Work around various PyTorch issues in support of convolution. * Enables the conv2d fwd test and ResA (which are both small). * Deletes resnet18 and vgg, which both run but generate output that crashes FileCheck and lit (or at least makes them take an eternity).
2020-10-19 12:32:29 +08:00
out << "Unsupported capture value returned from kernel '" << kernelName
<< "' (" << value.tagKind() << "): " << value;
throw std::invalid_argument(out.str());
}
if (value.isTensor()) {
resultIndexToTensorMap.emplace_back(resultCount++, value.toTensor());
}
mlirOperationStateAddResults(state, 1, &resultType);
}
MlirOperation AcapController::KernelCallBuilder::create() {
// Create operation.
MlirOperation op = state.createOperation();
parent.funcBuilder->getEntryBlockBuilder().insertBeforeTerminator(op);
// Map result tensors.
for (auto &it : resultIndexToTensorMap) {
MlirValue result = mlirOperationGetResult(op, it.first);
parent.funcBuilder->mapTensor(it.second, result);
}
return op;
}
Add boilerplate to do device capture (pytorch 1.6). * Uses the new dispatcher API. * Just prints to the console for the moment when an op is captured. * Executes the op through the existing implementation.
2020-09-26 09:13:16 +08:00
std::list<AcapController::Activation> &
AcapController::getThreadLocalActiveStack() {
static thread_local std::list<Activation> threadLocalActiveStack;
return threadLocalActiveStack;
}
py::object AcapController::contextEnter() {
auto &stack = getThreadLocalActiveStack();
stack.emplace_front(shared_from_this());
Activation &current = stack.front();
Work around various PyTorch issues in support of convolution. * Enables the conv2d fwd test and ResA (which are both small). * Deletes resnet18 and vgg, which both run but generate output that crashes FileCheck and lit (or at least makes them take an eternity).
2020-10-19 12:32:29 +08:00
c10::DispatchKeySet keySet{kAcapDispatchKey, kAcapGradDispatchKey};
current.includeGuard =
std::make_unique<c10::impl::IncludeDispatchKeyGuard>(keySet);
Add boilerplate to do device capture (pytorch 1.6). * Uses the new dispatcher API. * Just prints to the console for the moment when an op is captured. * Executes the op through the existing implementation.
2020-09-26 09:13:16 +08:00
return py::cast(this);
}
void AcapController::contextExit(py::object exc_type, py::object exc_val,
py::object exc_tb) {
auto &stack = getThreadLocalActiveStack();
if (stack.empty() || stack.front().controller.get() != this) {
throw py::raisePyError(PyExc_RuntimeError,
"Mismatched context manager __exit__");
}
stack.pop_front();
Implement torch.kernel_call capture. * Had to stop short of modifying the function return signature because of a missing C-API upstream. * Committing here is good enough for a test and will resolve the various TODOs about upstream APIs next.
2020-10-06 14:21:21 +08:00
if (!hasReturned) {
returns({});
}
}
void AcapController::returns(std::vector<at::Tensor> tensors) {
verifyHasNotReturned();
llvm::SmallVector<MlirType, 4> returnsTypes;
llvm::SmallVector<MlirValue, 4> returnsValues;
for (auto &tensor : tensors) {
MlirValue v = funcBuilder->lookupTensor(tensor);
Set PyTorch captured function return type. * Resolves various TODOs that required an LLVM change/bump. * Bumps LLVM to 4aa217160e5f06a96c6effc4950c3b402374de58
2020-10-08 01:14:34 +08:00
if (mlirValueIsNull(v)) {
Fix dispatch of arange. * Fixes #107 * I wouldn't say I love what had to be done here. Worth a conversation with the PT devs (probably as part of a rollup of a bunch of this stuff).
2020-11-13 09:42:19 +08:00
debugTrace(
"Return of imported-constant tensor (intentional memorization?)");
v = importTensorByValue(tensor);
Implement torch.kernel_call capture. * Had to stop short of modifying the function return signature because of a missing C-API upstream. * Committing here is good enough for a test and will resolve the various TODOs about upstream APIs next.
2020-10-06 14:21:21 +08:00
}
returnsTypes.push_back(mlirValueGetType(v));
returnsValues.push_back(v);
}
More progress on PyTorch acap device capture. * Now gets far enough to capture batch_norm. * Has some issues still with in-place ops. * Can materialize constants. * Includes an upgrade to PyTorch nightly, which has important bug fixes for fallback and boxed kernel dispatch. * Fixes #78, #79, #80. * Will do more testing in a follow-up once further bugs are fixed that facilitate getting at the other features.
2020-10-16 09:28:30 +08:00
MlirLocation loc = getCurrentLocation();
Implement torch.kernel_call capture. * Had to stop short of modifying the function return signature because of a missing C-API upstream. * Committing here is good enough for a test and will resolve the various TODOs about upstream APIs next.
2020-10-06 14:21:21 +08:00
OperationStateHolder s("std.return", loc);
Add remaining pieces to capture full example models. * Adds Basicpy List, Tuple, Dict types and plumbs through C API. * Started debugging the issues around aten::conv2d capture, but a PyTorch bug is suspected. * Was able to manually verify that the basic conv2d forward test captures correctly with a workaround. * Need to resolve some printing issues upstream and move these tests to an integration test target (they take ~seconds to run).
2020-10-17 08:38:07 +08:00
mlirOperationStateAddOperands(s, returnsValues.size(), returnsValues.data());
Implement torch.kernel_call capture. * Had to stop short of modifying the function return signature because of a missing C-API upstream. * Committing here is good enough for a test and will resolve the various TODOs about upstream APIs next.
2020-10-06 14:21:21 +08:00
funcBuilder->getEntryBlockBuilder().insertBeforeTerminator(
s.createOperation());
funcBuilder->rewriteFuncReturnTypes(returnsTypes);
hasReturned = true;
Add boilerplate to do device capture (pytorch 1.6). * Uses the new dispatcher API. * Just prints to the console for the moment when an op is captured. * Executes the op through the existing implementation.
2020-09-26 09:13:16 +08:00
}
Work around various PyTorch issues in support of convolution. * Enables the conv2d fwd test and ResA (which are both small). * Deletes resnet18 and vgg, which both run but generate output that crashes FileCheck and lit (or at least makes them take an eternity).
2020-10-19 12:32:29 +08:00
std::shared_ptr<AcapController>
AcapController::getCurrentThreadAcapController() {
Add boilerplate to do device capture (pytorch 1.6). * Uses the new dispatcher API. * Just prints to the console for the moment when an op is captured. * Executes the op through the existing implementation.
2020-09-26 09:13:16 +08:00
auto &stack = getThreadLocalActiveStack();
if (stack.empty())
return nullptr;
return stack.front().controller;
}
Implement torch.kernel_call capture. * Had to stop short of modifying the function return signature because of a missing C-API upstream. * Committing here is good enough for a test and will resolve the various TODOs about upstream APIs next.
2020-10-06 14:21:21 +08:00
void AcapController::verifyHasNotReturned() {
if (hasReturned) {
throw std::runtime_error(
"Function has already returned. Cannot trace more operations.");
}
}
Add boilerplate to do device capture (pytorch 1.6). * Uses the new dispatcher API. * Just prints to the console for the moment when an op is captured. * Executes the op through the existing implementation.
2020-09-26 09:13:16 +08:00
Implement torch.kernel_call capture. * Had to stop short of modifying the function return signature because of a missing C-API upstream. * Committing here is good enough for a test and will resolve the various TODOs about upstream APIs next.
2020-10-06 14:21:21 +08:00
/* static */
void AcapController::fallbackKernel(const OperatorHandle &opHandle,
Stack *stack) {
Fix dispatch of arange. * Fixes #107 * I wouldn't say I love what had to be done here. Worth a conversation with the PT devs (probably as part of a rollup of a bunch of this stuff).
2020-11-13 09:42:19 +08:00
auto redispatchCallback = [&]() {
// Exclude recursive dispatch to this kernel.
c10::impl::ExcludeDispatchKeyGuard exclusion(kAcapDispatchKey);
// Passthrough.
auto &dispatcher = c10::Dispatcher::singleton();
dispatcher.callBoxed(opHandle, stack);
};
Work around various PyTorch issues in support of convolution. * Enables the conv2d fwd test and ResA (which are both small). * Deletes resnet18 and vgg, which both run but generate output that crashes FileCheck and lit (or at least makes them take an eternity).
2020-10-19 12:32:29 +08:00
auto current = getCurrentThreadAcapController();
Implement torch.kernel_call capture. * Had to stop short of modifying the function return signature because of a missing C-API upstream. * Committing here is good enough for a test and will resolve the various TODOs about upstream APIs next.
2020-10-06 14:21:21 +08:00
if (!current) {
Fix dispatch of arange. * Fixes #107 * I wouldn't say I love what had to be done here. Worth a conversation with the PT devs (probably as part of a rollup of a bunch of this stuff).
2020-11-13 09:42:19 +08:00
redispatchCallback();
Implement torch.kernel_call capture. * Had to stop short of modifying the function return signature because of a missing C-API upstream. * Committing here is good enough for a test and will resolve the various TODOs about upstream APIs next.
2020-10-06 14:21:21 +08:00
return;
Add boilerplate to do device capture (pytorch 1.6). * Uses the new dispatcher API. * Just prints to the console for the moment when an op is captured. * Executes the op through the existing implementation.
2020-09-26 09:13:16 +08:00
}
Fix dispatch of arange. * Fixes #107 * I wouldn't say I love what had to be done here. Worth a conversation with the PT devs (probably as part of a rollup of a bunch of this stuff).
2020-11-13 09:42:19 +08:00
current->fallbackKernelImpl(opHandle, stack, redispatchCallback);
Implement torch.kernel_call capture. * Had to stop short of modifying the function return signature because of a missing C-API upstream. * Committing here is good enough for a test and will resolve the various TODOs about upstream APIs next.
2020-10-06 14:21:21 +08:00
}
Add boilerplate to do device capture (pytorch 1.6). * Uses the new dispatcher API. * Just prints to the console for the moment when an op is captured. * Executes the op through the existing implementation.
2020-09-26 09:13:16 +08:00
Work around various PyTorch issues in support of convolution. * Enables the conv2d fwd test and ResA (which are both small). * Deletes resnet18 and vgg, which both run but generate output that crashes FileCheck and lit (or at least makes them take an eternity).
2020-10-19 12:32:29 +08:00
at::Tensor AcapController::convolutionKernel(
const at::Tensor &input, const at::Tensor &weight,
const c10::optional<at::Tensor> &bias, const at::IntArrayRef stride,
const at::IntArrayRef padding, const at::IntArrayRef dilation,
const bool transposed, const at::IntArrayRef output_padding,
const int64_t groups) {
static c10::OperatorName opName{"aten::convolution", ""};
auto &dispatcher = c10::Dispatcher::singleton();
auto opHandle = dispatcher.findOp(opName);
assert(opHandle && "could not find convolution op");
Support optional args/returns and other odds and ends. * None's out Device? args. * Emits bool tensors if needed. * Adds some stderr tracing to better see what is going on. * Test case that exercises NLLLoss. * This test case emits something for backward calculations but there are some issues still to be worked out, so that part is left out of the test case. * Progress on #97
2020-10-30 08:41:15 +08:00
if (isDebugTraceEnabled()) {
std::stringstream s;
s << "Convolution (unboxed) dispatch: " << opHandle->schema();
debugTrace(s.str());
}
Work around various PyTorch issues in support of convolution. * Enables the conv2d fwd test and ResA (which are both small). * Deletes resnet18 and vgg, which both run but generate output that crashes FileCheck and lit (or at least makes them take an eternity).
2020-10-19 12:32:29 +08:00
auto opTyped = opHandle->typed<at::Tensor(
const at::Tensor &, const at::Tensor &, const c10::optional<at::Tensor> &,
const at::IntArrayRef, const at::IntArrayRef, const at::IntArrayRef,
const bool, const at::IntArrayRef, const int64_t)>();
// Exclude recursive calls: convolution is completely emitted by this
// kernel.
c10::DispatchKeySet keySet{kAcapDispatchKey, kAcapGradDispatchKey};
c10::impl::ExcludeDispatchKeyGuard exclusion(keySet);
auto current = getCurrentThreadAcapController();
if (!current) {
return opTyped.callWithDispatchKey(c10::DispatchKey::AutogradOther, input,
weight, bias, stride, padding, dilation,
transposed, output_padding, groups);
}
MlirContext context = current->funcBuilder->getContext();
MlirLocation loc = current->getCurrentLocation();
std::string kernelName{"aten::convolution"};
Expose signature metadata to ops and implement ATenRecognizeKernelsPass pass. * Two op interfaces, one for querying instance metadata and one for getting static data needed to construct an op from a generic form. * For torch.generic_kernel ops, metadata is splatted in during capture from Torch (it comes from the op registry, which will work for either device capture or graph import). * Moved the 'add' out of the generated set so I can experiment on it. It implements the TorchBuildableKernelOpInterface interface which provides its metadata. * The ATenRecognizeKernelsPass pass generically lowers from a torch.generic_kernel to recognized ops that implement the TorchBuildableKernelOpInterface, handling the various types of transformations that we allow at this stage.
2020-10-23 14:31:34 +08:00
KernelCallBuilder callBuilder{*current, context, loc, *opHandle};
Work around various PyTorch issues in support of convolution. * Enables the conv2d fwd test and ResA (which are both small). * Deletes resnet18 and vgg, which both run but generate output that crashes FileCheck and lit (or at least makes them take an eternity).
2020-10-19 12:32:29 +08:00
callBuilder.addOperand(IValue(input));
callBuilder.addOperand(IValue(weight));
// This is really sad: instead of storing a none in the optional, it stores
// an undefined tensor, which cannot convert to an IValue :(
// TODO: File PyTorch bug. Perhaps this is why they don't support boxing
// for it.
IValue biasIValue;
if (bias && bias->defined()) {
biasIValue = IValue(bias);
} else {
biasIValue = IValue(c10::optional<at::Tensor>());
}
callBuilder.addOperand(biasIValue);
callBuilder.addOperand(IValue(stride));
callBuilder.addOperand(IValue(padding));
callBuilder.addOperand(IValue(dilation));
callBuilder.addOperand(IValue(transposed));
callBuilder.addOperand(IValue(output_padding));
callBuilder.addOperand(IValue(groups));
auto result = opTyped.callWithDispatchKey(
c10::DispatchKey::AutogradOther, input, weight, bias, stride, padding,
dilation, transposed, output_padding, groups);
callBuilder.addResult(result);
callBuilder.create();
return result;
}
Capture backward conv and copy_ kernels. * This is sufficient to capture the forward and backward pass and gradients of a convolutional model with an nllloss. * As with the forward conv, the backward conv is a special case wrapped in an enigma on the PyTorch side. There aren't many like it, so special casing is just what we do. * When I traced this, I found that the copy_ op is not yet boxing compatible so I had to map it manually. If there are many more like this, I'll probably do something a bit more clever to reduce duplication. * This exposes new signature patterns that will need to be handled by the ATen lowering. Will take care of that next: It will be nice to have an e2e of a non-trivial case with full gradients. * Fixes #97.
2020-10-31 13:52:46 +08:00
std::tuple<at::Tensor, at::Tensor, at::Tensor>
AcapController::mklConvolutionBackward(
const at::Tensor &input, const at::Tensor &grad_output,
const at::Tensor &weight, const at::IntArrayRef padding,
const at::IntArrayRef stride, const at::IntArrayRef dilation,
const int64_t groups, std::array<bool, 3> output_mask) {
static c10::OperatorName opName{"aten::mkldnn_convolution_backward", ""};
auto &dispatcher = c10::Dispatcher::singleton();
auto opHandle = dispatcher.findOp(opName);
assert(opHandle && "could not find mkldnn_convolution_backward op");
if (isDebugTraceEnabled()) {
std::stringstream s;
s << "mkldnn_convolution_backward dispatch: " << opHandle->schema();
debugTrace(s.str());
}
auto opTyped = opHandle->typed<std::tuple<at::Tensor, at::Tensor, at::Tensor>(
const at::Tensor &input, const at::Tensor &grad_output,
const at::Tensor &weight, const at::IntArrayRef padding,
const at::IntArrayRef stride, const at::IntArrayRef dilation,
const int64_t groups, std::array<bool, 3> output_mask)>();
// Exclude recursive calls: convolution is completely emitted by this
// kernel.
c10::DispatchKeySet keySet{kAcapDispatchKey, kAcapGradDispatchKey};
c10::impl::ExcludeDispatchKeyGuard exclusion(keySet);
auto current = getCurrentThreadAcapController();
if (!current) {
return opTyped.callWithDispatchKey(c10::DispatchKey::AutogradOther, input,
grad_output, weight, padding, stride,
dilation, groups, output_mask);
}
// Emit the call as if to aten::convolution_overridable, the generic, full
// parameterized versions that backends are supposed to implement.
// Requires some parameter swizzling.
// It has the signature:
// convolution_backward_overrideable(Tensor grad_output, Tensor input,
// Tensor weight, int[] stride, int[] padding, int[] dilation,
// bool transposed, int[] output_padding, int groups,
// bool[3] output_mask) ->
// (Tensor grad_input, Tensor grad_weight, Tensor grad_bias)
MlirContext context = current->funcBuilder->getContext();
MlirLocation loc = current->getCurrentLocation();
std::string kernelName{"aten::convolution_backward"};
static c10::OperatorName emitOpName{"aten::convolution_backward_overrideable",
""};
auto emitOpHandle = dispatcher.findOp(emitOpName);
assert(emitOpHandle && "could not find convolution_backward_overrideable op");
KernelCallBuilder callBuilder{*current, context, loc, *emitOpHandle,
kernelName};
callBuilder.addOperand(IValue(grad_output));
callBuilder.addOperand(IValue(input));
callBuilder.addOperand(IValue(weight));
callBuilder.addOperand(IValue(stride));
callBuilder.addOperand(IValue(padding));
callBuilder.addOperand(IValue(dilation));
callBuilder.addOperand(IValue(false));
std::vector<int64_t> output_padding(padding.size()); // Not provided.
callBuilder.addOperand(IValue(at::IntArrayRef(output_padding)));
callBuilder.addOperand(IValue(groups));
callBuilder.addOperand(IValue(output_mask));
auto results = opTyped.callWithDispatchKey(
c10::DispatchKey::AutogradCPU, input, grad_output, weight, padding,
stride, dilation, groups, output_mask);
callBuilder.addResult(std::get<0>(results));
callBuilder.addResult(std::get<1>(results));
callBuilder.addResult(std::get<2>(results));
callBuilder.create();
return results;
}
at::Tensor &AcapController::copyUnderKernel(at::Tensor &self,
const at::Tensor &src,
bool non_blocking) {
static c10::OperatorName opName{"aten::copy_", ""};
auto &dispatcher = c10::Dispatcher::singleton();
auto opHandle = dispatcher.findOp(opName);
assert(opHandle && "could not find copy_ op");
if (isDebugTraceEnabled()) {
std::stringstream s;
s << "copy_ dispatch: " << opHandle->schema();
debugTrace(s.str());
}
auto opTyped = opHandle->typed<at::Tensor &(
at::Tensor & self, const at::Tensor &src, bool non_blocking)>();
// Exclude recursive calls.
c10::DispatchKeySet keySet{kAcapDispatchKey, kAcapGradDispatchKey};
c10::impl::ExcludeDispatchKeyGuard exclusion(keySet);
auto current = getCurrentThreadAcapController();
if (!current) {
return opTyped.callWithDispatchKey(c10::DispatchKey::AutogradOther, self,
src, non_blocking);
}
MlirContext context = current->funcBuilder->getContext();
MlirLocation loc = current->getCurrentLocation();
KernelCallBuilder callBuilder{*current, context, loc, *opHandle};
callBuilder.addOperand(IValue(self));
callBuilder.addOperand(IValue(src));
auto &result = opTyped.callWithDispatchKey(c10::DispatchKey::CPU, self, src,
non_blocking);
callBuilder.addResult(result);
callBuilder.create();
return result;
}
Fix dispatch of arange. * Fixes #107 * I wouldn't say I love what had to be done here. Worth a conversation with the PT devs (probably as part of a rollup of a bunch of this stuff).
2020-11-13 09:42:19 +08:00
at::Tensor AcapController::arangeBackendSelectKernel(
at::Scalar end, c10::optional<at::ScalarType> dtype,
c10::optional<at::Layout> layout, c10::optional<at::Device> device,
c10::optional<bool> pin_memory) {
static c10::OperatorName opName{"aten::arange", ""};
auto &dispatcher = c10::Dispatcher::singleton();
auto opHandle = dispatcher.findOp(opName);
assert(opHandle && "could not find arange op");
// Exclude recursive calls.
c10::DispatchKeySet keySet{kAcapDispatchKey, kAcapGradDispatchKey};
c10::impl::ExcludeDispatchKeyGuard exclusion(keySet);
// Dispatching in this fashion replicates the exact way that PyTorch
// built-in handlers dispatch to BackendSelect kernels.
auto targetDk = c10::computeDispatchKey(dtype, layout, device);
auto opTyped = opHandle->typed<at::Tensor(
at::Scalar end, c10::optional<at::ScalarType> dtype,
c10::optional<at::Layout> layout, c10::optional<at::Device> device,
c10::optional<bool> pin_memory)>();
return opTyped.callWithDispatchKey(targetDk, end, dtype, layout, device,
pin_memory);
More progress on PyTorch acap device capture. * Now gets far enough to capture batch_norm. * Has some issues still with in-place ops. * Can materialize constants. * Includes an upgrade to PyTorch nightly, which has important bug fixes for fallback and boxed kernel dispatch. * Fixes #78, #79, #80. * Will do more testing in a follow-up once further bugs are fixed that facilitate getting at the other features.
2020-10-16 09:28:30 +08:00
}
Fix dispatch of arange. * Fixes #107 * I wouldn't say I love what had to be done here. Worth a conversation with the PT devs (probably as part of a rollup of a bunch of this stuff).
2020-11-13 09:42:19 +08:00
MlirLocation AcapController::getCurrentLocation() {
return mlirLocationUnknownGet(funcBuilder->getContext());
Add boilerplate to do device capture (pytorch 1.6). * Uses the new dispatcher API. * Just prints to the console for the moment when an op is captured. * Executes the op through the existing implementation.
2020-09-26 09:13:16 +08:00
}
Fix dispatch of arange. * Fixes #107 * I wouldn't say I love what had to be done here. Worth a conversation with the PT devs (probably as part of a rollup of a bunch of this stuff).
2020-11-13 09:42:19 +08:00
void AcapController::fallbackKernelImpl(
const OperatorHandle &opHandle, Stack *stack,
std::function<void()> redispatchCallback) {
Implement torch.kernel_call capture. * Had to stop short of modifying the function return signature because of a missing C-API upstream. * Committing here is good enough for a test and will resolve the various TODOs about upstream APIs next.
2020-10-06 14:21:21 +08:00
verifyHasNotReturned();
Support optional args/returns and other odds and ends. * None's out Device? args. * Emits bool tensors if needed. * Adds some stderr tracing to better see what is going on. * Test case that exercises NLLLoss. * This test case emits something for backward calculations but there are some issues still to be worked out, so that part is left out of the test case. * Progress on #97
2020-10-30 08:41:15 +08:00
if (isDebugTraceEnabled()) {
std::stringstream s;
Fix dispatch of arange. * Fixes #107 * I wouldn't say I love what had to be done here. Worth a conversation with the PT devs (probably as part of a rollup of a bunch of this stuff).
2020-11-13 09:42:19 +08:00
s << "Fallback (boxed) dispatch: " << opHandle.schema()
<< " (stack size=" << stack->size() << ")";
Support optional args/returns and other odds and ends. * None's out Device? args. * Emits bool tensors if needed. * Adds some stderr tracing to better see what is going on. * Test case that exercises NLLLoss. * This test case emits something for backward calculations but there are some issues still to be worked out, so that part is left out of the test case. * Progress on #97
2020-10-30 08:41:15 +08:00
debugTrace(s.str());
}
Implement torch.kernel_call capture. * Had to stop short of modifying the function return signature because of a missing C-API upstream. * Committing here is good enough for a test and will resolve the various TODOs about upstream APIs next.
2020-10-06 14:21:21 +08:00
// Exclude recursive dispatch to this kernel.
c10::impl::ExcludeDispatchKeyGuard exclusion(kAcapDispatchKey);
const FunctionSchema &schema = opHandle.schema();
// Check for unsupported.
if (schema.is_vararg() || schema.is_varret()) {
throw std::invalid_argument(
"Cannot capture ops with variable arguments or returns");
}
MlirContext context = funcBuilder->getContext();
Work around various PyTorch issues in support of convolution. * Enables the conv2d fwd test and ResA (which are both small). * Deletes resnet18 and vgg, which both run but generate output that crashes FileCheck and lit (or at least makes them take an eternity).
2020-10-19 12:32:29 +08:00
MlirLocation loc = getCurrentLocation();
Implement torch.kernel_call capture. * Had to stop short of modifying the function return signature because of a missing C-API upstream. * Committing here is good enough for a test and will resolve the various TODOs about upstream APIs next.
2020-10-06 14:21:21 +08:00
auto kernelName = schema.name();
Expose signature metadata to ops and implement ATenRecognizeKernelsPass pass. * Two op interfaces, one for querying instance metadata and one for getting static data needed to construct an op from a generic form. * For torch.generic_kernel ops, metadata is splatted in during capture from Torch (it comes from the op registry, which will work for either device capture or graph import). * Moved the 'add' out of the generated set so I can experiment on it. It implements the TorchBuildableKernelOpInterface interface which provides its metadata. * The ATenRecognizeKernelsPass pass generically lowers from a torch.generic_kernel to recognized ops that implement the TorchBuildableKernelOpInterface, handling the various types of transformations that we allow at this stage.
2020-10-23 14:31:34 +08:00
KernelCallBuilder callBuilder{*this, context, loc, opHandle};
Implement torch.kernel_call capture. * Had to stop short of modifying the function return signature because of a missing C-API upstream. * Committing here is good enough for a test and will resolve the various TODOs about upstream APIs next.
2020-10-06 14:21:21 +08:00
// Map arguments to operands.
// This must be accumulated into the OperationState prior to re-dispatch
// since the stack is modified at that point.
size_t argCount = schema.arguments().size();
assert(stack->size() >= argCount && "stack too short");
for (auto argIt = stack->end() - argCount; argIt != stack->end(); ++argIt) {
Work around various PyTorch issues in support of convolution. * Enables the conv2d fwd test and ResA (which are both small). * Deletes resnet18 and vgg, which both run but generate output that crashes FileCheck and lit (or at least makes them take an eternity).
2020-10-19 12:32:29 +08:00
callBuilder.addOperand(*argIt);
Implement torch.kernel_call capture. * Had to stop short of modifying the function return signature because of a missing C-API upstream. * Committing here is good enough for a test and will resolve the various TODOs about upstream APIs next.
2020-10-06 14:21:21 +08:00
}
// Invoke the original kernel.
Fix dispatch of arange. * Fixes #107 * I wouldn't say I love what had to be done here. Worth a conversation with the PT devs (probably as part of a rollup of a bunch of this stuff).
2020-11-13 09:42:19 +08:00
redispatchCallback();
Implement torch.kernel_call capture. * Had to stop short of modifying the function return signature because of a missing C-API upstream. * Committing here is good enough for a test and will resolve the various TODOs about upstream APIs next.
2020-10-06 14:21:21 +08:00
// Map returns to results.
size_t returnCount = schema.returns().size();
assert(stack->size() >= returnCount && "stack too short");
for (auto returnIt = stack->end() - returnCount; returnIt != stack->end();
++returnIt) {
Work around various PyTorch issues in support of convolution. * Enables the conv2d fwd test and ResA (which are both small). * Deletes resnet18 and vgg, which both run but generate output that crashes FileCheck and lit (or at least makes them take an eternity).
2020-10-19 12:32:29 +08:00
callBuilder.addResult(*returnIt);
Implement torch.kernel_call capture. * Had to stop short of modifying the function return signature because of a missing C-API upstream. * Committing here is good enough for a test and will resolve the various TODOs about upstream APIs next.
2020-10-06 14:21:21 +08:00
}
Work around various PyTorch issues in support of convolution. * Enables the conv2d fwd test and ResA (which are both small). * Deletes resnet18 and vgg, which both run but generate output that crashes FileCheck and lit (or at least makes them take an eternity).
2020-10-19 12:32:29 +08:00
callBuilder.create();
Implement torch.kernel_call capture. * Had to stop short of modifying the function return signature because of a missing C-API upstream. * Committing here is good enough for a test and will resolve the various TODOs about upstream APIs next.
2020-10-06 14:21:21 +08:00
}
MlirValue AcapController::mapIValueToMlirValue(MlirLocation loc,
Work around various PyTorch issues in support of convolution. * Enables the conv2d fwd test and ResA (which are both small). * Deletes resnet18 and vgg, which both run but generate output that crashes FileCheck and lit (or at least makes them take an eternity).
2020-10-19 12:32:29 +08:00
const IValue &ival) {
Implement torch.kernel_call capture. * Had to stop short of modifying the function return signature because of a missing C-API upstream. * Committing here is good enough for a test and will resolve the various TODOs about upstream APIs next.
2020-10-06 14:21:21 +08:00
if (ival.isScalar()) {
return funcBuilder->getScalarConstant(loc, ival.toScalar());
}
if (ival.isTensor()) {
Support optional args/returns and other odds and ends. * None's out Device? args. * Emits bool tensors if needed. * Adds some stderr tracing to better see what is going on. * Test case that exercises NLLLoss. * This test case emits something for backward calculations but there are some issues still to be worked out, so that part is left out of the test case. * Progress on #97
2020-10-30 08:41:15 +08:00
auto tensor = ival.toTensor();
if (!tensor.defined()) {
// Optional tensors ("Tensor?" type) are represented as Tensor ivals
// that are undefined.
return funcBuilder->getNoneConstant(loc);
}
Implement torch.kernel_call capture. * Had to stop short of modifying the function return signature because of a missing C-API upstream. * Committing here is good enough for a test and will resolve the various TODOs about upstream APIs next.
2020-10-06 14:21:21 +08:00
// Is it an already mapped tensor?
MlirValue mappedValue = funcBuilder->lookupTensor(ival.toTensor());
More progress on PyTorch acap device capture. * Now gets far enough to capture batch_norm. * Has some issues still with in-place ops. * Can materialize constants. * Includes an upgrade to PyTorch nightly, which has important bug fixes for fallback and boxed kernel dispatch. * Fixes #78, #79, #80. * Will do more testing in a follow-up once further bugs are fixed that facilitate getting at the other features.
2020-10-16 09:28:30 +08:00
if (!mlirValueIsNull(mappedValue)) {
Implement torch.kernel_call capture. * Had to stop short of modifying the function return signature because of a missing C-API upstream. * Committing here is good enough for a test and will resolve the various TODOs about upstream APIs next.
2020-10-06 14:21:21 +08:00
return mappedValue;
}
More progress on PyTorch acap device capture. * Now gets far enough to capture batch_norm. * Has some issues still with in-place ops. * Can materialize constants. * Includes an upgrade to PyTorch nightly, which has important bug fixes for fallback and boxed kernel dispatch. * Fixes #78, #79, #80. * Will do more testing in a follow-up once further bugs are fixed that facilitate getting at the other features.
2020-10-16 09:28:30 +08:00
mappedValue = importTensorByValue(ival.toTensor());
assert(mappedValue.ptr);
return mappedValue;
}
if (ival.isBool()) {
// TODO: Switch to the numpy.bool type as that is a closer domain match.
return funcBuilder->getBoolConstant(loc, ival.toBool());
Implement torch.kernel_call capture. * Had to stop short of modifying the function return signature because of a missing C-API upstream. * Committing here is good enough for a test and will resolve the various TODOs about upstream APIs next.
2020-10-06 14:21:21 +08:00
}
Add remaining pieces to capture full example models. * Adds Basicpy List, Tuple, Dict types and plumbs through C API. * Started debugging the issues around aten::conv2d capture, but a PyTorch bug is suspected. * Was able to manually verify that the basic conv2d forward test captures correctly with a workaround. * Need to resolve some printing issues upstream and move these tests to an integration test target (they take ~seconds to run).
2020-10-17 08:38:07 +08:00
if (ival.isList()) {
auto list = ival.toList();
llvm::SmallVector<MlirValue, 4> elements;
Work around various PyTorch issues in support of convolution. * Enables the conv2d fwd test and ResA (which are both small). * Deletes resnet18 and vgg, which both run but generate output that crashes FileCheck and lit (or at least makes them take an eternity).
2020-10-19 12:32:29 +08:00
for (IValue element : list) {
Add remaining pieces to capture full example models. * Adds Basicpy List, Tuple, Dict types and plumbs through C API. * Started debugging the issues around aten::conv2d capture, but a PyTorch bug is suspected. * Was able to manually verify that the basic conv2d forward test captures correctly with a workaround. * Need to resolve some printing issues upstream and move these tests to an integration test target (they take ~seconds to run).
2020-10-17 08:38:07 +08:00
elements.push_back(mapIValueToMlirValue(loc, element));
}
Fix insertion point bug #102 The current code was inserting all build_list ops after the last constant op since it was assuming that all elements being passed in were constants. This patch replaces that patch with a new function that inserts the build_list ops before the terminator. Also modifies test_export_conv2d_fwd.py since its output no longer matches. TEST: Added test_export_cat.py which is the code in #102
2020-11-03 07:30:21 +08:00
return funcBuilder->buildList(loc, elements);
Add remaining pieces to capture full example models. * Adds Basicpy List, Tuple, Dict types and plumbs through C API. * Started debugging the issues around aten::conv2d capture, but a PyTorch bug is suspected. * Was able to manually verify that the basic conv2d forward test captures correctly with a workaround. * Need to resolve some printing issues upstream and move these tests to an integration test target (they take ~seconds to run).
2020-10-17 08:38:07 +08:00
}
if (ival.isNone()) {
return funcBuilder->getNoneConstant(loc);
}
Support optional args/returns and other odds and ends. * None's out Device? args. * Emits bool tensors if needed. * Adds some stderr tracing to better see what is going on. * Test case that exercises NLLLoss. * This test case emits something for backward calculations but there are some issues still to be worked out, so that part is left out of the test case. * Progress on #97
2020-10-30 08:41:15 +08:00
if (ival.isDevice()) {
// TODO: Do we need to model/preserve device? Currently, just None'ing
// it out.
return funcBuilder->getNoneConstant(loc);
}
Implement torch.kernel_call capture. * Had to stop short of modifying the function return signature because of a missing C-API upstream. * Committing here is good enough for a test and will resolve the various TODOs about upstream APIs next.
2020-10-06 14:21:21 +08:00
return {nullptr};
// TODO: Implement mappings for the whole set (relevant to this use case):
// _(Tensor)
// _(Double)
// _(Int)
// _(Tuple)
// _(String)
// _(Blob)
// _(GenericList)
// _(GenericDict)
// _(Future)
// _(Device)
// _(Object)
// _(PyObject)
// _(Uninitialized)
// _(Capsule)
// _(RRef)
// _(Generator)
}
MlirType AcapController::mapIValueToMlirType(MlirLocation loc,
Work around various PyTorch issues in support of convolution. * Enables the conv2d fwd test and ResA (which are both small). * Deletes resnet18 and vgg, which both run but generate output that crashes FileCheck and lit (or at least makes them take an eternity).
2020-10-19 12:32:29 +08:00
const IValue &ival) {
Implement torch.kernel_call capture. * Had to stop short of modifying the function return signature because of a missing C-API upstream. * Committing here is good enough for a test and will resolve the various TODOs about upstream APIs next.
2020-10-06 14:21:21 +08:00
if (ival.isScalar()) {
return typeMapper.mapScalarType(ival.toScalar().type());
}
if (ival.isTensor()) {
return typeMapper.forwardTensorToType(ival.toTensor());
}
More progress on PyTorch acap device capture. * Now gets far enough to capture batch_norm. * Has some issues still with in-place ops. * Can materialize constants. * Includes an upgrade to PyTorch nightly, which has important bug fixes for fallback and boxed kernel dispatch. * Fixes #78, #79, #80. * Will do more testing in a follow-up once further bugs are fixed that facilitate getting at the other features.
2020-10-16 09:28:30 +08:00
if (ival.isBool()) {
// TODO: Switch to the numpy.bool type as that is a closer domain match.
return mlirIntegerTypeGet(funcBuilder->getContext(), 1);
}
Add remaining pieces to capture full example models. * Adds Basicpy List, Tuple, Dict types and plumbs through C API. * Started debugging the issues around aten::conv2d capture, but a PyTorch bug is suspected. * Was able to manually verify that the basic conv2d forward test captures correctly with a workaround. * Need to resolve some printing issues upstream and move these tests to an integration test target (they take ~seconds to run).
2020-10-17 08:38:07 +08:00
if (ival.isList()) {
return npcompListTypeGet(funcBuilder->getContext());
}
if (ival.isNone()) {
return npcompNoneTypeGet(funcBuilder->getContext());
}
Support optional args/returns and other odds and ends. * None's out Device? args. * Emits bool tensors if needed. * Adds some stderr tracing to better see what is going on. * Test case that exercises NLLLoss. * This test case emits something for backward calculations but there are some issues still to be worked out, so that part is left out of the test case. * Progress on #97
2020-10-30 08:41:15 +08:00
if (ival.isDevice()) {
return npcompNoneTypeGet(funcBuilder->getContext());
}
Implement torch.kernel_call capture. * Had to stop short of modifying the function return signature because of a missing C-API upstream. * Committing here is good enough for a test and will resolve the various TODOs about upstream APIs next.
2020-10-06 14:21:21 +08:00
return {nullptr};
}
More progress on PyTorch acap device capture. * Now gets far enough to capture batch_norm. * Has some issues still with in-place ops. * Can materialize constants. * Includes an upgrade to PyTorch nightly, which has important bug fixes for fallback and boxed kernel dispatch. * Fixes #78, #79, #80. * Will do more testing in a follow-up once further bugs are fixed that facilitate getting at the other features.
2020-10-16 09:28:30 +08:00
MlirValue AcapController::importTensorByValue(at::Tensor tensor) {
using at::ScalarType;
auto throwUnsupportedTensorError = [&]() {
std::stringstream msg;
msg << "Unsupported import tensor type: " << tensor;
throw std::invalid_argument(msg.str());
};
// Get a C-contiguous form as we can bulk-load that into a DenseElementsAttr.
if (!tensor.is_contiguous())
tensor = tensor.contiguous();
// The flat number of bytes throws an exception for tensors that are not
// dense and accessible as such.
at::checkLayout(at::CheckedFrom("accessing contiguous"), tensor,
c10::Layout::Strided);
// Construct the ShapedType.
auto loc = getCurrentLocation();
Support optional args/returns and other odds and ends. * None's out Device? args. * Emits bool tensors if needed. * Adds some stderr tracing to better see what is going on. * Test case that exercises NLLLoss. * This test case emits something for backward calculations but there are some issues still to be worked out, so that part is left out of the test case. * Progress on #97
2020-10-30 08:41:15 +08:00
MlirType elementType;
if (tensor.scalar_type() == ScalarType::Bool) {
// Bool is a special case. When used as an element type, it must be i1.
// The generalized (non-Tensor) conversion, assumes that Bool is the
// Basicpy bool type.
elementType = mlirIntegerTypeGet(funcBuilder->getContext(), 1);
} else {
elementType = typeMapper.mapScalarType(tensor.scalar_type());
}
More progress on PyTorch acap device capture. * Now gets far enough to capture batch_norm. * Has some issues still with in-place ops. * Can materialize constants. * Includes an upgrade to PyTorch nightly, which has important bug fixes for fallback and boxed kernel dispatch. * Fixes #78, #79, #80. * Will do more testing in a follow-up once further bugs are fixed that facilitate getting at the other features.
2020-10-16 09:28:30 +08:00
llvm::SmallVector<int64_t, 4> shape(tensor.sizes().begin(),
tensor.sizes().end());
MlirType shapedType = mlirRankedTensorTypeGetChecked(
shape.size(), shape.data(), elementType, loc);
if (mlirTypeIsNull(shapedType)) {
throwUnsupportedTensorError();
}
// Import DenseElementsAttr data.
// TODO: Support bool tensors.
// TODO: More import formats in C-API.
MlirAttribute valueAttribute;
auto numElements = tensor.numel();
auto tensorData = tensor.data_ptr();
switch (tensor.scalar_type()) {
case ScalarType::Int:
valueAttribute = mlirDenseElementsAttrInt32Get(
shapedType, numElements, static_cast<const int32_t *>(tensorData));
break;
case ScalarType::Long:
valueAttribute = mlirDenseElementsAttrInt64Get(
shapedType, numElements, static_cast<const int64_t *>(tensorData));
break;
case ScalarType::Float:
valueAttribute = mlirDenseElementsAttrFloatGet(
shapedType, numElements, static_cast<const float *>(tensorData));
break;
case ScalarType::Double:
valueAttribute = mlirDenseElementsAttrDoubleGet(
shapedType, numElements, static_cast<const double *>(tensorData));
break;
Support optional args/returns and other odds and ends. * None's out Device? args. * Emits bool tensors if needed. * Adds some stderr tracing to better see what is going on. * Test case that exercises NLLLoss. * This test case emits something for backward calculations but there are some issues still to be worked out, so that part is left out of the test case. * Progress on #97
2020-10-30 08:41:15 +08:00
case ScalarType::Bool:
// TODO: Add a test specifically for bool and ensure consistency between
// storage format and load format
// (https://github.com/llvm/mlir-npcomp/issues/100).
valueAttribute = mlirDenseElementsAttrBoolGet(
shapedType, numElements, static_cast<const int *>(tensorData));
break;
More progress on PyTorch acap device capture. * Now gets far enough to capture batch_norm. * Has some issues still with in-place ops. * Can materialize constants. * Includes an upgrade to PyTorch nightly, which has important bug fixes for fallback and boxed kernel dispatch. * Fixes #78, #79, #80. * Will do more testing in a follow-up once further bugs are fixed that facilitate getting at the other features.
2020-10-16 09:28:30 +08:00
default:
throwUnsupportedTensorError();
}
MlirValue constTensorValue =
funcBuilder->getGeneralConstant(loc, valueAttribute);
// Create an array from the tensor constant via the
// numpy.create_array_from_tensor op.
MlirType constArrayType = npcompNdArrayTypeGetFromShaped(shapedType);
MlirOperationState state =
mlirOperationStateGet("numpy.create_array_from_tensor", loc);
mlirOperationStateAddOperands(&state, 1, &constTensorValue);
mlirOperationStateAddResults(&state, 1, &constArrayType);
MlirOperation constArrayOp = mlirOperationCreate(&state);
funcBuilder->getEntryBlockBuilder().insertBeforeTerminator(constArrayOp);
MlirValue constArrayValue = mlirOperationGetResult(constArrayOp, 0);
funcBuilder->mapTensor(tensor, constArrayValue);
return constArrayValue;
}
Fix dispatch of arange. * Fixes #107 * I wouldn't say I love what had to be done here. Worth a conversation with the PT devs (probably as part of a rollup of a bunch of this stuff).
2020-11-13 09:42:19 +08:00
TORCH_LIBRARY_IMPL(aten, BackendSelect, m) {
// PyTorch logs a warning when kernels are overriden, which is unavoidable
// for factory-function BackendSelect kernels (there is not yet a "safe"
// override mechanism). So, just silence it. Any of them here are coded
// to be a superset of the default functionality, and there are only a few.
auto orig_log_level = FLAGS_caffe2_log_level;
FLAGS_caffe2_log_level = c10::GLOG_ERROR;
// Disable capture of arange: causes it to memorize the resulting tensor.
m.impl("arange", &AcapController::arangeBackendSelectKernel);
// Restore log level.
FLAGS_caffe2_log_level = orig_log_level;
}
Add boilerplate to do device capture (pytorch 1.6). * Uses the new dispatcher API. * Just prints to the console for the moment when an op is captured. * Executes the op through the existing implementation.
2020-09-26 09:13:16 +08:00
TORCH_LIBRARY_IMPL(_, ACAP_DISPATCH_KEY, m) {
m.fallback(torch::CppFunction::makeFromBoxedFunction<
&AcapController::fallbackKernel>());
}
Add remaining pieces to capture full example models. * Adds Basicpy List, Tuple, Dict types and plumbs through C API. * Started debugging the issues around aten::conv2d capture, but a PyTorch bug is suspected. * Was able to manually verify that the basic conv2d forward test captures correctly with a workaround. * Need to resolve some printing issues upstream and move these tests to an integration test target (they take ~seconds to run).
2020-10-17 08:38:07 +08:00
Capture backward conv and copy_ kernels. * This is sufficient to capture the forward and backward pass and gradients of a convolutional model with an nllloss. * As with the forward conv, the backward conv is a special case wrapped in an enigma on the PyTorch side. There aren't many like it, so special casing is just what we do. * When I traced this, I found that the copy_ op is not yet boxing compatible so I had to map it manually. If there are many more like this, I'll probably do something a bit more clever to reduce duplication. * This exposes new signature patterns that will need to be handled by the ATen lowering. Will take care of that next: It will be nice to have an e2e of a non-trivial case with full gradients. * Fixes #97.
2020-10-31 13:52:46 +08:00
TORCH_LIBRARY_IMPL(aten, ACAP_DISPATCH_KEY, m) {
m.impl("copy_", &AcapController::copyUnderKernel);
}
Work around various PyTorch issues in support of convolution. * Enables the conv2d fwd test and ResA (which are both small). * Deletes resnet18 and vgg, which both run but generate output that crashes FileCheck and lit (or at least makes them take an eternity).
2020-10-19 12:32:29 +08:00
TORCH_LIBRARY_IMPL(aten, ACAP_GRAD_DISPATCH_KEY, m) {
// The at::convolution op is special in several ways. First, it presently
// does not support boxing, so all of the usual fanciness does not apply
// and it cannot be intercepted by generic fallthroughs, which is what
// would usually allow us to avoid intercepting it at the gradient phase.
// Second, the default implementation (see
// aten/src/ATen/native/Convolution.cpp) is very switchy based on hard-coded
// assumptions about device type. If we do nothing here, we will at best
// intercept an mkldnn_convolution, cudnn_convolution, etc on the backend
// dispatch keys. Non standard backends that don't have these switches
// just route to aten::convolution_overrideable (see the else in
// aten::convolution) as a convenience, but that is mostly a pass-through
// (except for 3d convolutions which contain a trailing squeeze that needs
// special casing). Therefore, we just intercept the aten::convolution op,
// record it specially, and then mask ourselves off and ask the CPU backend
// to invoke it. Not awesome.
// Presumably this is on someone's list to adapt to the dispatch machinery
// in a more appropriate way, but as the core of what the framework is,
// perhaps people are reticent to touch it. Maybe someday, this can go away.
m.impl_UNBOXED("convolution", &AcapController::convolutionKernel);
Capture backward conv and copy_ kernels. * This is sufficient to capture the forward and backward pass and gradients of a convolutional model with an nllloss. * As with the forward conv, the backward conv is a special case wrapped in an enigma on the PyTorch side. There aren't many like it, so special casing is just what we do. * When I traced this, I found that the copy_ op is not yet boxing compatible so I had to map it manually. If there are many more like this, I'll probably do something a bit more clever to reduce duplication. * This exposes new signature patterns that will need to be handled by the ATen lowering. Will take care of that next: It will be nice to have an e2e of a non-trivial case with full gradients. * Fixes #97.
2020-10-31 13:52:46 +08:00
// Sadly, there is no easy intercept point for the backwards convolution
// kernel which allows for chaining to an existing backend. And convolution
// is exceptionally special cased in this way, moreso than other ops.
// The "solution" is to intercept the backend specific backward convolution
// ops, emit it with the signature of the more generic
// "convolution_backward_overrideable" op, which is available for generic
// backends, and twiddle the parameters needed to get back to that form.
// For MKL, which is effectively the CPU implementation, this just means that
// some parameters are swapped and the full generality is not supported.
// The "right" answer at some point is probably just to implement a
// convolution kernel that fully does what is needed and delegates to an
// appropriate implementation behind the scenes.
m.impl("mkldnn_convolution_backward", AcapController::mklConvolutionBackward);
Add remaining pieces to capture full example models. * Adds Basicpy List, Tuple, Dict types and plumbs through C API. * Started debugging the issues around aten::conv2d capture, but a PyTorch bug is suspected. * Was able to manually verify that the basic conv2d forward test captures correctly with a workaround. * Need to resolve some printing issues upstream and move these tests to an integration test target (they take ~seconds to run).
2020-10-17 08:38:07 +08:00
}