mirror of https://github.com/llvm/torch-mlir
395 lines
17 KiB
C++
395 lines
17 KiB
C++
//===- ReduceOpVariants.cpp --------------------------------------*- C++-*-===//
|
|
//
|
|
// This file is licensed 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
|
|
// Also available under a BSD-style license. See LICENSE.
|
|
//
|
|
//===----------------------------------------------------------------------===//
|
|
|
|
#include "PassDetail.h"
|
|
|
|
#include "ReifyAbstractInterpCalculationsUtils.h"
|
|
#include "mlir/Transforms/DialectConversion.h"
|
|
#include "torch-mlir/Dialect/Torch/IR/TorchOps.h"
|
|
#include "torch-mlir/Dialect/Torch/Transforms/Passes.h"
|
|
#include "llvm/ADT/StringExtras.h"
|
|
|
|
using namespace mlir;
|
|
using namespace mlir::torch;
|
|
using namespace mlir::torch::Torch;
|
|
|
|
// Create an overwrite in a manner that preserves the
|
|
// `OverwriteTensorContentsOp` invariant that both arguments
|
|
// must have the same shape and dtype.
|
|
static void createOverwriteTensorContents(PatternRewriter &rewriter,
|
|
Location loc, Value overwriterTensor,
|
|
Value overwrittenTensor) {
|
|
Type overwriterTensorType = overwriterTensor.getType();
|
|
Type overwrittenTensorType = overwrittenTensor.getType()
|
|
.dyn_cast<NonValueTensorType>()
|
|
.getWithValueSemantics();
|
|
if (overwriterTensorType != overwrittenTensorType) {
|
|
overwriterTensor = rewriter.create<TensorStaticInfoCastOp>(
|
|
loc, overwrittenTensorType, overwriterTensor);
|
|
}
|
|
rewriter.create<OverwriteTensorContentsOp>(loc, overwriterTensor,
|
|
overwrittenTensor);
|
|
}
|
|
|
|
static Type getContainerOrTensorTypeWithValueSemantics(Type type) {
|
|
if (auto optionalType = type.dyn_cast<OptionalType>()) {
|
|
Type newContainedType = getContainerOrTensorTypeWithValueSemantics(
|
|
optionalType.getContainedType());
|
|
return OptionalType::get(newContainedType);
|
|
} else if (auto listType = type.dyn_cast<ListType>()) {
|
|
Type newContainedType =
|
|
getContainerOrTensorTypeWithValueSemantics(listType.getContainedType());
|
|
return ListType::get(newContainedType);
|
|
} else if (auto tensorType = type.dyn_cast<NonValueTensorType>()) {
|
|
return tensorType.getWithValueSemantics();
|
|
} else {
|
|
return nullptr;
|
|
}
|
|
}
|
|
|
|
static bool
|
|
operatorOpHasValueSemantics(OperatorOp opOp,
|
|
std::optional<SymbolTable> extraLibrary) {
|
|
if (!extraLibrary.has_value())
|
|
return false;
|
|
auto opName = opOp->getAttr("name").cast<StringAttr>().getValue();
|
|
std::string libFuncName = (mlir::torch::Torch::getLibraryFunctionPrefix(
|
|
LibraryFunctionKind::HasValueSemantics) +
|
|
Twine(opName))
|
|
.str();
|
|
auto libFunc = extraLibrary->lookup<func::FuncOp>(libFuncName);
|
|
return bool(libFunc);
|
|
}
|
|
|
|
namespace {
|
|
// Convert value semantic ops operating on mutable arrays to instead operate on
|
|
// immutable tensors.
|
|
class ConvertHasValueSemanticsOpsToValueTensors : public RewritePattern {
|
|
public:
|
|
ConvertHasValueSemanticsOpsToValueTensors(
|
|
MLIRContext *context, const std::optional<SymbolTable> &extraLibrary)
|
|
: RewritePattern(MatchAnyOpTypeTag(), /*benefit=*/1, context) {
|
|
this->extraLibrary = extraLibrary;
|
|
}
|
|
LogicalResult matchAndRewrite(Operation *op,
|
|
PatternRewriter &rewriter) const override {
|
|
if (isa<OperatorOp>(op)) {
|
|
if (!operatorOpHasValueSemantics(cast<OperatorOp>(op), extraLibrary)) {
|
|
return rewriter.notifyMatchFailure(op, "does not have value semantics");
|
|
}
|
|
} else if (!op->hasTrait<Torch::OpTrait::HasValueSemantics>()) {
|
|
return rewriter.notifyMatchFailure(op, "does not have value semantics");
|
|
}
|
|
|
|
rewriter.startOpModification(op);
|
|
// Convert all operands.
|
|
SmallVector<Value> newOperands;
|
|
for (OpOperand &opOperand : op->getOpOperands()) {
|
|
Type operandType = opOperand.get().getType();
|
|
if (operandType.isa<NonValueTensorType>()) {
|
|
opOperand.set(rewriter.create<CopyToValueTensorOp>(op->getLoc(),
|
|
opOperand.get()));
|
|
} else if (auto listType = operandType.dyn_cast<ListType>()) {
|
|
if (!(listType.getContainedType().isa<NonValueTensorType>() ||
|
|
listType.getContainedType().isa<OptionalType>()))
|
|
continue;
|
|
|
|
// Construct a new list whose elements are value tensors copied from
|
|
// the non-value tensors of the original list.
|
|
auto listConstruct =
|
|
opOperand.get().getDefiningOp<PrimListConstructOp>();
|
|
if (!listConstruct) {
|
|
rewriter.cancelOpModification(op);
|
|
return rewriter.notifyMatchFailure(
|
|
op, "unimplemented: list of non vtensor type not constructed "
|
|
"from list construct");
|
|
}
|
|
|
|
if (listConstruct.getElements().empty())
|
|
continue;
|
|
|
|
// TODO: Handle optional type in list type.
|
|
if (auto optionalType =
|
|
listType.getContainedType().dyn_cast<OptionalType>()) {
|
|
if (!llvm::all_of(listConstruct.getElements(), [](Value val) {
|
|
return val.getType().isa<NonValueTensorType, Torch::NoneType>();
|
|
})) {
|
|
rewriter.cancelOpModification(op);
|
|
return rewriter.notifyMatchFailure(
|
|
op, "unimplemented: list containing optional type is not "
|
|
"handled.");
|
|
}
|
|
}
|
|
|
|
auto newListElements = llvm::to_vector(llvm::map_range(
|
|
listConstruct.getElements(), [&](Value tensor) -> Value {
|
|
if (tensor.getType().isa<NonValueTensorType>()) {
|
|
return rewriter.create<CopyToValueTensorOp>(op->getLoc(),
|
|
tensor);
|
|
}
|
|
return tensor;
|
|
}));
|
|
|
|
Type newListType = getContainerOrTensorTypeWithValueSemantics(listType);
|
|
if (!newListType) {
|
|
rewriter.cancelOpModification(op);
|
|
return rewriter.notifyMatchFailure(
|
|
op, "Unable to convert list type to value semantics.");
|
|
}
|
|
opOperand.set(rewriter.create<PrimListConstructOp>(
|
|
op->getLoc(), newListType, newListElements));
|
|
} else if (auto optionalType = operandType.dyn_cast<OptionalType>()) {
|
|
// TODO: A more general way to handle the optional type is to
|
|
// introduce a `copy.to_optional_vtensor` op.
|
|
if (!optionalType.getContainedType().isa<NonValueTensorType>())
|
|
continue;
|
|
|
|
// Create a new optional value whose input is a value tensor copied
|
|
// from the non value tensor of the original optional value.
|
|
auto derefine = opOperand.get().getDefiningOp<DerefineOp>();
|
|
if (!derefine) {
|
|
rewriter.cancelOpModification(op);
|
|
return rewriter.notifyMatchFailure(
|
|
op, "unimplemented: optional of non vtensor type not from "
|
|
"derefine");
|
|
}
|
|
|
|
if (!derefine.getOperand().getType().isa<NonValueTensorType>())
|
|
continue;
|
|
auto newOperand = rewriter.create<CopyToValueTensorOp>(
|
|
op->getLoc(), derefine.getOperand());
|
|
opOperand.set(rewriter.create<DerefineOp>(
|
|
op->getLoc(), Torch::OptionalType::get(newOperand.getType()),
|
|
newOperand));
|
|
}
|
|
}
|
|
// Convert all results.
|
|
rewriter.setInsertionPointAfter(op);
|
|
for (Value result : op->getResults()) {
|
|
auto tensorType = result.getType().dyn_cast<NonValueTensorType>();
|
|
if (!tensorType)
|
|
continue;
|
|
result.setType(tensorType.getWithValueSemantics());
|
|
auto nonValueTensor =
|
|
rewriter.create<CopyToNonValueTensorOp>(op->getLoc(), result);
|
|
result.replaceAllUsesExcept(nonValueTensor, nonValueTensor);
|
|
}
|
|
rewriter.finalizeOpModification(op);
|
|
return success();
|
|
}
|
|
|
|
private:
|
|
std::optional<SymbolTable> extraLibrary;
|
|
};
|
|
} // namespace
|
|
|
|
// Reduce Ops without value semantics but the corresponding without trailing
|
|
// underscore variant doesn't exist.
|
|
namespace {
|
|
|
|
// int(ceil((end - start) / step))
|
|
Value calculateArangeResultNumElements(PatternRewriter &rewriter, Location loc,
|
|
Value start, Value end, Value step) {
|
|
Value sub = rewriter.create<AtenSubOp>(
|
|
loc, Torch::NumberType::get(rewriter.getContext()), end, start);
|
|
Value div = rewriter.create<AtenDivOp>(loc, sub, step);
|
|
return rewriter.create<AtenCeilFloatOp>(loc, div);
|
|
}
|
|
|
|
class ReduceNonValueSemanticOps : public RewritePattern {
|
|
public:
|
|
ReduceNonValueSemanticOps(MLIRContext *context)
|
|
: RewritePattern(MatchAnyOpTypeTag(), /*benefit=*/1, context) {}
|
|
LogicalResult matchAndRewrite(Operation *op,
|
|
PatternRewriter &rewriter) const override {
|
|
Location loc = op->getLoc();
|
|
MLIRContext *ctx = op->getContext();
|
|
if (isa<AtenBernoulli_FloatOp>(op)) {
|
|
Operation *newOp = rewriter.create<ValsemVariantAtenBernoulliFloatOp>(
|
|
loc, op->getResultTypes(), op->getOperands());
|
|
auto tensor =
|
|
rewriter.create<CopyToValueTensorOp>(loc, newOp->getResult(0));
|
|
createOverwriteTensorContents(rewriter, loc, tensor, op->getOperand(0));
|
|
rewriter.replaceOp(op, op->getOperand(0));
|
|
return success();
|
|
} else if (auto arangeOutOp = dyn_cast<AtenArangeStartOutOp>(op)) {
|
|
Value start = arangeOutOp.getStart();
|
|
Value end = arangeOutOp.getEnd();
|
|
Value step = arangeOutOp.getStep();
|
|
Value out = arangeOutOp.getOut();
|
|
|
|
// `overwrite.tensor.contents` cannot change the tensor shape,
|
|
// so `out` tensor should have same num_elements with result tensor.
|
|
// It means that we don't support code like:
|
|
// `x = torch.randn(12)`
|
|
// `y = torch.arange(13, out=x)`
|
|
Value resultNumElements =
|
|
calculateArangeResultNumElements(rewriter, loc, start, end, step);
|
|
Value outNumElements = rewriter.create<AtenNumelOp>(loc, out);
|
|
Value eqOrNot =
|
|
rewriter.create<AtenEqIntOp>(loc, resultNumElements, outNumElements);
|
|
rewriter.create<RuntimeAssertOp>(
|
|
loc, eqOrNot,
|
|
rewriter.getStringAttr("`out` tensor should have the same "
|
|
"num_elements with result tenosr"));
|
|
|
|
auto dtype = rewriter.create<PrimDtypeOp>(loc, out);
|
|
auto device = rewriter.create<PrimDeviceOp>(loc, out);
|
|
auto shape = rewriter.create<AtenSizeOp>(
|
|
loc, Torch::ListType::get(Torch::IntType::get(ctx)), out);
|
|
auto none = rewriter.create<ConstantNoneOp>(loc);
|
|
Value newArange = rewriter.create<AtenArangeStartStepOp>(
|
|
loc, arangeOutOp.getResult().getType(), start, end, step, dtype,
|
|
/*layout=*/none, device, /*pin_memory=*/none);
|
|
Value reshape = rewriter.create<AtenReshapeOp>(
|
|
loc, arangeOutOp.getResult().getType(), newArange, shape);
|
|
|
|
auto vtensor = rewriter.create<CopyToValueTensorOp>(loc, reshape);
|
|
createOverwriteTensorContents(rewriter, loc, vtensor, out);
|
|
rewriter.replaceOp(arangeOutOp, out);
|
|
return success();
|
|
} else {
|
|
return failure();
|
|
}
|
|
}
|
|
};
|
|
} // namespace
|
|
|
|
namespace {
|
|
// Reduce the "trailing underscore inplace variant" to the value semantic
|
|
// variant + an overwrite of the original "self" argument.
|
|
class ReduceTrailingUnderscoreInplaceVariant : public RewritePattern {
|
|
public:
|
|
ReduceTrailingUnderscoreInplaceVariant(MLIRContext *context)
|
|
: RewritePattern(MatchAnyOpTypeTag(), /*benefit=*/1, context) {}
|
|
LogicalResult matchAndRewrite(Operation *op,
|
|
PatternRewriter &rewriter) const override {
|
|
if (!op->hasTrait<Torch::OpTrait::IsTrailingUnderscoreInplaceVariant>())
|
|
return rewriter.notifyMatchFailure(op, "is not trailing_ variant");
|
|
|
|
SmallVector<StringRef> fragments;
|
|
llvm::SplitString(op->getName().getStringRef(), fragments, ".");
|
|
assert(fragments.size() >= 3 && fragments[2].ends_with("_") &&
|
|
"IsTrailingUnderscoreInplaceVariant incorrectly applied");
|
|
fragments[2] = fragments[2].drop_back();
|
|
std::string noUnderscoreName = llvm::join(fragments, ".");
|
|
|
|
OperationState state(op->getLoc(), noUnderscoreName);
|
|
state.addTypes(op->getResultTypes());
|
|
state.addOperands(op->getOperands());
|
|
state.addAttributes(op->getAttrDictionary().getValue());
|
|
// Note: No successors or regions. Torch JIT operators don't have any.
|
|
assert(op->getNumRegions() == 0 && op->getNumSuccessors() == 0 &&
|
|
"Torch JIT operators shouldn't have regions or successors");
|
|
|
|
Operation *newOp = rewriter.create(state);
|
|
// Note: need to convert result to first input's dtype because mix precision
|
|
// compute would result in different behaviors.
|
|
// For example:
|
|
// a = torch.randn(3, 3).half() # float16
|
|
// b = torch.randn(3, 3) # float32
|
|
// a += b # i.e. torch.ops.aten.add_(a, b), result is float16
|
|
// c = a + b # i.e. torch.ops.aten.add(a, b), result is float32
|
|
Value none = rewriter.create<ConstantNoneOp>(op->getLoc());
|
|
Value cstFalse = rewriter.create<ConstantBoolOp>(op->getLoc(), false);
|
|
auto aDtype = rewriter.create<PrimDtypeOp>(op->getLoc(), op->getOperand(0));
|
|
auto toDtype = rewriter.create<AtenToDtypeOp>(
|
|
op->getLoc(), newOp->getResult(0).getType(), newOp->getResult(0),
|
|
aDtype, /*non_blocking=*/cstFalse, /*copy=*/cstFalse,
|
|
/*memory_format=*/none);
|
|
auto tensor = rewriter.create<CopyToValueTensorOp>(op->getLoc(), toDtype);
|
|
createOverwriteTensorContents(rewriter, op->getLoc(), tensor,
|
|
op->getOperand(0));
|
|
rewriter.replaceOp(op, op->getOperand(0));
|
|
|
|
return success();
|
|
}
|
|
};
|
|
} // namespace
|
|
|
|
static LogicalResult
|
|
reduceNonValueTensorLiteralOpToValueTensorLiteralOp(NonValueTensorLiteralOp op,
|
|
PatternRewriter &rewriter) {
|
|
Value valueTensor =
|
|
rewriter.create<ValueTensorLiteralOp>(op->getLoc(), op.getValue());
|
|
Value tensor =
|
|
copyTensorToType(rewriter, op->getLoc(), op.getType(), valueTensor);
|
|
rewriter.replaceOp(op, {tensor});
|
|
return success();
|
|
}
|
|
|
|
namespace {
|
|
struct ReduceOpVariantsPass
|
|
: public ReduceOpVariantsBase<ReduceOpVariantsPass> {
|
|
ReduceOpVariantsPass() = default;
|
|
ReduceOpVariantsPass(StringRef extraLibrary) {
|
|
this->extraLibrary = extraLibrary.str();
|
|
}
|
|
void runOnOperation() override {
|
|
MLIRContext *context = &getContext();
|
|
RewritePatternSet patterns(context);
|
|
OwningOpRef<ModuleOp> extraLibraryModule =
|
|
ModuleOp::create(UnknownLoc::get(context));
|
|
std::optional<SymbolTable> extraLibraryModuleSymTable = std::nullopt;
|
|
if (!extraLibrary.empty()) {
|
|
if (failed(loadExtraLibrary(extraLibrary, extraLibraryModule))) {
|
|
emitError(getOperation()->getLoc(),
|
|
"Failed to load extra-library file at " + extraLibrary);
|
|
return signalPassFailure();
|
|
}
|
|
|
|
extraLibraryModuleSymTable =
|
|
SymbolTable(extraLibraryModule->getOperation());
|
|
}
|
|
patterns.add<ConvertHasValueSemanticsOpsToValueTensors>(
|
|
context, extraLibraryModuleSymTable);
|
|
patterns.add<ReduceTrailingUnderscoreInplaceVariant>(context);
|
|
patterns.add(reduceNonValueTensorLiteralOpToValueTensorLiteralOp);
|
|
patterns.add<ReduceNonValueSemanticOps>(context);
|
|
|
|
ConversionTarget target(*context);
|
|
target.addIllegalOp<NonValueTensorLiteralOp>();
|
|
target.addIllegalOp<AtenBernoulli_FloatOp>();
|
|
target.addIllegalOp<AtenArangeStartOutOp>();
|
|
target.markUnknownOpDynamicallyLegal([&extraLibraryModuleSymTable](
|
|
Operation *op) {
|
|
if (op->hasTrait<Torch::OpTrait::HasValueSemantics>() ||
|
|
(isa<OperatorOp>(op) &&
|
|
operatorOpHasValueSemantics(cast<OperatorOp>(op),
|
|
extraLibraryModuleSymTable))) {
|
|
auto hasValueSemantics = [](Type t) {
|
|
// TODO: Make this an allowlist based on a closed torch dialect
|
|
// type system.
|
|
if (auto tensorType = t.dyn_cast<NonValueTensorType>()) {
|
|
return false;
|
|
}
|
|
return true;
|
|
};
|
|
return llvm::all_of(op->getOperandTypes(), hasValueSemantics) &&
|
|
llvm::all_of(op->getResultTypes(), hasValueSemantics);
|
|
}
|
|
if (op->hasTrait<Torch::OpTrait::IsTrailingUnderscoreInplaceVariant>()) {
|
|
return false;
|
|
}
|
|
return true;
|
|
});
|
|
|
|
if (failed(applyPartialConversion(getOperation(), target,
|
|
std::move(patterns)))) {
|
|
return signalPassFailure();
|
|
}
|
|
}
|
|
};
|
|
} // namespace
|
|
|
|
std::unique_ptr<OperationPass<func::FuncOp>>
|
|
mlir::torch::Torch::createReduceOpVariantsPass(StringRef extraLibrary) {
|
|
return std::make_unique<ReduceOpVariantsPass>(extraLibrary);
|
|
}
|