mirror of https://github.com/llvm/torch-mlir
241 lines
9.3 KiB
C++
241 lines
9.3 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/Dialect/StandardOps/IR/Ops.h"
|
|
#include "mlir/IR/Builders.h"
|
|
#include "mlir/IR/BuiltinOps.h"
|
|
#include "mlir/IR/PatternMatch.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;
|
|
|
|
//===----------------------------------------------------------------------===//
|
|
// 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> ®ions) {
|
|
(void)operands;
|
|
|
|
if (!index.hasValue()) {
|
|
regions.emplace_back(®ion(), region().getArguments().slice(1));
|
|
return;
|
|
}
|
|
assert(*index == 0);
|
|
regions.emplace_back(®ion(), region().getArguments().slice(1));
|
|
regions.emplace_back(getResults());
|
|
}
|
|
|
|
//===----------------------------------------------------------------------===//
|
|
// DerefineOp
|
|
//===----------------------------------------------------------------------===//
|
|
|
|
bool DerefineOp::areCastCompatible(mlir::TypeRange inputs,
|
|
mlir::TypeRange outputs) {
|
|
return isValidSubtype(inputs[0], outputs[0]);
|
|
}
|
|
|
|
void DerefineOp::getCanonicalizationPatterns(RewritePatternSet &patterns,
|
|
MLIRContext *context) {
|
|
patterns.add(+[](DerefineOp op, PatternRewriter &rewriter) {
|
|
// TODO: Extend RefineTypes for this case and delete this canonicalization,
|
|
// since we don't want control flow or calls to randomly block this fold
|
|
// (this canonicalization pattern makes the compiler brittle to control flow
|
|
// and calls).
|
|
bool allAllowRefinement =
|
|
llvm::all_of(op.getResult().getUsers(), allowsTypeRefinement);
|
|
if (!allAllowRefinement)
|
|
return failure();
|
|
rewriter.replaceOp(op, op.getOperand());
|
|
return success();
|
|
});
|
|
}
|
|
|
|
//===----------------------------------------------------------------------===//
|
|
// Aten__Is__Op
|
|
//===----------------------------------------------------------------------===//
|
|
|
|
OpFoldResult Aten__Is__Op::fold(ArrayRef<Attribute> operands) {
|
|
auto lhsType = self().getType();
|
|
auto rhsType = obj().getType();
|
|
// If either type is a NoneType, make it be the lhsType.
|
|
if (rhsType.isa<Basicpy::NoneType>())
|
|
std::swap(lhsType, rhsType);
|
|
// TODO: Implement and use subtype infra for this.
|
|
// If neither type is a subtype of the other, then the result is false.
|
|
if (lhsType.isa<Basicpy::NoneType>() && !rhsType.isa<Torch::OptionalType>())
|
|
return IntegerAttr::get(IntegerType::get(getContext(), 1), 0);
|
|
return nullptr;
|
|
}
|
|
|
|
//===----------------------------------------------------------------------===//
|
|
// AtenLenTOp
|
|
//===----------------------------------------------------------------------===//
|
|
|
|
void AtenLenTOp::getCanonicalizationPatterns(RewritePatternSet &patterns,
|
|
MLIRContext *context) {
|
|
patterns.add(+[](AtenLenTOp op, PatternRewriter &rewriter) {
|
|
auto buildList = op.getOperand().getDefiningOp<Basicpy::BuildListOp>();
|
|
if (!buildList)
|
|
return rewriter.notifyMatchFailure(op, "operand not basicpy.build_list");
|
|
rewriter.replaceOpWithNewOp<::mlir::ConstantOp>(
|
|
op, rewriter.getI64IntegerAttr(buildList.getNumOperands()));
|
|
return success();
|
|
});
|
|
}
|
|
|
|
//===----------------------------------------------------------------------===//
|
|
// AtenSizeOp
|
|
//===----------------------------------------------------------------------===//
|
|
|
|
void AtenSizeOp::getCanonicalizationPatterns(RewritePatternSet &patterns,
|
|
MLIRContext *context) {
|
|
patterns.add(+[](AtenSizeOp op, PatternRewriter &rewriter) {
|
|
auto type = op.getOperand().getType().dyn_cast<RankedTensorType>();
|
|
if (!type)
|
|
return rewriter.notifyMatchFailure(op, "not a ranked tensor");
|
|
SmallVector<Value> listElements;
|
|
for (int64_t size : type.getShape()) {
|
|
listElements.push_back(rewriter.create<::mlir::ConstantOp>(
|
|
op->getLoc(), rewriter.getI64IntegerAttr(size)));
|
|
}
|
|
rewriter.replaceOpWithNewOp<Basicpy::BuildListOp>(
|
|
op, Basicpy::ListType::get(rewriter.getContext()), listElements);
|
|
return success();
|
|
});
|
|
}
|
|
|
|
#define GET_OP_CLASSES
|
|
#include "npcomp/Dialect/Torch/IR/TorchOps.cpp.inc"
|