torch-mlir/lib/Dialect/Torch/IR/TorchOps.cpp

183 lines
6.6 KiB
C++

//===----------------------------------------------------------------------===//
//
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
// See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
//
//===----------------------------------------------------------------------===//
#include "npcomp/Dialect/Torch/IR/TorchOps.h"
#include "mlir/IR/Builders.h"
#include "mlir/IR/BuiltinOps.h"
#include "npcomp/Dialect/Basicpy/IR/BasicpyDialect.h"
#include "npcomp/Dialect/Basicpy/IR/BasicpyOps.h"
#include "npcomp/Dialect/Numpy/IR/NumpyDialect.h"
#include "npcomp/Dialect/Numpy/IR/NumpyOps.h"
#include "llvm/ADT/StringMap.h"
using namespace mlir;
using namespace mlir::NPCOMP;
using namespace mlir::NPCOMP::Torch;
static SmallVector<StringRef, 4> strArrayAttrToVector(ArrayAttr array) {
SmallVector<StringRef, 4> strings;
strings.reserve(array.size());
for (auto stringAttr : array) {
strings.push_back(stringAttr.cast<StringAttr>().getValue());
}
return strings;
}
//===----------------------------------------------------------------------===//
// KernelCallOp
//===----------------------------------------------------------------------===//
KernelMetadata KernelCallOp::getTorchKernelMetadata() {
return KernelMetadata{
.kernelName = kernelName(),
.isVararg = sigIsVararg(),
.isVarret = sigIsVarret(),
.argTypes = strArrayAttrToVector(sigArgTypes()),
.returnTypes = strArrayAttrToVector(sigRetTypes()),
};
}
//===----------------------------------------------------------------------===//
// MethodOp
//===----------------------------------------------------------------------===//
LogicalResult MethodOp::verifySymbolUses(SymbolTableCollection &symbolTable) {
auto func = symbolTable.lookupNearestSymbolFrom<FuncOp>(*this, function());
if (!func)
return emitError() << "'@" << function()
<< "' does not reference a valid function";
if (func.getVisibility() != SymbolTable::Visibility::Private)
return emitError() << "'@" << function()
<< "' must reference a private function";
if (func.isDeclaration())
return emitError() << "'@" << function()
<< "' must reference a function that is defined (not "
"merely declared)";
auto expectedReceiverArgType = NnModuleType::get(
getContext(), getOperation()->getParentOfType<ClassTypeOp>().getName());
if (func.getType().getNumInputs() == 0 ||
func.getType().getInput(0) != expectedReceiverArgType) {
return emitError() << "the referenced function '" << function()
<< "' must have a first argument of type "
<< expectedReceiverArgType;
}
return success();
}
//===----------------------------------------------------------------------===//
// NnModuleOp
//===----------------------------------------------------------------------===//
static LogicalResult verify(NnModuleOp op) {
for (Operation &child : *op.getBody())
if (!isa<SlotOp, NnModuleTerminatorOp>(&child))
return child.emitOpError() << "is not allowed inside 'torch.nn_module'";
return success();
}
// PyTorch has a well-developed notion of subtyping.
//
// This is a restricted subset of it.
//
// TODO: Flesh this out.
bool isValidSubtype(Type subtype, Type type) {
if (subtype == type)
return true;
if (auto optional = type.dyn_cast<OptionalType>())
return subtype == optional.getContainedType() ||
subtype.isa<Basicpy::NoneType>();
return false;
}
LogicalResult NnModuleOp::verifySymbolUses(SymbolTableCollection &symbolTable) {
auto classType =
symbolTable.lookupNearestSymbolFrom<ClassTypeOp>(*this, getClassName());
if (!classType)
return emitError() << "'" << getClassName()
<< "' does not reference a valid class type";
auto attrs = llvm::to_vector<6>(getBody()->getOps<SlotOp>());
auto attrDefs = llvm::to_vector<6>(classType.getBody()->getOps<AttrOp>());
if (attrs.size() != attrDefs.size())
return emitError() << "number of 'torch.slot's in a 'torch.nn_module' must "
"match number of 'torch.attr's in "
"the corresponding 'torch.class_type'";
for (int i = 0, e = attrs.size(); i != e; i++) {
SlotOp attr = attrs[i];
AttrOp attrDef = attrDefs[i];
if (!isValidSubtype(attr.value().getType(), attrDef.type()) ||
attr.name() != attrDef.name()) {
return attr.emitOpError()
.append("is expected to match type and name of '",
attrDef.getOperation(), "'")
.attachNote(attrDef.getLoc())
.append("see torch.attr at corresponding index ", i, " here");
}
}
return success();
}
//===----------------------------------------------------------------------===//
// ClassTypeOp
//===----------------------------------------------------------------------===//
static LogicalResult verify(ClassTypeOp op) {
llvm::StringMap<Operation *> namesToOps;
for (Operation &child : op.getBody()->without_terminator()) {
if (!isa<AttrOp, MethodOp>(&child))
return child.emitOpError() << "is not allowed inside `torch.class_type`";
StringRef name;
if (auto attr = dyn_cast<AttrOp>(child))
name = attr.name();
else
name = cast<MethodOp>(child).name();
auto itAndWasInserted = namesToOps.insert({name, &child});
auto it = itAndWasInserted.first;
bool wasInserted = itAndWasInserted.second;
if (!wasInserted) {
auto diag = op.emitOpError().append(
"has duplicate attr/method with name '", name, "'");
diag.attachNote(it->second->getLoc())
.append("see first conflicting attr/method here");
diag.attachNote(child.getLoc())
.append("see second conflicting attr/method here");
return failure();
}
}
return success();
}
//===----------------------------------------------------------------------===//
// PrimLoopOp
//===----------------------------------------------------------------------===//
OperandRange PrimLoopOp::getSuccessorEntryOperands(unsigned index) {
assert(index == 0);
return iterArgsInit();
}
void PrimLoopOp::getSuccessorRegions(
Optional<unsigned> index, ArrayRef<Attribute> operands,
SmallVectorImpl<RegionSuccessor> &regions) {
(void)operands;
if (!index.hasValue()) {
regions.emplace_back(&region(), region().getArguments().slice(1));
return;
}
assert(*index == 0);
regions.emplace_back(&region(), region().getArguments().slice(1));
regions.emplace_back(getResults());
}
#define GET_OP_CLASSES
#include "npcomp/Dialect/Torch/IR/TorchOps.cpp.inc"