mirror of https://github.com/llvm/torch-mlir
488 lines
19 KiB
C++
488 lines
19 KiB
C++
//===- RecognizeKernels.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
|
|
//
|
|
//===----------------------------------------------------------------------===//
|
|
|
|
#include "PassDetail.h"
|
|
|
|
#include "mlir/IR/PatternMatch.h"
|
|
#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
|
|
#include "npcomp/Dialect/ATen/IR/ATenDialect.h"
|
|
#include "npcomp/Dialect/ATen/Transforms/Passes.h"
|
|
#include "npcomp/Dialect/Basicpy/IR/BasicpyDialect.h"
|
|
#include "npcomp/Dialect/Numpy/IR/NumpyDialect.h"
|
|
#include "npcomp/Dialect/Numpy/IR/NumpyOps.h"
|
|
#include "npcomp/Dialect/Torch/IR/OpInterfaces.h"
|
|
#include "npcomp/Dialect/Torch/IR/TorchOps.h"
|
|
#include "llvm/ADT/StringMap.h"
|
|
#include "llvm/Support/Debug.h"
|
|
|
|
#define DEBUG_TYPE "aten-recognize-kernels"
|
|
|
|
using namespace mlir;
|
|
using namespace mlir::NPCOMP;
|
|
using namespace mlir::NPCOMP::aten;
|
|
using namespace mlir::NPCOMP::Torch;
|
|
|
|
namespace {
|
|
|
|
bool isTorchTensorType(StringRef torchType) {
|
|
return torchType == "Tensor" || torchType == "Tensor?";
|
|
}
|
|
|
|
bool isTorchOptionalType(StringRef torchType) {
|
|
return torchType.endswith("?");
|
|
}
|
|
|
|
struct TypeConversion {
|
|
Type targetType;
|
|
std::function<Value(Location loc, Value originalValue,
|
|
PatternRewriter &rewriter)>
|
|
callback;
|
|
};
|
|
|
|
/// Converts a Torch argument type so it is compatible with a target type
|
|
/// and flags. "Source" refers to the original operand value/type. "Target"
|
|
/// refers to the new op's operand value/type.
|
|
Optional<TypeConversion>
|
|
convertTorchArgType(StringRef sourceTorchType, StringRef targetTorchType,
|
|
KernelValueConversion::BitMask flag, Type sourceMlirType) {
|
|
using KVC = KernelValueConversion::BitMask;
|
|
// Default trivial case.
|
|
if (sourceTorchType == targetTorchType && flag == 0)
|
|
return TypeConversion{sourceMlirType, nullptr};
|
|
|
|
// Immutable tensor conversion.
|
|
if (flag & KVC::kImmutableTensor) {
|
|
// TODO: Support the kPromoteScalar flag.
|
|
if (!isTorchTensorType(sourceTorchType) ||
|
|
!isTorchTensorType(targetTorchType))
|
|
return None;
|
|
|
|
// If the target is optional and the type is NoneType, passthrough.
|
|
if (isTorchOptionalType(targetTorchType) &&
|
|
sourceMlirType.isa<Basicpy::NoneType>())
|
|
return TypeConversion{sourceMlirType, nullptr};
|
|
|
|
// Already immutable.
|
|
if (sourceMlirType.isa<TensorType>())
|
|
return TypeConversion{sourceMlirType, nullptr};
|
|
|
|
// Convert NdArray type.
|
|
if (auto ndArrayType = sourceMlirType.dyn_cast<Numpy::NdArrayType>()) {
|
|
auto tensorType = ndArrayType.toTensorType();
|
|
auto callback = [=](Location loc, Value originalValue,
|
|
PatternRewriter &rewriter) -> Value {
|
|
return rewriter.create<Numpy::CopyToTensorOp>(loc, tensorType,
|
|
originalValue);
|
|
};
|
|
return TypeConversion{tensorType, callback};
|
|
}
|
|
|
|
return None;
|
|
}
|
|
|
|
// TODO: Special case promotions and conversions.
|
|
return None;
|
|
}
|
|
|
|
/// Converts a Torch result type so it is compatible with the target type and
|
|
/// flags of a new op. "Source" refers to the original result value/type.
|
|
/// "Target" refers to the new ops's result value/type. The conversions
|
|
/// supported for results are, in general, much more constrained than those
|
|
/// supported for operands since these vary far less. The returned conversion
|
|
/// callback will convert from the target's type to the source's type.
|
|
Optional<TypeConversion>
|
|
convertTorchReturnType(StringRef sourceTorchType, StringRef targetTorchType,
|
|
KernelValueConversion::BitMask flag,
|
|
Type sourceMlirType) {
|
|
using KVC = KernelValueConversion::BitMask;
|
|
// Default trivial case.
|
|
if (sourceTorchType == targetTorchType && flag == 0) {
|
|
LLVM_DEBUG(llvm::dbgs() << " * Return types already match\n");
|
|
return TypeConversion{sourceMlirType, nullptr};
|
|
}
|
|
|
|
// Immutable tensor conversion.
|
|
if (flag & KVC::kImmutableTensor) {
|
|
LLVM_DEBUG(llvm::dbgs()
|
|
<< " * Return conversion flag kImmutableTensor\n");
|
|
if (!isTorchTensorType(sourceTorchType) ||
|
|
!isTorchTensorType(targetTorchType)) {
|
|
LLVM_DEBUG(llvm::dbgs()
|
|
<< " * Source or target not a Tensor type\n");
|
|
return None;
|
|
}
|
|
|
|
// Already immutable.
|
|
if (sourceMlirType.isa<TensorType>()) {
|
|
LLVM_DEBUG(llvm::dbgs() << " * Source is already immutable\n");
|
|
return TypeConversion{sourceMlirType, nullptr};
|
|
}
|
|
|
|
// Convert NdArray type.
|
|
if (sourceMlirType.isa<Basicpy::NoneType>() &&
|
|
isTorchOptionalType(targetTorchType)) {
|
|
LLVM_DEBUG(llvm::dbgs() << " * None Tensor type passthrough\n");
|
|
return TypeConversion{sourceMlirType, nullptr};
|
|
} else if (auto ndArrayType =
|
|
sourceMlirType.dyn_cast<Numpy::NdArrayType>()) {
|
|
auto tensorType = ndArrayType.toTensorType();
|
|
auto callback = [=](Location loc, Value newOpResultValue,
|
|
PatternRewriter &rewriter) -> Value {
|
|
return rewriter.create<Numpy::CreateArrayFromTensorOp>(
|
|
loc, ndArrayType, newOpResultValue);
|
|
};
|
|
LLVM_DEBUG(llvm::dbgs() << " * Convert return type\n");
|
|
return TypeConversion{tensorType, callback};
|
|
} else {
|
|
LLVM_DEBUG(llvm::dbgs()
|
|
<< " * Return type is not a supported tensor type\n");
|
|
return None;
|
|
}
|
|
}
|
|
|
|
LLVM_DEBUG(llvm::dbgs() << " * Return type conversion fallthrough\n");
|
|
return None;
|
|
}
|
|
|
|
/// Transforms from torch.kernel_call to recognized ops that implement the
|
|
/// TorchBuildableKernelOpInterface.
|
|
class KernelCallTransformer {
|
|
struct CandidateTransform;
|
|
|
|
public:
|
|
KernelCallTransformer(MLIRContext &context) : context(context) {}
|
|
|
|
template <typename DialectTy> void addDialectOps() {
|
|
Dialect *dialect = context.getOrLoadDialect<DialectTy>();
|
|
// TODO: This is not efficient. We should have a mechanism for dialects to
|
|
// track their own ops and allow a more fine grained mechanism.
|
|
auto allOps = context.getRegisteredOperations();
|
|
for (AbstractOperation *absOp : allOps) {
|
|
if (&absOp->dialect != dialect)
|
|
continue;
|
|
auto *concept = absOp->getInterface<TorchBuildableKernelOpInterface>();
|
|
if (!concept)
|
|
continue;
|
|
const BuildKernelMetadata &buildMetadata =
|
|
concept->getTorchBuildKernelMetadata();
|
|
addBuildableOperation(absOp->name, buildMetadata);
|
|
}
|
|
}
|
|
|
|
void addBuildableOperation(Identifier opName,
|
|
const BuildKernelMetadata &buildMetadata) {
|
|
LLVM_DEBUG(llvm::dbgs()
|
|
<< "Register kernel call translation for: " << opName << "\n");
|
|
{
|
|
CandidateTransformList &candidates =
|
|
kernelTransforms[buildMetadata.kernelName];
|
|
candidates.emplace_back(opName, buildMetadata);
|
|
}
|
|
|
|
for (StringRef aliasKernelName : buildMetadata.aliasKernelNames) {
|
|
CandidateTransformList &candidates = kernelTransforms[aliasKernelName];
|
|
candidates.emplace_back(opName, buildMetadata);
|
|
}
|
|
|
|
if (buildMetadata.inplaceVariantKernelName) {
|
|
CandidateTransformList &candidates =
|
|
kernelTransforms[*buildMetadata.inplaceVariantKernelName];
|
|
candidates.emplace_back(opName, buildMetadata);
|
|
}
|
|
}
|
|
|
|
LogicalResult transformKernelCall(KernelCallOp kernelCall,
|
|
PatternRewriter &rewriter) const {
|
|
StringRef kernelName = kernelCall.kernelName();
|
|
LLVM_DEBUG(llvm::dbgs()
|
|
<< "Evaluate kernel transform '" << kernelName << "':\n");
|
|
auto it = kernelTransforms.find(kernelName);
|
|
if (it == kernelTransforms.end()) {
|
|
LLVM_DEBUG(llvm::dbgs() << " - No candidate ops for kernel name\n");
|
|
return failure();
|
|
}
|
|
|
|
const CandidateTransformList &candidates = it->second;
|
|
for (const CandidateTransform &candidate : candidates) {
|
|
if (succeeded(rewriteForCandidateOp(kernelCall, candidate, rewriter)))
|
|
return success();
|
|
}
|
|
return failure();
|
|
}
|
|
|
|
LogicalResult rewriteForCandidateOp(KernelCallOp kernelCall,
|
|
const CandidateTransform &candidate,
|
|
PatternRewriter &rewriter) const {
|
|
LLVM_DEBUG(llvm::dbgs()
|
|
<< " * Evaluate op " << candidate.targetOpName << "\n");
|
|
Torch::KernelMetadata sourceMetadata = kernelCall.getTorchKernelMetadata();
|
|
// Fail on presently unsupported cases.
|
|
if (sourceMetadata.isVararg || candidate.buildMetadata.isVararg) {
|
|
LLVM_DEBUG(llvm::dbgs() << " - Skip candidate op: vararg kernels "
|
|
"presently not supported\n");
|
|
return failure();
|
|
}
|
|
if (sourceMetadata.isVarret || candidate.buildMetadata.isVarret) {
|
|
LLVM_DEBUG(llvm::dbgs() << " - Skip candidate op: varret kernels "
|
|
"presently not supported\n");
|
|
return failure();
|
|
}
|
|
|
|
// In none of the special forms do return arity mismatch.
|
|
if (sourceMetadata.returnTypes.size() !=
|
|
candidate.buildMetadata.returnTypes.size()) {
|
|
LLVM_DEBUG(llvm::dbgs()
|
|
<< " - Skip candidate op: return arity mismatch\n");
|
|
return failure();
|
|
}
|
|
|
|
bool sourceHasTrailingOutRef =
|
|
candidate.buildMetadata.promoteTrailingOutTensor &&
|
|
sourceMetadata.argTypes.size() ==
|
|
candidate.buildMetadata.argTypes.size() + 1;
|
|
if (sourceHasTrailingOutRef ||
|
|
sourceMetadata.argTypes.size() ==
|
|
candidate.buildMetadata.argTypes.size()) {
|
|
// Arg arity match.
|
|
LLVM_DEBUG(llvm::dbgs() << " + Candidate matches based on arity\n");
|
|
|
|
return rewriteMatchingArity(
|
|
kernelCall, sourceMetadata, candidate,
|
|
/*fixedArgArity=*/candidate.buildMetadata.argTypes.size(),
|
|
/*fixedRetArity=*/candidate.buildMetadata.returnTypes.size(),
|
|
/*sourceHasTrailingOutRef=*/sourceHasTrailingOutRef, rewriter);
|
|
}
|
|
|
|
return failure();
|
|
}
|
|
|
|
LogicalResult rewriteMatchingArity(KernelCallOp kernelCall,
|
|
const Torch::KernelMetadata sourceMetadata,
|
|
const CandidateTransform &candidate,
|
|
size_t fixedArgArity, size_t fixedRetArity,
|
|
bool sourceHasTrailingOutRef,
|
|
PatternRewriter &rewriter) const {
|
|
using KVC = KernelValueConversion::BitMask;
|
|
// Pre-conditions.
|
|
if (sourceHasTrailingOutRef)
|
|
assert((sourceMetadata.argTypes.size() ==
|
|
candidate.buildMetadata.argTypes.size() + 1) &&
|
|
"arg arity mismatch for trailing outref conversion");
|
|
else
|
|
assert(sourceMetadata.argTypes.size() ==
|
|
candidate.buildMetadata.argTypes.size() &&
|
|
"arg arity mismatch");
|
|
bool isInplaceVariant =
|
|
candidate.buildMetadata.inplaceVariantKernelName &&
|
|
kernelCall.kernelName() ==
|
|
*candidate.buildMetadata.inplaceVariantKernelName;
|
|
|
|
// Convert fixed return types.
|
|
using PostConversionCallback = std::function<void()>;
|
|
SmallVector<PostConversionCallback, 4> postConversionCallbacks;
|
|
struct ConversionInfo {
|
|
Value originalValue;
|
|
TypeConversion conversion;
|
|
};
|
|
SmallVector<Type, 4> resultTypes;
|
|
SmallVector<ConversionInfo, 4> resultConversions;
|
|
for (size_t i = 0; i < fixedRetArity; ++i) {
|
|
StringRef sourceTorchType = sourceMetadata.returnTypes[i];
|
|
StringRef targetTorchType = candidate.buildMetadata.returnTypes[i];
|
|
KVC flag = candidate.buildMetadata.getReturnConversion(i);
|
|
Value sourceValue = kernelCall.getResult(i);
|
|
Type sourceMlirType = kernelCall.getResultTypes()[i];
|
|
if (flag & KVC::kDropReturnAndAliasArg0) {
|
|
// Reduce result arity and alias any uses to arg0.
|
|
if (kernelCall.args().empty()) {
|
|
LLVM_DEBUG(llvm::dbgs()
|
|
<< " - Cannot alias arg0 (no arguments)\n");
|
|
return failure();
|
|
}
|
|
Value arg0 = kernelCall.args()[0];
|
|
postConversionCallbacks.push_back(
|
|
[sourceValue, arg0]() { sourceValue.replaceAllUsesWith(arg0); });
|
|
} else {
|
|
// General, arity-preserving type conversion.
|
|
auto conversion = convertTorchReturnType(
|
|
sourceTorchType, targetTorchType, flag, sourceMlirType);
|
|
if (!conversion) {
|
|
LLVM_DEBUG(llvm::dbgs()
|
|
<< " - Return type[" << i << "] incompatible: source="
|
|
<< sourceTorchType << ", target=" << targetTorchType
|
|
<< ", flag=" << flag << "\n");
|
|
return failure();
|
|
}
|
|
resultTypes.push_back(conversion->targetType);
|
|
resultConversions.push_back({sourceValue, std::move(*conversion)});
|
|
}
|
|
}
|
|
|
|
// Convert fixed arg types.
|
|
SmallVector<ConversionInfo, 4> operandInfos;
|
|
for (size_t i = 0, operandIndex = 0; i < fixedArgArity; ++i) {
|
|
// Drop this arg?
|
|
if (candidate.buildMetadata.argConversions[i] & KVC::kDrop)
|
|
continue;
|
|
if (kernelCall.getNumOperands() <= operandIndex) {
|
|
LLVM_DEBUG(llvm::dbgs()
|
|
<< " - Arg operand " << i
|
|
<< " does not exist in kernel call (missing default?)\n");
|
|
return failure();
|
|
}
|
|
|
|
// Normal type conversion of the operand.
|
|
operandInfos.emplace_back();
|
|
ConversionInfo &info = operandInfos.back();
|
|
info.originalValue = kernelCall.getOperand(operandIndex++);
|
|
Type sourceMlirType = info.originalValue.getType();
|
|
auto conversion = convertTorchArgType(
|
|
/*sourceTorchType=*/sourceMetadata.argTypes[i],
|
|
/*targetTorchType=*/candidate.buildMetadata.argTypes[i],
|
|
/*flag=*/candidate.buildMetadata.getArgConversion(i),
|
|
/*sourceMlirType=*/sourceMlirType);
|
|
if (!conversion) {
|
|
LLVM_DEBUG(llvm::dbgs()
|
|
<< " - Arg type[" << i
|
|
<< "] incompatible: source=" << sourceMetadata.argTypes[i]
|
|
<< ", target=" << candidate.buildMetadata.argTypes[i]
|
|
<< ", flag=" << candidate.buildMetadata.getArgConversion(i)
|
|
<< "\n");
|
|
return failure();
|
|
}
|
|
info.conversion = std::move(*conversion);
|
|
}
|
|
|
|
// Match criteria are satisfied. The IR can now be rewritten.
|
|
// Materialize conversions for operands.
|
|
SmallVector<Value, 4> operands;
|
|
for (ConversionInfo &info : operandInfos) {
|
|
if (!info.conversion.callback) {
|
|
// Identity conversion.
|
|
operands.push_back(info.originalValue);
|
|
} else {
|
|
// Computed conversion.
|
|
Value newValue = info.conversion.callback(kernelCall.getLoc(),
|
|
info.originalValue, rewriter);
|
|
operands.push_back(newValue);
|
|
}
|
|
}
|
|
|
|
// Create the op.
|
|
OperationState state(kernelCall.getLoc(), candidate.targetOpName, operands,
|
|
resultTypes, {});
|
|
Operation *newOp = rewriter.createOperation(state);
|
|
|
|
// Materialize conversions for results.
|
|
// For out params, we need to save off the converted first result -- we will
|
|
// just RAUW it with the out param later.
|
|
Value firstResultConverted;
|
|
for (auto it : llvm::enumerate(resultConversions)) {
|
|
ConversionInfo &info = it.value();
|
|
Value origOpResultValue = info.originalValue;
|
|
Value newOpResultValue = newOp->getOpResult(it.index());
|
|
Value convertedValue = newOpResultValue;
|
|
if (info.conversion.callback) {
|
|
// Conversion required.
|
|
convertedValue = info.conversion.callback(kernelCall.getLoc(),
|
|
newOpResultValue, rewriter);
|
|
}
|
|
origOpResultValue.replaceAllUsesWith(convertedValue);
|
|
if (it.index() == 0)
|
|
firstResultConverted = convertedValue;
|
|
}
|
|
|
|
// Post conversion callbacks.
|
|
for (auto &callback : postConversionCallbacks)
|
|
callback();
|
|
|
|
if (sourceHasTrailingOutRef || isInplaceVariant) {
|
|
assert(newOp->getNumResults() > 0 &&
|
|
newOp->getResultTypes()[0].isa<TensorType>() &&
|
|
"must have tensor first result");
|
|
LLVM_DEBUG(llvm::dbgs()
|
|
<< " - Ovewriting out param with result tensor.\n");
|
|
Value out;
|
|
if (sourceHasTrailingOutRef)
|
|
out = kernelCall.getOperand(fixedArgArity);
|
|
else // isInplaceVariant
|
|
out = kernelCall.getOperand(0);
|
|
rewriter.create<Numpy::OverwriteArrayOp>(kernelCall.getLoc(),
|
|
newOp->getResult(0), out);
|
|
assert(firstResultConverted && "must have a first result");
|
|
firstResultConverted.replaceAllUsesWith(out);
|
|
}
|
|
|
|
// Done.
|
|
rewriter.eraseOp(kernelCall);
|
|
return success();
|
|
}
|
|
|
|
private:
|
|
struct CandidateTransform {
|
|
CandidateTransform(Identifier targetOpName,
|
|
const BuildKernelMetadata &buildMetadata)
|
|
: targetOpName(targetOpName), buildMetadata(buildMetadata) {}
|
|
Identifier targetOpName;
|
|
const BuildKernelMetadata &buildMetadata;
|
|
};
|
|
using CandidateTransformList = SmallVector<CandidateTransform, 1>;
|
|
|
|
MLIRContext &context;
|
|
// Map of the torch.kernel_call op's kernelName() attribute to a list of
|
|
// candidate transforms describing how to map to a specific, recognized op.
|
|
// Note that a single generic kernel name can map to more than one candidate
|
|
// kernels based on signature. PyTorch has many such overloads that vary
|
|
// by signature. Some are handled transparently (such as special handling
|
|
// for variants with an out= parameter), while others may map differently
|
|
// (i.e. variants that operate on scalars vs tensors, have different arities,
|
|
// etc).
|
|
llvm::StringMap<CandidateTransformList> kernelTransforms;
|
|
};
|
|
|
|
class RecognizeOpPattern : public OpRewritePattern<Torch::KernelCallOp> {
|
|
public:
|
|
RecognizeOpPattern(MLIRContext *context,
|
|
const KernelCallTransformer &callTransformer)
|
|
: OpRewritePattern(context), callTransformer(callTransformer) {}
|
|
|
|
LogicalResult matchAndRewrite(Torch::KernelCallOp kernelCall,
|
|
PatternRewriter &rewriter) const override {
|
|
return callTransformer.transformKernelCall(kernelCall, rewriter);
|
|
}
|
|
|
|
private:
|
|
const KernelCallTransformer &callTransformer;
|
|
};
|
|
|
|
class ATenRecognizeKernelsPass
|
|
: public ATenRecognizeKernelsBase<ATenRecognizeKernelsPass> {
|
|
void getDependentDialects(DialectRegistry ®istry) const override {
|
|
registry.insert<ATenDialect>();
|
|
}
|
|
|
|
void runOnOperation() override {
|
|
MLIRContext *context = &getContext();
|
|
KernelCallTransformer transformer(*context);
|
|
transformer.addDialectOps<ATenDialect>();
|
|
|
|
RewritePatternSet patterns(context);
|
|
patterns.add<RecognizeOpPattern>(context, transformer);
|
|
if (failed(
|
|
applyPatternsAndFoldGreedily(getOperation(), std::move(patterns))))
|
|
signalPassFailure();
|
|
}
|
|
};
|
|
|
|
} // namespace
|
|
|
|
std::unique_ptr<OperationPass<FuncOp>>
|
|
mlir::NPCOMP::aten::createRecognizeKernelsPass() {
|
|
return std::make_unique<ATenRecognizeKernelsPass>();
|
|
}
|