mirror of https://github.com/llvm/torch-mlir
97 lines
3.7 KiB
C++
97 lines
3.7 KiB
C++
//===- NumpyOps.cpp - Core numpy dialect ops --------------------*- 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 "npcomp/Dialect/Numpy/IR/NumpyOps.h"
|
|
#include "mlir/IR/Builders.h"
|
|
#include "mlir/IR/FunctionImplementation.h"
|
|
#include "mlir/IR/OpImplementation.h"
|
|
#include "mlir/IR/PatternMatch.h"
|
|
#include "npcomp/Dialect/Basicpy/IR/BasicpyDialect.h"
|
|
#include "npcomp/Dialect/Numpy/IR/NumpyDialect.h"
|
|
|
|
using namespace mlir;
|
|
using namespace mlir::NPCOMP;
|
|
using namespace mlir::NPCOMP::Numpy;
|
|
|
|
//----------------------------------------------------------------------------//
|
|
// Type inference
|
|
//----------------------------------------------------------------------------//
|
|
|
|
/// Adds constraints to relating a unary op that accepts and returns either
|
|
/// tensor or ndarray types where the dtype should be the same.
|
|
/// Type constraints are added on the dtype, not the outer object type.
|
|
static void constrainUnaryDtypeInvariantOp(Typing::CPA::Context &context,
|
|
Value source, Value dest,
|
|
Operation *op) {
|
|
auto &env = context.getCurrentEnvironment();
|
|
auto *sourceTn =
|
|
llvm::dyn_cast<Typing::CPA::ObjectValueType>(env.mapValueToType(source));
|
|
auto *destTn =
|
|
llvm::dyn_cast<Typing::CPA::ObjectValueType>(env.mapValueToType(dest));
|
|
if (sourceTn && destTn && sourceTn->getFieldCount() == 1 &&
|
|
destTn->getFieldCount() == 1) {
|
|
context.getConstraint(sourceTn->getFieldTypes().front(),
|
|
destTn->getFieldTypes().front());
|
|
}
|
|
}
|
|
|
|
void CreateArrayFromTensorOp::addCPAConstraints(Typing::CPA::Context &context) {
|
|
constrainUnaryDtypeInvariantOp(context, source(), dest(), *this);
|
|
}
|
|
|
|
void CopyToTensorOp::addCPAConstraints(Typing::CPA::Context &context) {
|
|
constrainUnaryDtypeInvariantOp(context, source(), dest(), *this);
|
|
}
|
|
|
|
void BuiltinUfuncCallOp::addCPAConstraints(Typing::CPA::Context &context) {
|
|
// TODO: This should really be a function call chosen so as to promote
|
|
// arguments. For now, though, we just say that the result is constrained
|
|
// to the inputs. Note that not all ufuncs transfer types like this.
|
|
// We just pretend this is two unary functions that write into the output.
|
|
for (auto input : inputs()) {
|
|
constrainUnaryDtypeInvariantOp(context, input, output(), *this);
|
|
}
|
|
}
|
|
|
|
//----------------------------------------------------------------------------//
|
|
// CreateArrayFromTensorOp
|
|
//----------------------------------------------------------------------------//
|
|
|
|
namespace {
|
|
/// Match create_array_from_tensor -> copy_to_tensor and elide in favor
|
|
/// of the original tensor.
|
|
class ElideCreateRedundantArrayFromTensor
|
|
: public OpRewritePattern<CopyToTensorOp> {
|
|
public:
|
|
using OpRewritePattern::OpRewritePattern;
|
|
LogicalResult matchAndRewrite(CopyToTensorOp op,
|
|
PatternRewriter &rewriter) const {
|
|
auto createArrayOp =
|
|
dyn_cast_or_null<CreateArrayFromTensorOp>(op.source().getDefiningOp());
|
|
if (createArrayOp && createArrayOp.dest().hasOneUse()) {
|
|
rewriter.replaceOp(op, createArrayOp.source());
|
|
}
|
|
return success();
|
|
}
|
|
};
|
|
} // namespace
|
|
|
|
void CopyToTensorOp::getCanonicalizationPatterns(
|
|
OwningRewritePatternList &patterns, MLIRContext *context) {
|
|
patterns.insert<ElideCreateRedundantArrayFromTensor>(context);
|
|
}
|
|
|
|
namespace mlir {
|
|
namespace NPCOMP {
|
|
namespace Numpy {
|
|
#define GET_OP_CLASSES
|
|
#include "npcomp/Dialect/Numpy/IR/NumpyOps.cpp.inc"
|
|
} // namespace Numpy
|
|
} // namespace NPCOMP
|
|
} // namespace mlir
|