mirror of https://github.com/llvm/torch-mlir
Add ATen Dialect (#16)
This patch adds a dialect intended to be used as a frontend dialect to facilitate lowering from "A Tensor Library" in torch/pytorch. This patch includes several passes that are useful in conjuction with the dialect: --aten-layer-name: Generates layer names for each operation, which are not present in the original pytorch. --aten-to-std: Lower the ATen dialect into standard dialect function calls. --return-elimination-pass: convert functions (primarily the toplevel function) to pass return values by reference. This simplifies pytorch integration. --aten-op-report: generate a textual report about the model --liveness-report Future patches will implement actual integration with the pytorch jit to intercept and generates MLIR in this dialect, then lower the resulting MLIR into function calls through aten-layer-name -> aten-to-std -> return-elimination -> std-to-llvm. The result would then jitted using the LLVM jit, linked against a runtime library which makes calls back into pytorch to implement all the layers. Co-authored-by: Jeff Fifield <jeff.fifield@xilinx.com> Co-authored-by: Jeff Fifield <jeff.fifield@xilinx.com>pull/31/head
stephenneuendorffer
GitHub
parent
14f614396d
commit
bb668e6e26
|
|
@ -0,0 +1,160 @@
|
|||
//===- ATen.td ---------------------------------------------*- tablegen -*-===//
|
||||
//
|
||||
// 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 "mlir/IR/OpBase.td"
|
||||
|
||||
#ifndef ATEN_OPS
|
||||
#define ATEN_OPS
|
||||
|
||||
include "mlir/Interfaces/SideEffectInterfaces.td"
|
||||
include "npcomp/Dialect/ATen/ATenOpInterface.td"
|
||||
|
||||
def aten_Dialect : Dialect {
|
||||
let name = "aten";
|
||||
let cppNamespace = "aten";
|
||||
}
|
||||
|
||||
// TODO: convert to "let results =" style
|
||||
|
||||
class aten_Op<string mnemonic, list<OpTrait> traits = [StatisticsOpInterface]> :
|
||||
Op<aten_Dialect, mnemonic, traits>;
|
||||
|
||||
|
||||
// Most ops are automatically generated from pytorch specs.
|
||||
include "npcomp/Dialect/ATen/ATenOps.td"
|
||||
|
||||
|
||||
def aten_BatchNormOp: aten_Op<"batch_norm", [NoSideEffect, StatisticsOpInterface]>,
|
||||
Results<(outs AnyTensor:$output, AnyTensor:$save_mean, AnyTensor:$save_invstd)> {
|
||||
let arguments = (
|
||||
ins AnyType:$arg0,
|
||||
AnyType:$arg1,
|
||||
AnyType:$arg2,
|
||||
AnyType:$arg3,
|
||||
AnyType:$arg4,
|
||||
AnyType:$arg5,
|
||||
AnyType:$arg6,
|
||||
AnyType:$arg7,
|
||||
AnyType:$arg8
|
||||
);
|
||||
|
||||
let summary = "BatchNorm operator";
|
||||
let description = [{
|
||||
BatchNorm operator
|
||||
}];
|
||||
let extraClassDeclaration = [{
|
||||
std::map<std::string, uint64_t> getStatistics();
|
||||
}];
|
||||
}
|
||||
|
||||
// We have list constants, which come out of pytorch. Represent them using
|
||||
// our own constant-like type, which gets lowered to std_ConstantOp later.
|
||||
def aten_ConstantOp: aten_Op<"constant", [NoSideEffect]>,
|
||||
Results<(outs AnyType)> {
|
||||
let summary = "Constant operator";
|
||||
let description = [{
|
||||
Constant operator
|
||||
}];
|
||||
|
||||
}
|
||||
|
||||
// Our jit library only supports 6 argument convolutions, rather than 9
|
||||
// arguments supported by pytorch. This operation allows us to represent this
|
||||
// limitation temporarily.
|
||||
def aten_ConvolutionOp: aten_Op<"_convolution", [NoSideEffect, StatisticsOpInterface]>,
|
||||
Results<(outs AnyTensor)> {
|
||||
let arguments = (
|
||||
ins AnyTensor:$input,
|
||||
AnyTensor:$weight,
|
||||
AnyTensor:$bias,
|
||||
AnyType:$stride,
|
||||
AnyType:$padding,
|
||||
AnyType:$dilation
|
||||
);
|
||||
|
||||
let summary = "Convolution operator";
|
||||
let description = [{
|
||||
Convolution operator
|
||||
}];
|
||||
let extraClassDeclaration = [{
|
||||
std::map<std::string, uint64_t> getStatistics();
|
||||
uint64_t getOperandTransferVolume(unsigned int idx, bool read);
|
||||
uint64_t getResultTransferVolume(unsigned int idx, bool read);
|
||||
}];
|
||||
}
|
||||
|
||||
// Our jit library only supports 6 argument convolutions, rather than 9
|
||||
// arguments supported by pytorch. This operation allows us to represent this
|
||||
// limitation temporarily.
|
||||
def aten_ConvolutionBackwardOp: aten_Op<"_convolution_backward", [NoSideEffect, StatisticsOpInterface]>,
|
||||
Results<(outs AnyTensor:$dx, AnyTensor:$dw, AnyTensor:$db)> {
|
||||
let arguments = (
|
||||
ins AnyTensor:$grad_output,
|
||||
AnyTensor:$input,
|
||||
AnyTensor:$weight,
|
||||
AnyType:$stride,
|
||||
AnyType:$padding,
|
||||
AnyType:$dilation
|
||||
);
|
||||
|
||||
let summary = "ConvolutionBackward operator";
|
||||
let description = [{
|
||||
ConvolutionBackward operator
|
||||
}];
|
||||
let extraClassDeclaration = [{
|
||||
std::map<std::string, uint64_t> getStatistics();
|
||||
}];
|
||||
}
|
||||
|
||||
|
||||
def aten_FlattenOp: aten_Op<"flatten", [NoSideEffect, StatisticsOpInterface]>,
|
||||
Results<(outs AnyTensor)> {
|
||||
let arguments = (
|
||||
ins AnyType:$arg0,
|
||||
AnyType:$arg1,
|
||||
AnyType:$arg2
|
||||
);
|
||||
|
||||
let summary = "Flatten operator";
|
||||
let description = [{
|
||||
Flatten operator
|
||||
}];
|
||||
let extraClassDeclaration = [{
|
||||
std::map<std::string, uint64_t> getStatistics();
|
||||
}];
|
||||
}
|
||||
|
||||
def aten_MaxPool2dOp: aten_Op<"max_pool2d", [NoSideEffect, StatisticsOpInterface]>,
|
||||
Results<(outs AnyTensor)> {
|
||||
let arguments = (
|
||||
ins AnyType:$arg0,
|
||||
AnyType:$arg1,
|
||||
AnyType:$arg2,
|
||||
AnyType:$arg3,
|
||||
AnyType:$arg4,
|
||||
AnyType:$arg5
|
||||
);
|
||||
|
||||
let summary = "MaxPool2d operator";
|
||||
let description = [{
|
||||
MaxPool2d operator
|
||||
}];
|
||||
let extraClassDeclaration = [{
|
||||
std::map<std::string, uint64_t> getStatistics();
|
||||
}];
|
||||
}
|
||||
|
||||
def aten_TypeCastOp : aten_Op<"type_cast", [NoSideEffect]>,
|
||||
Results<(outs AnyType)> {
|
||||
let summary = "TypeCast operator";
|
||||
let arguments = (
|
||||
ins AnyType:$x
|
||||
);
|
||||
}
|
||||
|
||||
#endif
|
||||
|
|
@ -0,0 +1,131 @@
|
|||
//===- ATenDialect.h --------------------------------------------*- 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
|
||||
//
|
||||
//===----------------------------------------------------------------------===//
|
||||
|
||||
#ifndef NPCOMP_DIALECT_ATEN_DIALECT_H
|
||||
#define NPCOMP_DIALECT_ATEN_DIALECT_H
|
||||
|
||||
#include "mlir/IR/Builders.h"
|
||||
#include "mlir/IR/Dialect.h"
|
||||
#include "mlir/IR/Function.h"
|
||||
#include "mlir/IR/OpDefinition.h"
|
||||
#include "mlir/IR/OpImplementation.h"
|
||||
#include "mlir/IR/StandardTypes.h"
|
||||
#include "mlir/IR/TypeSupport.h"
|
||||
#include "mlir/IR/Types.h"
|
||||
#include "mlir/Interfaces/SideEffectInterfaces.h"
|
||||
|
||||
#include <map>
|
||||
|
||||
namespace mlir {
|
||||
namespace NPCOMP {
|
||||
namespace aten {
|
||||
|
||||
/// The ATenDialect models 'A Tensor library' from Pytorch. The intention
|
||||
/// is to provide an abstraction which is isomorphic with datastructures
|
||||
/// returned from the pytorch jit, enabling integration with Pytorch models.
|
||||
/// Most of the actual operation definitions in tablegen are themselves
|
||||
/// generated from C APIs exported by Pytorch.
|
||||
class ATenDialect : public mlir::Dialect {
|
||||
public:
|
||||
explicit ATenDialect(mlir::MLIRContext *ctx);
|
||||
static StringRef getDialectNamespace() { return "aten"; }
|
||||
|
||||
/// Parse a type registered to this dialect. Overridding this method is
|
||||
/// required for dialects that have custom types.
|
||||
/// Technically this is only needed to be able to round-trip to textual IR.
|
||||
mlir::Type parseType(DialectAsmParser &parser) const override;
|
||||
|
||||
/// Print a type registered to this dialect. Overridding this method is
|
||||
/// only required for dialects that have custom types.
|
||||
/// Technically this is only needed to be able to round-trip to textual IR.
|
||||
void printType(mlir::Type type, DialectAsmPrinter &os) const override;
|
||||
};
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
/////////////////////// Custom Types for the Dialect ///////////////////////////
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
namespace detail {
|
||||
struct ATenListTypeStorage;
|
||||
}
|
||||
|
||||
/// LLVM-style RTTI: one entry per subclass to allow dyn_cast/isa.
|
||||
enum ATenTypeKind {
|
||||
// The enum starts at the range reserved for this dialect.
|
||||
ATEN_TYPE = mlir::Type::FIRST_PRIVATE_EXPERIMENTAL_0_TYPE,
|
||||
ATEN_LIST,
|
||||
};
|
||||
|
||||
/// A variadic list of arguments in ATen
|
||||
class ATenListType : public mlir::Type::TypeBase<ATenListType, mlir::Type,
|
||||
detail::ATenListTypeStorage> {
|
||||
public:
|
||||
using Base::Base;
|
||||
|
||||
/// Return the type of individual elements in the array.
|
||||
mlir::Type getElementType();
|
||||
|
||||
/// Get the unique instance of this Type from the context.
|
||||
static ATenListType get(mlir::Type elementType);
|
||||
|
||||
/// Support method to enable LLVM-style RTTI type casting.
|
||||
static bool kindof(unsigned kind) { return kind == ATenTypeKind::ATEN_LIST; }
|
||||
};
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
//////////////////// Custom Operations for the Dialect /////////////////////////
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
namespace {
|
||||
|
||||
// Return the tensor volume (i.e., the number of elements) of the given shaped
|
||||
// type. If the type does not have a rank, return 1. If the type doesn't
|
||||
// have a static shape, return 0.
|
||||
uint64_t getTensorVolume(const ShapedType ty) {
|
||||
if (!ty.hasRank())
|
||||
return 1;
|
||||
|
||||
if (!ty.hasStaticShape())
|
||||
return 0;
|
||||
|
||||
uint64_t volume = 1;
|
||||
for (auto &d : ty.getShape())
|
||||
volume *= d;
|
||||
return volume;
|
||||
}
|
||||
|
||||
// Return the tensor volume (i.e., the number of elements) of the given type.
|
||||
// If the type doesn't have a shape, return 1. If the type is shaped, but
|
||||
// does not have a rank, return 1. If the type is shaped, but doesn't have a
|
||||
// static shape, return 0.
|
||||
uint64_t getTensorVolume(const Type ty) {
|
||||
if (auto t = ty.dyn_cast<ShapedType>()) {
|
||||
return getTensorVolume(t);
|
||||
} else {
|
||||
return 1;
|
||||
}
|
||||
}
|
||||
|
||||
} // namespace
|
||||
} // namespace aten
|
||||
} // namespace NPCOMP
|
||||
} // namespace mlir
|
||||
|
||||
#include "npcomp/Dialect/ATen/ATenOpInterfaces.h"
|
||||
|
||||
namespace mlir {
|
||||
namespace NPCOMP {
|
||||
namespace aten {
|
||||
// include TableGen generated Op definitions
|
||||
#define GET_OP_CLASSES
|
||||
#include "npcomp/Dialect/ATen/ATen.h.inc"
|
||||
|
||||
} // namespace aten
|
||||
} // namespace NPCOMP
|
||||
} // namespace mlir
|
||||
|
||||
#endif
|
||||
|
|
@ -0,0 +1,29 @@
|
|||
//===- ATenLayerNamePass.h --------------------------------------*- 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
|
||||
//
|
||||
//===----------------------------------------------------------------------===//
|
||||
|
||||
#ifndef NPCOMP_DIALECT_ATEN_LAYERNAMEPASS_H
|
||||
#define NPCOMP_DIALECT_ATEN_LAYERNAMEPASS_H
|
||||
|
||||
#include <memory>
|
||||
|
||||
namespace mlir {
|
||||
class Pass;
|
||||
} // namespace mlir
|
||||
|
||||
namespace mlir {
|
||||
namespace NPCOMP {
|
||||
namespace aten {
|
||||
|
||||
std::unique_ptr<mlir::Pass> createATenLayerNamePass();
|
||||
void registerATenLayerNamePass();
|
||||
|
||||
} // namespace aten
|
||||
} // namespace NPCOMP
|
||||
} // namespace mlir
|
||||
|
||||
#endif
|
||||
|
|
@ -0,0 +1,29 @@
|
|||
//===- ATenLoweringPass.h ---------------------------------------*- 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
|
||||
//
|
||||
//===----------------------------------------------------------------------===//
|
||||
|
||||
#ifndef NPCOMP_DIALECT_ATEN_LOWERING_H
|
||||
#define NPCOMP_DIALECT_ATEN_LOWERING_H
|
||||
|
||||
#include <memory>
|
||||
|
||||
namespace mlir {
|
||||
class Pass;
|
||||
class DialectConversion;
|
||||
} // namespace mlir
|
||||
|
||||
namespace mlir {
|
||||
namespace NPCOMP {
|
||||
namespace aten {
|
||||
|
||||
std::unique_ptr<mlir::Pass> createATenLoweringPass();
|
||||
void registerATenLoweringPass();
|
||||
} // namespace aten
|
||||
} // namespace NPCOMP
|
||||
} // namespace mlir
|
||||
|
||||
#endif // NPCOMP_DIALECT_ATEN_LOWERING_H
|
||||
|
|
@ -0,0 +1,70 @@
|
|||
//===- ATenOpInterface.td ----------------------------------*- tablegen -*-===//
|
||||
//
|
||||
// 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 "mlir/IR/OpBase.td"
|
||||
|
||||
#ifndef ATEN_OP_INTERFACES
|
||||
#define ATEN_OP_INTERFACES
|
||||
|
||||
def StatisticsOpInterface : OpInterface<"StatisticsOpInterface"> {
|
||||
let description = [{
|
||||
This interface allows ops to expose a static operation profile,
|
||||
describing the computational behavior of their function.
|
||||
}];
|
||||
|
||||
let methods = [
|
||||
InterfaceMethod<[{
|
||||
Return statistics about the compute requirements of an op. The return
|
||||
value maps an arbitrary set of statistic names to an integer value.
|
||||
Users are currently expected to accept any statistic names and statistic
|
||||
names are arbitrary for different operations. In the future this
|
||||
interface could be expanded and standardized.
|
||||
}],
|
||||
"std::map<std::string, uint64_t>", "getStatistics"
|
||||
>,
|
||||
|
||||
InterfaceMethod<[{
|
||||
Return memory transfer requirements of the operation for the
|
||||
operand with the given index.
|
||||
}],
|
||||
"uint64_t", "getOperandTransferVolume",
|
||||
(ins "unsigned int":$idx, "bool":$read), /*methodBody=*/[{}], [{
|
||||
ConcreteOp *op = static_cast<ConcreteOp *>(this);
|
||||
if (!read) return 0;
|
||||
Value v = op->getODSOperands(idx).front();
|
||||
auto ty = v.getType();
|
||||
return mlir::NPCOMP::aten::getTensorVolume(ty);
|
||||
}]>,
|
||||
|
||||
InterfaceMethod<[{
|
||||
Return memory transfer requirements of the operation for the
|
||||
result with the given index.
|
||||
}],
|
||||
"uint64_t", "getResultTransferVolume",
|
||||
(ins "unsigned int":$idx, "bool":$write), /*methodBody=*/[{}], [{
|
||||
ConcreteOp *op = static_cast<ConcreteOp *>(this);
|
||||
if (!write) return 0;
|
||||
Value v = op->getODSResults(idx).front();
|
||||
auto ty = v.getType();
|
||||
return mlir::NPCOMP::aten::getTensorVolume(ty);
|
||||
}]>,
|
||||
|
||||
];
|
||||
}
|
||||
|
||||
def AnyScalarOrTensor : TypeConstraint<Or<[AnySignlessInteger.predicate,
|
||||
AnyFloat.predicate,
|
||||
AnyTensor.predicate]>,
|
||||
"scalar-or-tensor">;
|
||||
|
||||
def AnyScalar : TypeConstraint<Or<[AnySignlessInteger.predicate,
|
||||
AnyFloat.predicate]>,
|
||||
"scalar">;
|
||||
|
||||
|
||||
#endif
|
||||
|
|
@ -0,0 +1,20 @@
|
|||
//===- ATenOpInterfaces.h ---------------------------------------*- 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
|
||||
//
|
||||
//===----------------------------------------------------------------------===//
|
||||
|
||||
#ifndef NPCOMP_DIALECT_ATEN_OPINTERFACES_H
|
||||
#define NPCOMP_DIALECT_ATEN_OPINTERFACES_H
|
||||
|
||||
#include "mlir/IR/Types.h"
|
||||
|
||||
namespace mlir {
|
||||
namespace NPCOMP {
|
||||
#include "npcomp/Dialect/ATen/ATenOpInterfaces.h.inc"
|
||||
} // namespace aten
|
||||
} // namespace NPCOMP
|
||||
|
||||
#endif
|
||||
|
|
@ -0,0 +1,32 @@
|
|||
//===- ATenOpReport.h -------------------------------------------*- 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
|
||||
//
|
||||
//===----------------------------------------------------------------------===//
|
||||
|
||||
#ifndef NPCOMP_DIALECT_ATEN_OPREPORT_H
|
||||
#define NPCOMP_DIALECT_ATEN_OPREPORT_H
|
||||
|
||||
#include <memory>
|
||||
|
||||
namespace mlir {
|
||||
class Pass;
|
||||
} // namespace mlir
|
||||
|
||||
namespace mlir {
|
||||
namespace NPCOMP {
|
||||
namespace aten {
|
||||
|
||||
// Generate report on standard error.
|
||||
std::unique_ptr<mlir::Pass> createATenOpReportPass();
|
||||
// Return the report in the given output string.
|
||||
std::unique_ptr<mlir::Pass> createATenOpReportPass(std::string &output);
|
||||
void registerATenOpReportPass();
|
||||
|
||||
} // namespace aten
|
||||
} // namespace NPCOMP
|
||||
} // namespace mlir
|
||||
|
||||
#endif
|
||||
|
|
@ -0,0 +1,276 @@
|
|||
//===- ATenOpStatisticsUtils.h ----------------------------------*- 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
|
||||
//
|
||||
//===----------------------------------------------------------------------===//
|
||||
|
||||
#ifndef NPCOMP_DIALECT_ATEN_OPSTATISTICSUTILS_H
|
||||
#define NPCOMP_DIALECT_ATEN_OPSTATISTICSUTILS_H
|
||||
|
||||
#include "npcomp/Dialect/ATen/ATenDialect.h"
|
||||
|
||||
#include "llvm/Support/Debug.h"
|
||||
#include "mlir/IR/StandardTypes.h"
|
||||
#include "mlir/IR/Types.h"
|
||||
|
||||
#define DEBUG_TYPE "aten-op-stats"
|
||||
|
||||
/// This file generally contains utility methods that factor common code
|
||||
/// out from operations that implement Statisticsopinterface.
|
||||
|
||||
namespace mlir {
|
||||
namespace NPCOMP {
|
||||
namespace aten {
|
||||
|
||||
// Return the op statistics for conv2d-like operations.
|
||||
template <class T> std::map<std::string, uint64_t> getConv2dStatistics(T *o, uint64_t groups) {
|
||||
|
||||
std::map<std::string, uint64_t> toReturn;
|
||||
|
||||
TensorType resultTy = o->getResult().getType().template cast<TensorType>();
|
||||
TensorType inputTy = o->input().getType().template cast<TensorType>();
|
||||
TensorType weightTy = o->weight().getType().template cast<TensorType>();
|
||||
TensorType biasTy = o->bias().getType().template cast<TensorType>();
|
||||
|
||||
uint64_t ofm_volume = getTensorVolume(resultTy);
|
||||
uint64_t ofm_depth = resultTy.getShape()[1];
|
||||
|
||||
uint64_t ifm_depth = inputTy.getShape()[1];
|
||||
uint64_t kernel_height = weightTy.getShape()[2];
|
||||
uint64_t kernel_width = weightTy.getShape()[3];
|
||||
|
||||
// Number of forward MACs per pixel =
|
||||
// kernel_width * kernel_height * ifm_depth / groups
|
||||
uint64_t MACs_per_OFM = (ifm_depth / groups) * kernel_height * kernel_width;
|
||||
uint64_t total_MACs = ofm_volume * MACs_per_OFM;
|
||||
|
||||
uint64_t ifm_volume = getTensorVolume(inputTy);
|
||||
uint64_t weight_volume = getTensorVolume(weightTy);
|
||||
uint64_t bias_volume = getTensorVolume(biasTy);
|
||||
|
||||
// Should be gated on whether there is bias at all
|
||||
toReturn["ops:+"] = ofm_volume;
|
||||
toReturn["ops:MAC"] = total_MACs;
|
||||
|
||||
toReturn["operand:0:activation_in"] = ifm_volume;
|
||||
toReturn["result:0:activation_out"] = ofm_volume;
|
||||
toReturn["operand:1:parameters_in:weight"] = weight_volume;
|
||||
toReturn["operand:2:parameters_in:bias"] = bias_volume;
|
||||
|
||||
toReturn["reads"] = weight_volume + bias_volume + ifm_volume;
|
||||
toReturn["writes"] = ofm_volume;
|
||||
|
||||
return toReturn;
|
||||
}
|
||||
|
||||
// Return the op statistics for conv2dBackward-like operations.
|
||||
template<typename T>
|
||||
std::map<std::string, uint64_t> getConv2dBackwardStatistics(T op, uint64_t groups) {
|
||||
|
||||
std::map<std::string, uint64_t> toReturn;
|
||||
TensorType dx_out_resultTy = op.getResult(0).getType().template cast<TensorType>();
|
||||
uint64_t dx_out_volume = getTensorVolume(dx_out_resultTy);
|
||||
|
||||
TensorType weightTy = op.getOperand(2).getType().template cast<TensorType>();
|
||||
uint64_t weight_volume = getTensorVolume(weightTy);
|
||||
uint64_t loss_in_depth = weightTy.getShape()[0];
|
||||
uint64_t kernel_width = weightTy.getShape()[2];
|
||||
uint64_t kernel_height = weightTy.getShape()[3];
|
||||
|
||||
uint64_t MACs_per_loss =
|
||||
(loss_in_depth / groups) * kernel_height * kernel_width;
|
||||
|
||||
uint64_t total_MACs = dx_out_volume * MACs_per_loss;
|
||||
|
||||
TensorType ifmTy = op.getOperand(1).getType().template cast<TensorType>();
|
||||
uint64_t ifm_volume = getTensorVolume(ifmTy);
|
||||
auto ifm_shape = ifmTy.getShape();
|
||||
|
||||
uint64_t ifm_bwh = ifm_shape[0] * ifm_shape[2] *
|
||||
ifm_shape[3]; // Batch * height * width: the depth is in
|
||||
// the weight shape already
|
||||
total_MACs += ifm_bwh * weight_volume;
|
||||
|
||||
TensorType dx_inTy = op.getOperand(0).getType().template cast<TensorType>();
|
||||
uint64_t dx_in_volume = getTensorVolume(dx_inTy);
|
||||
toReturn["ops:+"] = dx_in_volume;
|
||||
|
||||
// Reads: Conv_backward reads 3 tensors: the loss in, the activation in and
|
||||
// the transposed weights
|
||||
toReturn["reads"] = dx_in_volume + ifm_volume + weight_volume;
|
||||
|
||||
// Writes: Conv_backward writes 3 tensors: the loss out, gradients for the
|
||||
// weights, and gradients for the biases
|
||||
TensorType biasTy = op.getResult(2).getType().template cast<TensorType>();
|
||||
uint64_t bias_volume = getTensorVolume(biasTy);
|
||||
toReturn["writes"] = dx_out_volume + weight_volume + bias_volume;
|
||||
|
||||
toReturn["ops:MAC"] = total_MACs;
|
||||
toReturn["operand:0:activation_in"] = dx_in_volume;
|
||||
toReturn["operand:1:activation_in"] = ifm_volume;
|
||||
toReturn["operand:2:parameters_in:weight"] = weight_volume;
|
||||
|
||||
toReturn["result:0:grad:dx"] = dx_out_volume;
|
||||
toReturn["result:1:grad:dw"] = weight_volume;
|
||||
toReturn["result:2:grad:db"] = bias_volume;
|
||||
|
||||
return toReturn;
|
||||
}
|
||||
|
||||
|
||||
// Return a model of the number of bytes needed to represent the operand of
|
||||
// the given convolution-like operation with the given index. The shape is
|
||||
// assumed to be in NCHW order with a simple tiled model of data reuse. TODO:
|
||||
// convert this to a target-specific interface.
|
||||
template <class T>
|
||||
uint64_t getConv2dOperandTransferVolume(T *o, unsigned int idx, bool read) {
|
||||
|
||||
if (!read)
|
||||
return 0;
|
||||
|
||||
double vol = getTensorVolume(o->getOperand(idx).getType());
|
||||
|
||||
TensorType inputTy = o->input().getType().template cast<TensorType>();
|
||||
TensorType weightTy = o->weight().getType().template cast<TensorType>();
|
||||
TensorType resultTy = o->getResult().getType().template cast<TensorType>();
|
||||
|
||||
float filter_width = weightTy.getShape()[2];
|
||||
float filter_height = weightTy.getShape()[3];
|
||||
|
||||
float batch_sw = inputTy.getShape()[0];
|
||||
float ifm_depth_sw = inputTy.getShape()[1];
|
||||
float ih = inputTy.getShape()[2];
|
||||
float iw = inputTy.getShape()[3];
|
||||
|
||||
float ofm_depth_sw = resultTy.getShape()[1];
|
||||
|
||||
const float batch_hw = 4;
|
||||
const float ifm_depth_hw = 32;
|
||||
const float ofm_depth_hw = 32;
|
||||
|
||||
const float ifm_tile_height = 4;
|
||||
const float ifm_tile_width = 4;
|
||||
const float ofm_tile_height = 4;
|
||||
const float ofm_tile_width = 4;
|
||||
|
||||
float ifm_aperture = ifm_tile_height - ceilf(filter_height / 2.0f);
|
||||
float ifm_overlap = ceilf(filter_height / 2.0f);
|
||||
|
||||
float bl = ceilf(batch_sw / batch_hw);
|
||||
float ol = ceilf(ofm_depth_sw / ofm_depth_hw);
|
||||
float il = ceilf(ifm_depth_sw / ifm_depth_hw);
|
||||
|
||||
float ifm_overhead = 1.0f;
|
||||
float weight_overhead = 1.0f;
|
||||
if (filter_width > 1) {
|
||||
ifm_overhead =
|
||||
ol * ifm_tile_height * ((ih - ifm_overlap) / (ih * ifm_aperture));
|
||||
weight_overhead = bl;
|
||||
} else {
|
||||
ifm_overhead = ol;
|
||||
}
|
||||
|
||||
if (idx == 0) {
|
||||
LLVM_DEBUG(llvm::outs() << "ifm_overhead:" << ifm_overhead << "\n");
|
||||
return vol * ifm_overhead;
|
||||
}
|
||||
if (idx == 1) {
|
||||
LLVM_DEBUG(llvm::outs() << "weight_overhead:" << weight_overhead << "\n");
|
||||
return vol * weight_overhead;
|
||||
}
|
||||
return vol;
|
||||
}
|
||||
|
||||
// Return a model of the number of bytes needed to represent the result of
|
||||
// the given convolution-like operation with the given index. The shape is
|
||||
// assumed to be in NCHW order with a simple tiled model of data reuse. TODO:
|
||||
// convert this to a target-specific interface.
|
||||
template <class T>
|
||||
uint64_t getConv2dResultTransferVolume(T *o, unsigned int idx, bool write) {
|
||||
|
||||
TensorType inputTy = o->input().getType().template cast<TensorType>();
|
||||
TensorType resultTy = o->getResult().getType().template cast<TensorType>();
|
||||
TensorType weightTy = o->weight().getType().template cast<TensorType>();
|
||||
float filter_width = weightTy.getShape()[2];
|
||||
// float filter_height = weightTy.getShape()[3];
|
||||
|
||||
float ifm_depth_sw = inputTy.getShape()[1];
|
||||
const float ifm_depth_hw = 32;
|
||||
|
||||
float il = ceilf(ifm_depth_sw / ifm_depth_hw);
|
||||
|
||||
float write_output_overhead = 1.0f;
|
||||
float read_output_cost = 1.0f;
|
||||
|
||||
if (filter_width > 1) {
|
||||
write_output_overhead = il;
|
||||
read_output_cost = il;
|
||||
}
|
||||
|
||||
double vol = getTensorVolume(resultTy);
|
||||
|
||||
if (write) {
|
||||
LLVM_DEBUG(llvm::outs()
|
||||
<< "write_output_overhead:" << write_output_overhead << "\n");
|
||||
return vol * write_output_overhead;
|
||||
} else {
|
||||
LLVM_DEBUG(llvm::outs() << "read_output_cost:" << read_output_cost << "\n");
|
||||
return vol * read_output_cost;
|
||||
}
|
||||
}
|
||||
|
||||
// Return the op statistics for matrixmultiply-like operations.
|
||||
template<typename T>
|
||||
std::map<std::string, uint64_t> getMMOpStatistics(T op) {
|
||||
|
||||
std::map<std::string, uint64_t> toReturn;
|
||||
|
||||
TensorType resultTy = op.getResult().getType().template cast<TensorType>();
|
||||
uint64_t ofm_volume = getTensorVolume(resultTy);
|
||||
|
||||
// Use the weight tensor to find the number of input neurons
|
||||
TensorType lossTy = op.getOperand(0).getType().template cast<TensorType>();
|
||||
TensorType weightTy = op.getOperand(1).getType().template cast<TensorType>();
|
||||
uint64_t num_input_neurons = weightTy.getShape()[0];
|
||||
uint64_t total_MACs = ofm_volume * num_input_neurons;
|
||||
toReturn["ops:MAC"] = total_MACs;
|
||||
|
||||
uint64_t loss_in_volume = getTensorVolume(lossTy);
|
||||
uint64_t weight_volume = getTensorVolume(weightTy);
|
||||
toReturn["reads"] = loss_in_volume + weight_volume;
|
||||
toReturn["writes"] = ofm_volume;
|
||||
|
||||
toReturn["operand:0:activation_in"] = loss_in_volume;
|
||||
toReturn["operand:1:activation_in"] = weight_volume;
|
||||
toReturn["result:0:activation_out"] = ofm_volume;
|
||||
return toReturn;
|
||||
}
|
||||
|
||||
// Return the op statistics for ReLU-like operations.
|
||||
template <typename T>
|
||||
std::map<std::string, uint64_t> getReLUOpStatistics(T op) {
|
||||
|
||||
std::map<std::string, uint64_t> toReturn;
|
||||
|
||||
TensorType inputTy = op.getOperand().getType().template cast<TensorType>();
|
||||
TensorType resultTy = op.getResult().getType().template cast<TensorType>();
|
||||
|
||||
uint64_t in_volume = getTensorVolume(inputTy);
|
||||
uint64_t out_volume = getTensorVolume(resultTy);
|
||||
|
||||
toReturn["operand:0:activation_in"] = in_volume;
|
||||
toReturn["result:0:activation_out"] = out_volume;
|
||||
toReturn["reads"] = in_volume;
|
||||
toReturn["writes"] = out_volume;
|
||||
toReturn["ops:>"] = out_volume;
|
||||
|
||||
return toReturn;
|
||||
}
|
||||
|
||||
} // namespace aten
|
||||
} // namespace NPCOMP
|
||||
} // namespace mlir
|
||||
|
||||
#endif
|
||||
|
|
@ -0,0 +1,733 @@
|
|||
// This file is (mostly) automatically generated. Please do not edit.
|
||||
//===- ATenOps.td ------------------------------------------*- tablegen -*-===//
|
||||
//
|
||||
// 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
|
||||
//
|
||||
//===----------------------------------------------------------------------===//
|
||||
|
||||
#ifndef ATEN_OP_DEFS
|
||||
#define ATEN_OP_DEFS
|
||||
|
||||
def aten_AddOp: aten_Op<"add", [NoSideEffect, StatisticsOpInterface]>,
|
||||
Results<(outs AnyTensor)> {
|
||||
let arguments = (
|
||||
ins AnyTensor:$self,
|
||||
AnyTensor:$other,
|
||||
AnyScalar:$alpha
|
||||
);
|
||||
let summary = "aten add operator";
|
||||
let description = [{
|
||||
AddOp
|
||||
aten add operator
|
||||
}];
|
||||
let extraClassDeclaration = [{
|
||||
std::map<std::string, uint64_t> getStatistics();
|
||||
}];
|
||||
}
|
||||
|
||||
def aten_AddUnderOp: aten_Op<"add_", [NoSideEffect, StatisticsOpInterface]>,
|
||||
Results<(outs AnyTensor)> {
|
||||
let arguments = (
|
||||
ins AnyTensor:$self,
|
||||
AnyTensor:$other,
|
||||
AnyScalar:$alpha
|
||||
);
|
||||
let summary = "aten add_ operator";
|
||||
let description = [{
|
||||
AddUnderOp
|
||||
aten add_ operator
|
||||
}];
|
||||
let extraClassDeclaration = [{
|
||||
std::map<std::string, uint64_t> getStatistics();
|
||||
}];
|
||||
}
|
||||
|
||||
def aten_AsStridedOp: aten_Op<"as_strided", [NoSideEffect, StatisticsOpInterface]>,
|
||||
Results<(outs AnyTensor)> {
|
||||
let arguments = (
|
||||
ins AnyTensor:$self,
|
||||
AnyType:$size,
|
||||
AnyType:$stride
|
||||
);
|
||||
let summary = "aten as_strided operator";
|
||||
let description = [{
|
||||
AsStridedOp
|
||||
aten as_strided operator
|
||||
}];
|
||||
let extraClassDeclaration = [{
|
||||
std::map<std::string, uint64_t> getStatistics();
|
||||
}];
|
||||
}
|
||||
|
||||
def aten_ConvolutionOverrideableOp: aten_Op<"convolution_overrideable", [NoSideEffect, StatisticsOpInterface]>,
|
||||
Results<(outs AnyTensor)> {
|
||||
let arguments = (
|
||||
ins AnyTensor:$input,
|
||||
AnyTensor:$weight,
|
||||
AnyTensor:$bias,
|
||||
AnyType:$stride,
|
||||
AnyType:$padding,
|
||||
AnyType:$dilation,
|
||||
AnyScalar:$transposed,
|
||||
AnyType:$output_padding,
|
||||
AnyScalar:$groups
|
||||
);
|
||||
let summary = "aten convolution_overrideable operator";
|
||||
let description = [{
|
||||
ConvolutionOverrideableOp
|
||||
aten convolution_overrideable operator
|
||||
}];
|
||||
let extraClassDeclaration = [{
|
||||
std::map<std::string, uint64_t> getStatistics();
|
||||
}];
|
||||
}
|
||||
|
||||
def aten_ConvolutionBackwardOverrideableOp: aten_Op<"convolution_backward_overrideable", [NoSideEffect, StatisticsOpInterface]>,
|
||||
Results<(outs AnyTensor, AnyTensor, AnyTensor)> {
|
||||
let arguments = (
|
||||
ins AnyTensor:$grad_output,
|
||||
AnyTensor:$input,
|
||||
AnyTensor:$weight,
|
||||
AnyType:$stride,
|
||||
AnyType:$padding,
|
||||
AnyType:$dilation,
|
||||
AnyScalar:$transposed,
|
||||
AnyType:$output_padding,
|
||||
AnyScalar:$groups
|
||||
);
|
||||
let summary = "aten convolution_backward_overrideable operator";
|
||||
let description = [{
|
||||
ConvolutionBackwardOverrideableOp
|
||||
aten convolution_backward_overrideable operator
|
||||
}];
|
||||
let extraClassDeclaration = [{
|
||||
std::map<std::string, uint64_t> getStatistics();
|
||||
}];
|
||||
}
|
||||
|
||||
def aten_DivOp: aten_Op<"div", [NoSideEffect, StatisticsOpInterface]>,
|
||||
Results<(outs AnyTensor)> {
|
||||
let arguments = (
|
||||
ins AnyTensor:$self,
|
||||
AnyTensor:$other
|
||||
);
|
||||
let summary = "aten div operator";
|
||||
let description = [{
|
||||
DivOp
|
||||
aten div operator
|
||||
}];
|
||||
let extraClassDeclaration = [{
|
||||
std::map<std::string, uint64_t> getStatistics();
|
||||
}];
|
||||
}
|
||||
|
||||
def aten_DivUnderOp: aten_Op<"div_", [NoSideEffect, StatisticsOpInterface]>,
|
||||
Results<(outs AnyTensor)> {
|
||||
let arguments = (
|
||||
ins AnyTensor:$self,
|
||||
AnyTensor:$other
|
||||
);
|
||||
let summary = "aten div_ operator";
|
||||
let description = [{
|
||||
DivUnderOp
|
||||
aten div_ operator
|
||||
}];
|
||||
let extraClassDeclaration = [{
|
||||
std::map<std::string, uint64_t> getStatistics();
|
||||
}];
|
||||
}
|
||||
|
||||
def aten_ExpandOp: aten_Op<"expand", [NoSideEffect, StatisticsOpInterface]>,
|
||||
Results<(outs AnyTensor)> {
|
||||
let arguments = (
|
||||
ins AnyTensor:$self,
|
||||
AnyType:$size,
|
||||
AnyScalar:$implicit
|
||||
);
|
||||
let summary = "aten expand operator";
|
||||
let description = [{
|
||||
ExpandOp
|
||||
aten expand operator
|
||||
}];
|
||||
let extraClassDeclaration = [{
|
||||
std::map<std::string, uint64_t> getStatistics();
|
||||
}];
|
||||
}
|
||||
|
||||
def aten_LogSoftmaxOp: aten_Op<"_log_softmax", [NoSideEffect]>,
|
||||
Results<(outs AnyTensor)> {
|
||||
let arguments = (
|
||||
ins AnyTensor:$self,
|
||||
AnyScalar:$dim,
|
||||
AnyScalar:$half_to_float
|
||||
);
|
||||
let summary = "aten _log_softmax operator";
|
||||
let description = [{
|
||||
LogSoftmaxOp
|
||||
aten _log_softmax operator
|
||||
}];
|
||||
}
|
||||
|
||||
def aten_LogSoftmaxBackwardDataOp: aten_Op<"_log_softmax_backward_data", [NoSideEffect]>,
|
||||
Results<(outs AnyTensor)> {
|
||||
let arguments = (
|
||||
ins AnyTensor:$grad_output,
|
||||
AnyTensor:$output,
|
||||
AnyScalar:$dim,
|
||||
AnyTensor:$self
|
||||
);
|
||||
let summary = "aten _log_softmax_backward_data operator";
|
||||
let description = [{
|
||||
LogSoftmaxBackwardDataOp
|
||||
aten _log_softmax_backward_data operator
|
||||
}];
|
||||
}
|
||||
|
||||
def aten_MeanOp: aten_Op<"mean", [NoSideEffect, StatisticsOpInterface]>,
|
||||
Results<(outs AnyTensor)> {
|
||||
let arguments = (
|
||||
ins AnyTensor:$self
|
||||
);
|
||||
let summary = "aten mean operator";
|
||||
let description = [{
|
||||
MeanOp
|
||||
aten mean operator
|
||||
}];
|
||||
let extraClassDeclaration = [{
|
||||
std::map<std::string, uint64_t> getStatistics();
|
||||
}];
|
||||
}
|
||||
|
||||
def aten_MmOp: aten_Op<"mm", [NoSideEffect, StatisticsOpInterface]>,
|
||||
Results<(outs AnyTensor)> {
|
||||
let arguments = (
|
||||
ins AnyTensor:$self,
|
||||
AnyTensor:$mat2
|
||||
);
|
||||
let summary = "aten mm operator";
|
||||
let description = [{
|
||||
MmOp
|
||||
aten mm operator
|
||||
}];
|
||||
let extraClassDeclaration = [{
|
||||
std::map<std::string, uint64_t> getStatistics();
|
||||
}];
|
||||
}
|
||||
|
||||
def aten_MulOp: aten_Op<"mul", [NoSideEffect, StatisticsOpInterface]>,
|
||||
Results<(outs AnyTensor)> {
|
||||
let arguments = (
|
||||
ins AnyTensor:$self,
|
||||
AnyTensor:$other
|
||||
);
|
||||
let summary = "aten mul operator";
|
||||
let description = [{
|
||||
MulOp
|
||||
aten mul operator
|
||||
}];
|
||||
let extraClassDeclaration = [{
|
||||
std::map<std::string, uint64_t> getStatistics();
|
||||
}];
|
||||
}
|
||||
|
||||
def aten_MulUnderOp: aten_Op<"mul_", [NoSideEffect, StatisticsOpInterface]>,
|
||||
Results<(outs AnyTensor)> {
|
||||
let arguments = (
|
||||
ins AnyTensor:$self,
|
||||
AnyTensor:$other
|
||||
);
|
||||
let summary = "aten mul_ operator";
|
||||
let description = [{
|
||||
MulUnderOp
|
||||
aten mul_ operator
|
||||
}];
|
||||
let extraClassDeclaration = [{
|
||||
std::map<std::string, uint64_t> getStatistics();
|
||||
}];
|
||||
}
|
||||
|
||||
def aten_NativeBatchNormOp: aten_Op<"native_batch_norm", [NoSideEffect, StatisticsOpInterface]>,
|
||||
// FIXME: not quite automatically generated names
|
||||
Results<(outs AnyTensor:$output, AnyTensor:$save_mean, AnyTensor:$save_invstd)> {
|
||||
let arguments = (
|
||||
ins AnyTensor:$input,
|
||||
AnyTensor:$weight,
|
||||
AnyTensor:$bias,
|
||||
AnyTensor:$running_mean,
|
||||
AnyTensor:$running_var,
|
||||
AnyScalar:$training,
|
||||
AnyScalar:$momentum,
|
||||
AnyScalar:$eps
|
||||
);
|
||||
let summary = "aten native_batch_norm operator";
|
||||
let description = [{
|
||||
NativeBatchNormOp
|
||||
aten native_batch_norm operator
|
||||
}];
|
||||
let extraClassDeclaration = [{
|
||||
std::map<std::string, uint64_t> getStatistics();
|
||||
}];
|
||||
}
|
||||
|
||||
def aten_NativeBatchNormBackwardOp: aten_Op<"native_batch_norm_backward", [NoSideEffect, StatisticsOpInterface]>,
|
||||
// FIXME: not quite automatically generated
|
||||
Results<(outs AnyTensor:$dx, AnyTensor:$dm, AnyTensor:$dv)> {
|
||||
let arguments = (
|
||||
ins AnyTensor:$grad_out,
|
||||
AnyTensor:$input,
|
||||
AnyTensor:$weight,
|
||||
AnyTensor:$running_mean,
|
||||
AnyTensor:$running_var,
|
||||
AnyTensor:$save_mean,
|
||||
AnyTensor:$save_invstd,
|
||||
AnyScalar:$train,
|
||||
AnyScalar:$eps
|
||||
);
|
||||
let summary = "aten native_batch_norm_backward operator";
|
||||
let description = [{
|
||||
NativeBatchNormBackwardOp
|
||||
aten native_batch_norm_backward operator
|
||||
}];
|
||||
let extraClassDeclaration = [{
|
||||
std::map<std::string, uint64_t> getStatistics();
|
||||
}];
|
||||
}
|
||||
|
||||
def aten_NegOp: aten_Op<"neg", [NoSideEffect, StatisticsOpInterface]>,
|
||||
Results<(outs AnyTensor)> {
|
||||
let arguments = (
|
||||
ins AnyTensor:$self
|
||||
);
|
||||
let summary = "aten neg operator";
|
||||
let description = [{
|
||||
NegOp
|
||||
aten neg operator
|
||||
}];
|
||||
let extraClassDeclaration = [{
|
||||
std::map<std::string, uint64_t> getStatistics();
|
||||
}];
|
||||
}
|
||||
|
||||
def aten_ReluOp: aten_Op<"relu", [NoSideEffect, StatisticsOpInterface]>,
|
||||
Results<(outs AnyTensor)> {
|
||||
let arguments = (
|
||||
ins AnyTensor:$self
|
||||
);
|
||||
let summary = "aten relu operator";
|
||||
let description = [{
|
||||
ReluOp
|
||||
aten relu operator
|
||||
}];
|
||||
let extraClassDeclaration = [{
|
||||
std::map<std::string, uint64_t> getStatistics();
|
||||
}];
|
||||
}
|
||||
|
||||
def aten_ReluUnderOp: aten_Op<"relu_", [NoSideEffect, StatisticsOpInterface]>,
|
||||
Results<(outs AnyTensor)> {
|
||||
let arguments = (
|
||||
ins AnyTensor:$self
|
||||
);
|
||||
let summary = "aten relu_ operator";
|
||||
let description = [{
|
||||
ReluUnderOp
|
||||
aten relu_ operator
|
||||
}];
|
||||
let extraClassDeclaration = [{
|
||||
std::map<std::string, uint64_t> getStatistics();
|
||||
}];
|
||||
}
|
||||
|
||||
def aten_SizeOp: aten_Op<"size", [NoSideEffect, StatisticsOpInterface]>,
|
||||
Results<(outs AnyTensor)> {
|
||||
let arguments = (
|
||||
ins AnyTensor:$self,
|
||||
AnyScalar:$dim
|
||||
);
|
||||
let summary = "aten size operator";
|
||||
let description = [{
|
||||
SizeOp
|
||||
aten size operator
|
||||
}];
|
||||
let extraClassDeclaration = [{
|
||||
std::map<std::string, uint64_t> getStatistics();
|
||||
}];
|
||||
}
|
||||
|
||||
def aten_SqueezeOp: aten_Op<"squeeze", [NoSideEffect, StatisticsOpInterface]>,
|
||||
Results<(outs AnyTensor)> {
|
||||
let arguments = (
|
||||
ins AnyTensor:$self,
|
||||
AnyScalar:$dim
|
||||
);
|
||||
let summary = "aten squeeze operator";
|
||||
let description = [{
|
||||
SqueezeOp
|
||||
aten squeeze operator
|
||||
}];
|
||||
let extraClassDeclaration = [{
|
||||
std::map<std::string, uint64_t> getStatistics();
|
||||
}];
|
||||
}
|
||||
|
||||
def aten_SumOp: aten_Op<"sum", [NoSideEffect, StatisticsOpInterface]>,
|
||||
Results<(outs AnyTensor)> {
|
||||
let arguments = (
|
||||
ins AnyTensor:$self,
|
||||
AnyType:$dim,
|
||||
AnyScalar:$keepdim
|
||||
);
|
||||
let summary = "aten sum operator";
|
||||
let description = [{
|
||||
SumOp
|
||||
aten sum operator
|
||||
}];
|
||||
let extraClassDeclaration = [{
|
||||
std::map<std::string, uint64_t> getStatistics();
|
||||
}];
|
||||
}
|
||||
|
||||
def aten_TOp: aten_Op<"t", [NoSideEffect, StatisticsOpInterface]>,
|
||||
Results<(outs AnyTensor)> {
|
||||
let arguments = (
|
||||
ins AnyTensor:$self
|
||||
);
|
||||
let summary = "aten t operator";
|
||||
let description = [{
|
||||
TOp
|
||||
aten t operator
|
||||
}];
|
||||
let extraClassDeclaration = [{
|
||||
std::map<std::string, uint64_t> getStatistics();
|
||||
}];
|
||||
}
|
||||
|
||||
def aten_ThresholdBackwardOp: aten_Op<"threshold_backward", [NoSideEffect, StatisticsOpInterface]>,
|
||||
Results<(outs AnyTensor)> {
|
||||
let arguments = (
|
||||
ins AnyTensor:$grad_output,
|
||||
AnyTensor:$self,
|
||||
AnyScalar:$threshold
|
||||
);
|
||||
let summary = "aten threshold_backward operator";
|
||||
let description = [{
|
||||
ThresholdBackwardOp
|
||||
aten threshold_backward operator
|
||||
}];
|
||||
let extraClassDeclaration = [{
|
||||
std::map<std::string, uint64_t> getStatistics();
|
||||
}];
|
||||
}
|
||||
|
||||
def aten_UnsqueezeOp: aten_Op<"unsqueeze", [NoSideEffect, StatisticsOpInterface]>,
|
||||
Results<(outs AnyTensor)> {
|
||||
let arguments = (
|
||||
ins AnyTensor:$self,
|
||||
AnyScalar:$dim
|
||||
);
|
||||
let summary = "aten unsqueeze operator";
|
||||
let description = [{
|
||||
UnsqueezeOp
|
||||
aten unsqueeze operator
|
||||
}];
|
||||
let extraClassDeclaration = [{
|
||||
std::map<std::string, uint64_t> getStatistics();
|
||||
}];
|
||||
}
|
||||
|
||||
def aten_SubOp: aten_Op<"sub", [NoSideEffect, StatisticsOpInterface]>,
|
||||
Results<(outs AnyTensor)> {
|
||||
let arguments = (
|
||||
ins AnyTensor:$self,
|
||||
AnyTensor:$other,
|
||||
AnyScalar:$alpha
|
||||
);
|
||||
let summary = "aten sub operator";
|
||||
let description = [{
|
||||
SubOp
|
||||
aten sub operator
|
||||
}];
|
||||
let extraClassDeclaration = [{
|
||||
std::map<std::string, uint64_t> getStatistics();
|
||||
}];
|
||||
}
|
||||
|
||||
def aten_SubUnderOp: aten_Op<"sub_", [NoSideEffect, StatisticsOpInterface]>,
|
||||
Results<(outs AnyTensor)> {
|
||||
let arguments = (
|
||||
ins AnyTensor:$self,
|
||||
AnyTensor:$other,
|
||||
AnyScalar:$alpha
|
||||
);
|
||||
let summary = "aten sub_ operator";
|
||||
let description = [{
|
||||
SubUnderOp
|
||||
aten sub_ operator
|
||||
}];
|
||||
let extraClassDeclaration = [{
|
||||
std::map<std::string, uint64_t> getStatistics();
|
||||
}];
|
||||
}
|
||||
|
||||
def aten_AddmmOp: aten_Op<"addmm", [NoSideEffect, StatisticsOpInterface]>,
|
||||
Results<(outs AnyTensor)> {
|
||||
let arguments = (
|
||||
ins AnyTensor:$self,
|
||||
AnyTensor:$mat1,
|
||||
AnyTensor:$mat2,
|
||||
AnyScalar:$beta,
|
||||
AnyScalar:$alpha
|
||||
);
|
||||
let summary = "aten addmm operator";
|
||||
let description = [{
|
||||
AddmmOp
|
||||
aten addmm operator
|
||||
}];
|
||||
let extraClassDeclaration = [{
|
||||
std::map<std::string, uint64_t> getStatistics();
|
||||
}];
|
||||
}
|
||||
|
||||
def aten_ViewOp: aten_Op<"view", [NoSideEffect, StatisticsOpInterface]>,
|
||||
Results<(outs AnyTensor)> {
|
||||
let arguments = (
|
||||
ins AnyTensor:$self,
|
||||
AnyType:$size
|
||||
);
|
||||
let summary = "aten view operator";
|
||||
let description = [{
|
||||
ViewOp
|
||||
aten view operator
|
||||
}];
|
||||
let extraClassDeclaration = [{
|
||||
std::map<std::string, uint64_t> getStatistics();
|
||||
}];
|
||||
}
|
||||
|
||||
def aten_GatherOp: aten_Op<"gather", [NoSideEffect, StatisticsOpInterface]>,
|
||||
Results<(outs AnyTensor)> {
|
||||
let arguments = (
|
||||
ins AnyTensor:$self,
|
||||
AnyScalar:$dim,
|
||||
AnyTensor:$index,
|
||||
AnyScalar:$sparse_grad
|
||||
);
|
||||
let summary = "aten gather operator";
|
||||
let description = [{
|
||||
GatherOp
|
||||
aten gather operator
|
||||
}];
|
||||
let extraClassDeclaration = [{
|
||||
std::map<std::string, uint64_t> getStatistics();
|
||||
}];
|
||||
}
|
||||
|
||||
def aten_NllLossForwardOp: aten_Op<"nll_loss_forward", [NoSideEffect, StatisticsOpInterface]>,
|
||||
Results<(outs AnyTensor, AnyTensor)> {
|
||||
let arguments = (
|
||||
ins AnyTensor:$self,
|
||||
AnyTensor:$target,
|
||||
AnyTensor:$weight,
|
||||
AnyScalar:$reduction,
|
||||
AnyScalar:$ignore_index
|
||||
);
|
||||
let summary = "aten nll_loss_forward operator";
|
||||
let description = [{
|
||||
NllLossForwardOp
|
||||
aten nll_loss_forward operator
|
||||
}];
|
||||
let extraClassDeclaration = [{
|
||||
std::map<std::string, uint64_t> getStatistics();
|
||||
}];
|
||||
}
|
||||
|
||||
def aten_NllLossBackwardOp: aten_Op<"nll_loss_backward", [NoSideEffect, StatisticsOpInterface]>,
|
||||
Results<(outs AnyTensor)> {
|
||||
let arguments = (
|
||||
ins AnyTensor:$grad_output,
|
||||
AnyTensor:$self,
|
||||
AnyTensor:$target,
|
||||
AnyTensor:$weight,
|
||||
AnyScalar:$reduction,
|
||||
AnyScalar:$ignore_index,
|
||||
AnyTensor:$total_weight
|
||||
);
|
||||
let summary = "aten nll_loss_backward operator";
|
||||
let description = [{
|
||||
NllLossBackwardOp
|
||||
aten nll_loss_backward operator
|
||||
}];
|
||||
let extraClassDeclaration = [{
|
||||
std::map<std::string, uint64_t> getStatistics();
|
||||
}];
|
||||
}
|
||||
|
||||
def aten_NllLoss2dForwardOp: aten_Op<"nll_loss2d_forward", [NoSideEffect, StatisticsOpInterface]>,
|
||||
Results<(outs AnyTensor, AnyTensor)> {
|
||||
let arguments = (
|
||||
ins AnyTensor:$self,
|
||||
AnyTensor:$target,
|
||||
AnyTensor:$weight,
|
||||
AnyScalar:$reduction,
|
||||
AnyScalar:$ignore_index
|
||||
);
|
||||
let summary = "aten nll_loss2d_forward operator";
|
||||
let description = [{
|
||||
NllLoss2dForwardOp
|
||||
aten nll_loss2d_forward operator
|
||||
}];
|
||||
let extraClassDeclaration = [{
|
||||
std::map<std::string, uint64_t> getStatistics();
|
||||
}];
|
||||
}
|
||||
|
||||
def aten_NllLoss2dBackwardOp: aten_Op<"nll_loss2d_backward", [NoSideEffect, StatisticsOpInterface]>,
|
||||
Results<(outs AnyTensor)> {
|
||||
let arguments = (
|
||||
ins AnyTensor:$grad_output,
|
||||
AnyTensor:$self,
|
||||
AnyTensor:$target,
|
||||
AnyTensor:$weight,
|
||||
AnyScalar:$reduction,
|
||||
AnyScalar:$ignore_index,
|
||||
AnyTensor:$total_weight
|
||||
);
|
||||
let summary = "aten nll_loss2d_backward operator";
|
||||
let description = [{
|
||||
NllLoss2dBackwardOp
|
||||
aten nll_loss2d_backward operator
|
||||
}];
|
||||
let extraClassDeclaration = [{
|
||||
std::map<std::string, uint64_t> getStatistics();
|
||||
}];
|
||||
}
|
||||
|
||||
def aten_HardtanhOp: aten_Op<"hardtanh", [NoSideEffect, StatisticsOpInterface]>,
|
||||
Results<(outs AnyTensor)> {
|
||||
let arguments = (
|
||||
ins AnyTensor:$self,
|
||||
AnyScalar:$min_val,
|
||||
AnyScalar:$max_val
|
||||
);
|
||||
let summary = "aten hardtanh operator";
|
||||
let description = [{
|
||||
HardtanhOp
|
||||
aten hardtanh operator
|
||||
}];
|
||||
let extraClassDeclaration = [{
|
||||
std::map<std::string, uint64_t> getStatistics();
|
||||
}];
|
||||
}
|
||||
|
||||
def aten_HardtanhBackwardOp: aten_Op<"hardtanh_backward", [NoSideEffect, StatisticsOpInterface]>,
|
||||
Results<(outs AnyTensor)> {
|
||||
let arguments = (
|
||||
ins AnyTensor:$grad_output,
|
||||
AnyTensor:$self,
|
||||
AnyScalar:$min_val,
|
||||
AnyScalar:$max_val
|
||||
);
|
||||
let summary = "aten hardtanh_backward operator";
|
||||
let description = [{
|
||||
HardtanhBackwardOp
|
||||
aten hardtanh_backward operator
|
||||
}];
|
||||
let extraClassDeclaration = [{
|
||||
std::map<std::string, uint64_t> getStatistics();
|
||||
}];
|
||||
}
|
||||
|
||||
def aten_HardtanhUnderOp: aten_Op<"hardtanh_", [NoSideEffect, StatisticsOpInterface]>,
|
||||
Results<(outs AnyTensor)> {
|
||||
let arguments = (
|
||||
ins AnyTensor:$self,
|
||||
AnyScalar:$min_val,
|
||||
AnyScalar:$max_val
|
||||
);
|
||||
let summary = "aten hardtanh_ operator";
|
||||
let description = [{
|
||||
HardtanhUnderOp
|
||||
aten hardtanh_ operator
|
||||
}];
|
||||
let extraClassDeclaration = [{
|
||||
std::map<std::string, uint64_t> getStatistics();
|
||||
}];
|
||||
}
|
||||
|
||||
def aten_AdaptiveAvgPool2dOp: aten_Op<"_adaptive_avg_pool2d", [NoSideEffect, StatisticsOpInterface]>,
|
||||
Results<(outs AnyTensor)> {
|
||||
let arguments = (
|
||||
ins AnyTensor:$self,
|
||||
AnyType:$output_size
|
||||
);
|
||||
let summary = "aten _adaptive_avg_pool2d operator";
|
||||
let description = [{
|
||||
AdaptiveAvgPool2dOp
|
||||
aten _adaptive_avg_pool2d operator
|
||||
}];
|
||||
let extraClassDeclaration = [{
|
||||
std::map<std::string, uint64_t> getStatistics();
|
||||
}];
|
||||
}
|
||||
|
||||
def aten_AdaptiveAvgPool2dBackwardOp: aten_Op<"_adaptive_avg_pool2d_backward", [NoSideEffect, StatisticsOpInterface]>,
|
||||
Results<(outs AnyTensor)> {
|
||||
let arguments = (
|
||||
ins AnyTensor:$grad_output,
|
||||
AnyTensor:$self
|
||||
);
|
||||
let summary = "aten _adaptive_avg_pool2d_backward operator";
|
||||
let description = [{
|
||||
AdaptiveAvgPool2dBackwardOp
|
||||
aten _adaptive_avg_pool2d_backward operator
|
||||
}];
|
||||
let extraClassDeclaration = [{
|
||||
std::map<std::string, uint64_t> getStatistics();
|
||||
}];
|
||||
}
|
||||
|
||||
def aten_MaxPool2dWithIndicesOp: aten_Op<"max_pool2d_with_indices", [NoSideEffect, StatisticsOpInterface]>,
|
||||
Results<(outs AnyTensor, AnyTensor)> {
|
||||
let arguments = (
|
||||
ins AnyTensor:$self,
|
||||
AnyType:$kernel_size,
|
||||
AnyType:$stride,
|
||||
AnyType:$padding,
|
||||
AnyType:$dilation,
|
||||
AnyScalar:$ceil_mode
|
||||
);
|
||||
let summary = "aten max_pool2d_with_indices operator";
|
||||
let description = [{
|
||||
MaxPool2dWithIndicesOp
|
||||
aten max_pool2d_with_indices operator
|
||||
}];
|
||||
let extraClassDeclaration = [{
|
||||
std::map<std::string, uint64_t> getStatistics();
|
||||
}];
|
||||
}
|
||||
|
||||
def aten_MaxPool2dWithIndicesBackwardOp: aten_Op<"max_pool2d_with_indices_backward", [NoSideEffect, StatisticsOpInterface]>,
|
||||
Results<(outs AnyTensor)> {
|
||||
let arguments = (
|
||||
ins AnyTensor:$grad_output,
|
||||
AnyTensor:$self,
|
||||
AnyType:$kernel_size,
|
||||
AnyType:$stride,
|
||||
AnyType:$padding,
|
||||
AnyType:$dilation,
|
||||
AnyScalar:$ceil_mode,
|
||||
AnyTensor:$indices
|
||||
);
|
||||
let summary = "aten max_pool2d_with_indices_backward operator";
|
||||
let description = [{
|
||||
MaxPool2dWithIndicesBackwardOp
|
||||
aten max_pool2d_with_indices_backward operator
|
||||
}];
|
||||
let extraClassDeclaration = [{
|
||||
std::map<std::string, uint64_t> getStatistics();
|
||||
}];
|
||||
}
|
||||
|
||||
#endif
|
||||
|
|
@ -0,0 +1,28 @@
|
|||
//===- ATenPasses.h ---------------------------------------------*- 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
|
||||
//
|
||||
//===----------------------------------------------------------------------===//
|
||||
|
||||
#ifndef NPCOMP_DIALECT_ATEN_PASSES_H
|
||||
#define NPCOMP_DIALECT_ATEN_PASSES_H
|
||||
|
||||
#include "npcomp/Dialect/ATen/ATenLayerNamePass.h"
|
||||
#include "npcomp/Dialect/ATen/ATenLoweringPass.h"
|
||||
#include "npcomp/Dialect/ATen/ATenOpReport.h"
|
||||
#include "npcomp/Dialect/ATen/ReturnEliminationPass.h"
|
||||
|
||||
namespace mlir {
|
||||
namespace NPCOMP {
|
||||
namespace aten {
|
||||
// #define GEN_PASS_CLASSES
|
||||
// #include "npcomp/Dialect/ATen/ATenPasses.h.inc"
|
||||
|
||||
void registerATenPasses();
|
||||
} // namespace aten
|
||||
} // namespace NPCOMP
|
||||
} // namespace mlir
|
||||
|
||||
#endif // NPCOMP_DIALECT_ATEN_PASSES_H
|
||||
|
|
@ -0,0 +1,31 @@
|
|||
//===- ATenToStd.h ----------------------------------------------*- 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
|
||||
//
|
||||
//===----------------------------------------------------------------------===//
|
||||
|
||||
#ifndef NPCOMP_DIALECT_ATEN_TO_STD_H
|
||||
#define NPCOMP_DIALECT_ATEN_TO_STD_H
|
||||
|
||||
#include "mlir/Transforms/DialectConversion.h"
|
||||
|
||||
namespace mlir {
|
||||
namespace NPCOMP {
|
||||
namespace aten {
|
||||
|
||||
class ATenDialect;
|
||||
|
||||
} // namespace aten
|
||||
} // namespace NPCOMP
|
||||
} // namespace mlir
|
||||
|
||||
namespace mlir {
|
||||
|
||||
void populateATenToStdPatterns(MLIRContext *context,
|
||||
OwningRewritePatternList &patterns);
|
||||
|
||||
} // namespace mlir
|
||||
|
||||
#endif // NPCOMP_DIALECT_ATEN_TO_STD_H
|
||||
|
|
@ -0,0 +1,34 @@
|
|||
//===- ATenToStd.td ----------------------------------------*- tablegen -*-===//
|
||||
//
|
||||
// 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
|
||||
//
|
||||
//===----------------------------------------------------------------------===//
|
||||
|
||||
#ifdef MLIR_ATEN_TO_STD_TD
|
||||
#else
|
||||
#define MLIR_ATEN_TO_STD_TD
|
||||
|
||||
#ifdef STANDARD_OPS
|
||||
#else
|
||||
include "mlir/Dialect/StandardOps/IR/Ops.td"
|
||||
#endif // STANDARD_OPS
|
||||
|
||||
#ifdef ATEN_OPS
|
||||
#else
|
||||
include "ATen.td"
|
||||
#endif
|
||||
|
||||
// The pytorch convolution operator has 9 arguments, but we only have a jit
|
||||
// library that supports the first six at the moment.
|
||||
def : Pat<(aten_ConvolutionOverrideableOp $a1, $a2, $a3, $a4, $a5, $a6,
|
||||
$a7, $a8, $a9),
|
||||
(aten_ConvolutionOp $a1, $a2, $a3, $a4, $a5, $a6)>;
|
||||
|
||||
def : Pat<(aten_ConvolutionBackwardOverrideableOp $a1, $a2, $a3, $a4, $a5, $a6,
|
||||
$a7, $a8, $a9),
|
||||
(aten_ConvolutionBackwardOp $a1, $a2, $a3, $a4, $a5, $a6)>;
|
||||
|
||||
|
||||
#endif
|
||||
|
|
@ -0,0 +1,17 @@
|
|||
include_directories(${PROJECT_SOURCE_DIR}/dialect)
|
||||
|
||||
add_mlir_dialect(ATen aten)
|
||||
set(LLVM_TARGET_DEFINITIONS ATen.td)
|
||||
mlir_tablegen(ATenEnums.h.inc -gen-enum-decls)
|
||||
mlir_tablegen(ATenEnums.cpp.inc -gen-enum-defs)
|
||||
add_public_tablegen_target(MLIRATenEnumsIncGen)
|
||||
|
||||
set(LLVM_TARGET_DEFINITIONS ATenOpInterface.td)
|
||||
mlir_tablegen(ATenOpInterfaces.h.inc -gen-op-interface-decls)
|
||||
mlir_tablegen(ATenOpInterfaces.cpp.inc -gen-op-interface-defs)
|
||||
add_public_tablegen_target(MLIRATenOpInterfacesIncGen)
|
||||
add_dependencies(mlir-generic-headers MLIRATenOpInterfacesIncGen)
|
||||
|
||||
set(LLVM_TARGET_DEFINITIONS ATenToStd.td)
|
||||
mlir_tablegen(ATenToStd.cpp.inc -gen-rewriters)
|
||||
add_public_tablegen_target(MLIRATenToStdIncGen)
|
||||
|
|
@ -0,0 +1,40 @@
|
|||
//===- LivenessReport.h -----------------------------------------*- 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
|
||||
//
|
||||
//===----------------------------------------------------------------------===//
|
||||
|
||||
#ifndef NPCOMP_DIALECT_ATEN_LIVENESSREPORT_H
|
||||
#define NPCOMP_DIALECT_ATEN_LIVENESSREPORT_H
|
||||
|
||||
#include <string>
|
||||
|
||||
namespace mlir {
|
||||
namespace NPCOMP {
|
||||
namespace aten {
|
||||
|
||||
struct LivenessReport {
|
||||
|
||||
public:
|
||||
LivenessReport(mlir::ModuleOp &module) : module(module) {}
|
||||
|
||||
std::string generateTextReport();
|
||||
std::string emitJSONReport();
|
||||
llvm::DenseMap<Value, std::vector<Operation *>> &getLiveness() {
|
||||
return livenessIntervals;
|
||||
};
|
||||
|
||||
private:
|
||||
void resolveLiveness();
|
||||
|
||||
mlir::ModuleOp &module;
|
||||
llvm::DenseMap<Value, std::vector<Operation *>> livenessIntervals;
|
||||
};
|
||||
|
||||
} // namespace aten
|
||||
} // namespace NPCOMP
|
||||
} // namespace mlir
|
||||
|
||||
#endif
|
||||
|
|
@ -0,0 +1,28 @@
|
|||
//===- ReturnEliminationPass.h ----------------------------------*- 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
|
||||
//
|
||||
//===----------------------------------------------------------------------===//
|
||||
|
||||
#ifndef NPCOMP_DIALECT_ATEN_RETURNELIMINATIONPASS_H
|
||||
#define NPCOMP_DIALECT_ATEN_RETURNELIMINATIONPASS_H
|
||||
|
||||
#include <memory>
|
||||
|
||||
namespace mlir {
|
||||
class Pass;
|
||||
} // namespace mlir
|
||||
|
||||
namespace mlir {
|
||||
namespace NPCOMP {
|
||||
namespace aten {
|
||||
|
||||
std::unique_ptr<mlir::Pass> createReturnEliminationPass();
|
||||
void registerReturnEliminationPass();
|
||||
} // namespace aten
|
||||
} // namespace NPCOMP
|
||||
} // namespace mlir
|
||||
|
||||
#endif
|
||||
|
|
@ -1,3 +1,4 @@
|
|||
add_subdirectory(ATen)
|
||||
add_subdirectory(Basicpy)
|
||||
add_subdirectory(Npcomprt)
|
||||
add_subdirectory(Numpy)
|
||||
|
|
|
|||
|
|
@ -43,6 +43,7 @@ add_mlir_library(NPCOMPInitAll
|
|||
NPCOMPTCP
|
||||
NPCOMPTCF
|
||||
NPCOMPNpcomprt
|
||||
NPCOMPATenDialect
|
||||
NPCOMPBasicpyDialect
|
||||
NPCOMPBasicpyPasses
|
||||
NPCOMPNumpyDialect
|
||||
|
|
|
|||
|
|
@ -0,0 +1,106 @@
|
|||
//===- ATenDialect.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 "npcomp/Dialect/ATen/ATenDialect.h"
|
||||
#include "mlir/IR/DialectImplementation.h"
|
||||
|
||||
using namespace mlir;
|
||||
|
||||
namespace mlir {
|
||||
namespace NPCOMP {
|
||||
namespace aten {
|
||||
|
||||
namespace detail {
|
||||
|
||||
/// This class holds the implementation of the ATenListType.
|
||||
/// It is intended to be uniqued based on its content and owned by the context.
|
||||
struct ATenListTypeStorage : public mlir::TypeStorage {
|
||||
ATenListTypeStorage(Type elementType) : elementType(elementType) {}
|
||||
|
||||
/// The hash key used for uniquing.
|
||||
using KeyTy = mlir::Type;
|
||||
bool operator==(const KeyTy &key) const { return key == getElementType(); }
|
||||
|
||||
/// This is a factory method to create our type storage. It is only
|
||||
/// invoked after looking up the type in the context using the key and not
|
||||
/// finding it.
|
||||
static ATenListTypeStorage *construct(mlir::TypeStorageAllocator &allocator,
|
||||
const KeyTy &key) {
|
||||
|
||||
// Allocate the instance for the ATenListTypeStorage itself
|
||||
auto *storage = allocator.allocate<ATenListTypeStorage>();
|
||||
// Initialize the instance using placement new.
|
||||
return new (storage) ATenListTypeStorage(key);
|
||||
}
|
||||
|
||||
Type getElementType() const { return elementType; }
|
||||
|
||||
private:
|
||||
Type elementType;
|
||||
};
|
||||
} // namespace detail
|
||||
|
||||
ATenListType ATenListType::get(mlir::Type elemType) {
|
||||
return Base::get(elemType.getContext(), ATenTypeKind::ATEN_LIST, elemType);
|
||||
}
|
||||
|
||||
mlir::Type ATenListType::getElementType() {
|
||||
return getImpl()->getElementType();
|
||||
}
|
||||
|
||||
mlir::Type ATenDialect::parseType(DialectAsmParser &parser) const {
|
||||
Location loc = parser.getEncodedSourceLoc(parser.getNameLoc());
|
||||
|
||||
// All types start with an identifier that we switch on.
|
||||
StringRef typeNameSpelling;
|
||||
if (failed(parser.parseKeyword(&typeNameSpelling)))
|
||||
return nullptr;
|
||||
|
||||
if (typeNameSpelling == "list") {
|
||||
if (failed(parser.parseLess()))
|
||||
return nullptr;
|
||||
Type t;
|
||||
if (failed(parser.parseType(t)))
|
||||
return nullptr;
|
||||
if (failed(parser.parseGreater()))
|
||||
return nullptr;
|
||||
return ATenListType::get(t);
|
||||
}
|
||||
|
||||
parser.emitError(parser.getCurrentLocation(),
|
||||
"Invalid ATen type '" + typeNameSpelling + "'");
|
||||
return nullptr;
|
||||
}
|
||||
|
||||
/// Print a ATenListType
|
||||
void ATenDialect::printType(mlir::Type type, DialectAsmPrinter &os) const {
|
||||
auto ty = type.dyn_cast<ATenListType>();
|
||||
if (!ty) {
|
||||
os << "unknown aten type";
|
||||
return;
|
||||
}
|
||||
os << "list<";
|
||||
os.getStream() << ty.getElementType();
|
||||
os << ">";
|
||||
}
|
||||
|
||||
ATenDialect::ATenDialect(mlir::MLIRContext *ctx) : mlir::Dialect("aten", ctx) {
|
||||
addTypes<ATenListType>();
|
||||
addOperations<
|
||||
#define GET_OP_LIST
|
||||
#include "npcomp/Dialect/ATen/ATen.cpp.inc"
|
||||
>();
|
||||
}
|
||||
|
||||
#define GET_OP_CLASSES
|
||||
#include "npcomp/Dialect/ATen/ATen.cpp.inc"
|
||||
|
||||
} // namespace aten
|
||||
#include "npcomp/Dialect/ATen/ATenOpInterfaces.cpp.inc"
|
||||
} // namespace NPCOMP
|
||||
} // namespace mlir
|
||||
|
|
@ -0,0 +1,826 @@
|
|||
//===- ATenDialectOpStats.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 "npcomp/Dialect/ATen/ATenDialect.h"
|
||||
#include "npcomp/Dialect/ATen/ATenOpStatisticsUtils.h"
|
||||
|
||||
#include "llvm/Support/Debug.h"
|
||||
|
||||
#include "mlir/IR/StandardTypes.h"
|
||||
#include "mlir/IR/Types.h"
|
||||
|
||||
#include <iostream>
|
||||
|
||||
#define DEBUG_TYPE "aten-op-stats"
|
||||
|
||||
// This file contains the StatisticsOpInterface implementations
|
||||
// for ATDialect operations
|
||||
|
||||
using namespace mlir;
|
||||
|
||||
namespace {
|
||||
|
||||
std::vector<uint64_t> unpackListConstant(Value op) {
|
||||
std::vector<uint64_t> v;
|
||||
auto co = cast<mlir::NPCOMP::aten::ConstantOp>(op.getDefiningOp());
|
||||
DenseElementsAttr a = co.template getAttrOfType<DenseElementsAttr>("value");
|
||||
for (auto i : a.getIntValues())
|
||||
v.push_back(i.getSExtValue());
|
||||
return v;
|
||||
};
|
||||
|
||||
} // namespace
|
||||
|
||||
namespace mlir {
|
||||
namespace NPCOMP {
|
||||
namespace aten {
|
||||
|
||||
std::map<std::string, uint64_t> AdaptiveAvgPool2dOp::getStatistics() {
|
||||
std::map<std::string, uint64_t> toReturn;
|
||||
// FIXME: unimplemented
|
||||
toReturn["reads"] = -1;
|
||||
toReturn["writes"] = -1;
|
||||
return toReturn;
|
||||
}
|
||||
std::map<std::string, uint64_t> AdaptiveAvgPool2dBackwardOp::getStatistics() {
|
||||
std::map<std::string, uint64_t> toReturn;
|
||||
// FIXME: unimplemented
|
||||
toReturn["reads"] = -1;
|
||||
toReturn["writes"] = -1;
|
||||
return toReturn;
|
||||
}
|
||||
|
||||
// add
|
||||
std::map<std::string, uint64_t> AddOp::getStatistics() {
|
||||
|
||||
std::map<std::string, uint64_t> toReturn;
|
||||
|
||||
TensorType resultTy = getResult().getType().cast<TensorType>();
|
||||
TensorType aType = getOperand(0).getType().cast<TensorType>();
|
||||
Type bType = getOperand(1).getType();
|
||||
|
||||
uint64_t ofm_volume = getTensorVolume(resultTy);
|
||||
|
||||
toReturn["ops:+"] = ofm_volume;
|
||||
toReturn["result:0:activation_out"] = ofm_volume;
|
||||
|
||||
// Find the size of the A and B operands
|
||||
uint64_t a_volume = getTensorVolume(aType);
|
||||
uint64_t b_volume = getTensorVolume(bType);
|
||||
|
||||
toReturn["operand:0:activation_in"] = a_volume;
|
||||
toReturn["operand:1:activation_in"] = b_volume;
|
||||
|
||||
toReturn["reads"] = a_volume + b_volume;
|
||||
toReturn["writes"] = ofm_volume;
|
||||
|
||||
return toReturn;
|
||||
}
|
||||
|
||||
// add_
|
||||
std::map<std::string, uint64_t> AddUnderOp::getStatistics() {
|
||||
|
||||
std::map<std::string, uint64_t> toReturn;
|
||||
|
||||
TensorType resultTy = getResult().getType().cast<TensorType>();
|
||||
TensorType aType = getOperand(0).getType().cast<TensorType>();
|
||||
Type bType = getOperand(1).getType();
|
||||
|
||||
uint64_t ofm_volume = getTensorVolume(resultTy);
|
||||
|
||||
toReturn["ops:+"] = ofm_volume;
|
||||
toReturn["result:0:activation_out"] = ofm_volume;
|
||||
|
||||
// Find the size of the A and B operands
|
||||
uint64_t a_volume = getTensorVolume(aType);
|
||||
uint64_t b_volume = getTensorVolume(bType);
|
||||
|
||||
toReturn["operand:0:activation_in"] = a_volume;
|
||||
toReturn["operand:1:activation_in"] = b_volume;
|
||||
|
||||
toReturn["reads"] = a_volume + b_volume;
|
||||
toReturn["writes"] = ofm_volume;
|
||||
|
||||
return toReturn;
|
||||
}
|
||||
|
||||
// addmm
|
||||
std::map<std::string, uint64_t> AddmmOp::getStatistics() {
|
||||
|
||||
std::map<std::string, uint64_t> toReturn;
|
||||
|
||||
// For linear, we need the number of output neurons and the number of input
|
||||
// neurons Then the number of forward MACs is input * output And the number of
|
||||
// adds is output if there is bias
|
||||
|
||||
TensorType resultTy = getResult().getType().cast<TensorType>();
|
||||
TensorType biasTy = getOperand(0).getType().cast<TensorType>();
|
||||
TensorType inputTy = getOperand(1).getType().cast<TensorType>();
|
||||
TensorType weightTy = getOperand(2).getType().cast<TensorType>();
|
||||
|
||||
uint64_t num_output_neurons = resultTy.getShape()[1];
|
||||
uint64_t ofm_volume = getTensorVolume(resultTy);
|
||||
|
||||
// Use the weight tensor to find the number of input neurons
|
||||
uint64_t num_input_neurons = weightTy.getShape()[0];
|
||||
uint64_t total_MACs = ofm_volume * num_input_neurons;
|
||||
uint64_t weight_volume = getTensorVolume(weightTy);
|
||||
|
||||
uint64_t ifm_volume = getTensorVolume(inputTy);
|
||||
|
||||
toReturn["ops:MAC"] = total_MACs;
|
||||
toReturn["ops:+"] =
|
||||
ofm_volume; // Should be gated on whether there is bias at all
|
||||
toReturn["operand:1:activation_in"] = ifm_volume;
|
||||
toReturn["result:0:activation_out"] = ofm_volume;
|
||||
toReturn["operand:0:parameters_in:bias"] = getTensorVolume(biasTy);
|
||||
toReturn["operand:2:parameters_in:weight"] = weight_volume;
|
||||
|
||||
toReturn["reads"] = ifm_volume + weight_volume + num_output_neurons;
|
||||
toReturn["writes"] = ofm_volume;
|
||||
|
||||
return toReturn;
|
||||
}
|
||||
|
||||
// as_strided can be zero overhead
|
||||
std::map<std::string, uint64_t> AsStridedOp::getStatistics() {
|
||||
std::map<std::string, uint64_t> toReturn;
|
||||
toReturn["reads"] = 0;
|
||||
toReturn["writes"] = 0;
|
||||
toReturn["operand:0:activation_in"] = 0;
|
||||
toReturn["result:0:activation_out"] = 0;
|
||||
return toReturn;
|
||||
}
|
||||
|
||||
// batch_norm
|
||||
std::map<std::string, uint64_t> BatchNormOp::getStatistics() {
|
||||
|
||||
std::map<std::string, uint64_t> toReturn;
|
||||
|
||||
TensorType resultTy = getResult(0).getType().cast<TensorType>();
|
||||
uint64_t op_volume = getTensorVolume(resultTy);
|
||||
uint64_t weight_volume = getTensorVolume(getOperand(1).getType());
|
||||
uint64_t bias_volume = getTensorVolume(getOperand(2).getType());
|
||||
toReturn["operand:0:activation_in"] = op_volume;
|
||||
toReturn["result:0:activation_out"] = op_volume;
|
||||
toReturn["operand:1:parameters_in:weight"] = weight_volume;
|
||||
toReturn["operand:2:parameters_in:bias"] = bias_volume;
|
||||
|
||||
// Now for the arithmetic. Assume variance is calculated as sum of squares
|
||||
uint64_t ifm_depth = resultTy.getShape()[1];
|
||||
|
||||
toReturn["ops:+"] = op_volume; // Add up for mean
|
||||
toReturn["ops:*"] = op_volume; // Square for variance
|
||||
toReturn["ops:+"] += op_volume; // Add up squares for variance
|
||||
|
||||
toReturn["ops:*"] += ifm_depth; // Calc channel means
|
||||
toReturn["ops:-"] += ifm_depth; // Calc channel vars
|
||||
toReturn["ops:*"] += ifm_depth; // Calc channel vars
|
||||
|
||||
toReturn["ops:sqrt"] = ifm_depth; // Convert to SD
|
||||
toReturn["ops:/"] = ifm_depth; // Get the reciprocal
|
||||
|
||||
toReturn["ops:+"] += op_volume; // Subtract mean off each pixel
|
||||
toReturn["ops:*"] += op_volume; // Multiply by 1/SD for each pixel
|
||||
|
||||
toReturn["ops:+"] += op_volume; // Bias
|
||||
toReturn["ops:*"] += op_volume; // Scale
|
||||
|
||||
toReturn["reads"] = op_volume + weight_volume + bias_volume;
|
||||
toReturn["writes"] = op_volume;
|
||||
|
||||
return toReturn;
|
||||
}
|
||||
|
||||
// _convolution
|
||||
std::map<std::string, uint64_t> ConvolutionOp::getStatistics() {
|
||||
return getConv2dStatistics(this, /*groups*/ 1);
|
||||
}
|
||||
std::map<std::string, uint64_t> ConvolutionOverrideableOp::getStatistics() {
|
||||
// FIXME
|
||||
auto co = cast<mlir::NPCOMP::aten::ConstantOp>(groups().getDefiningOp());
|
||||
auto ia = co.template getAttrOfType<IntegerAttr>("value");
|
||||
uint64_t groups = ia.getValue().getZExtValue();
|
||||
|
||||
return getConv2dStatistics(this, groups);
|
||||
}
|
||||
|
||||
uint64_t ConvolutionOp::getOperandTransferVolume(unsigned int idx, bool read) {
|
||||
return getConv2dOperandTransferVolume<ConvolutionOp>(this, idx, read);
|
||||
}
|
||||
|
||||
uint64_t ConvolutionOp::getResultTransferVolume(unsigned int idx, bool write) {
|
||||
return getConv2dResultTransferVolume<ConvolutionOp>(this, idx, write);
|
||||
}
|
||||
|
||||
// _convolution_backward
|
||||
std::map<std::string, uint64_t> ConvolutionBackwardOp::getStatistics() {
|
||||
return getConv2dBackwardStatistics(*this, 1);
|
||||
}
|
||||
std::map<std::string, uint64_t> ConvolutionBackwardOverrideableOp::getStatistics() {
|
||||
auto co = cast<mlir::NPCOMP::aten::ConstantOp>(groups().getDefiningOp());
|
||||
auto ia = co.template getAttrOfType<IntegerAttr>("value");
|
||||
uint64_t groups = ia.getValue().getZExtValue();
|
||||
|
||||
return getConv2dBackwardStatistics(*this, groups);
|
||||
}
|
||||
|
||||
// div
|
||||
std::map<std::string, uint64_t> DivOp::getStatistics() {
|
||||
|
||||
std::map<std::string, uint64_t> toReturn;
|
||||
|
||||
TensorType resultTy = getResult().getType().cast<TensorType>();
|
||||
TensorType aType = getOperand(0).getType().cast<TensorType>();
|
||||
Type bType = getOperand(1).getType();
|
||||
|
||||
uint64_t ofm_volume = getTensorVolume(resultTy);
|
||||
toReturn["ops:/"] = ofm_volume;
|
||||
|
||||
toReturn["result:0:activation_out"] = ofm_volume;
|
||||
|
||||
// Find the size of the A and B operands
|
||||
uint64_t a_volume = getTensorVolume(aType);
|
||||
uint64_t b_volume = getTensorVolume(bType);
|
||||
|
||||
toReturn["operand:0:activation_in"] = a_volume;
|
||||
toReturn["operand:1:activation_in"] = b_volume;
|
||||
|
||||
toReturn["reads"] = a_volume + b_volume;
|
||||
toReturn["writes"] = ofm_volume;
|
||||
|
||||
return toReturn;
|
||||
}
|
||||
|
||||
// div_
|
||||
std::map<std::string, uint64_t> DivUnderOp::getStatistics() {
|
||||
|
||||
std::map<std::string, uint64_t> toReturn;
|
||||
|
||||
TensorType resultTy = getResult().getType().cast<TensorType>();
|
||||
TensorType aType = getOperand(0).getType().cast<TensorType>();
|
||||
Type bType = getOperand(1).getType();
|
||||
|
||||
uint64_t ofm_volume = getTensorVolume(resultTy);
|
||||
toReturn["ops:/"] = ofm_volume;
|
||||
toReturn["result:0:activation_out"] = ofm_volume;
|
||||
|
||||
// Find the size of the A and B operands
|
||||
uint64_t a_volume = getTensorVolume(aType);
|
||||
uint64_t b_volume = getTensorVolume(bType);
|
||||
|
||||
toReturn["operand:0:activation_in"] = a_volume;
|
||||
toReturn["operand:1:activation_in"] = b_volume;
|
||||
|
||||
toReturn["reads"] = a_volume + b_volume;
|
||||
toReturn["writes"] = ofm_volume;
|
||||
|
||||
return toReturn;
|
||||
}
|
||||
|
||||
// expand can be zero overhead
|
||||
std::map<std::string, uint64_t> ExpandOp::getStatistics() {
|
||||
std::map<std::string, uint64_t> toReturn;
|
||||
toReturn["reads"] = toReturn["operand:0:activation_in"] = 0;
|
||||
toReturn["writes"] = toReturn["result:0:activation_out"] = 0;
|
||||
return toReturn;
|
||||
}
|
||||
|
||||
// flatten can be zero overhead
|
||||
std::map<std::string, uint64_t> FlattenOp::getStatistics() {
|
||||
std::map<std::string, uint64_t> toReturn;
|
||||
toReturn["reads"] = toReturn["operand:0:activation_in"] = 0;
|
||||
toReturn["writes"] = toReturn["result:0:activation_out"] = 0;
|
||||
return toReturn;
|
||||
}
|
||||
|
||||
std::map<std::string, uint64_t> GatherOp::getStatistics() {
|
||||
std::map<std::string, uint64_t> toReturn;
|
||||
// FIXME: unimplemented
|
||||
toReturn["reads"] = -1;
|
||||
toReturn["writes"] = -1;
|
||||
return toReturn;
|
||||
}
|
||||
|
||||
// hardtanh
|
||||
std::map<std::string, uint64_t> HardtanhOp::getStatistics() {
|
||||
|
||||
std::map<std::string, uint64_t> toReturn;
|
||||
|
||||
TensorType inputTy = getOperand(0).getType().cast<TensorType>();
|
||||
TensorType resultTy = getResult().getType().cast<TensorType>();
|
||||
|
||||
uint64_t in_volume = getTensorVolume(inputTy);
|
||||
uint64_t out_volume = getTensorVolume(resultTy);
|
||||
|
||||
toReturn["operand:0:activation_in"] = in_volume;
|
||||
toReturn["result:0:activation_out"] = out_volume;
|
||||
toReturn["reads"] = in_volume;
|
||||
toReturn["writes"] = out_volume;
|
||||
toReturn["ops:>"] = out_volume;
|
||||
|
||||
return toReturn;
|
||||
}
|
||||
|
||||
// hardtanh_
|
||||
std::map<std::string, uint64_t> HardtanhUnderOp::getStatistics() {
|
||||
|
||||
std::map<std::string, uint64_t> toReturn;
|
||||
|
||||
TensorType inputTy = getOperand(0).getType().cast<TensorType>();
|
||||
TensorType resultTy = getResult().getType().cast<TensorType>();
|
||||
|
||||
uint64_t in_volume = getTensorVolume(inputTy);
|
||||
uint64_t out_volume = getTensorVolume(resultTy);
|
||||
|
||||
toReturn["operand:0:activation_in"] = in_volume;
|
||||
toReturn["result:0:activation_out"] = out_volume;
|
||||
toReturn["reads"] = in_volume;
|
||||
toReturn["writes"] = out_volume;
|
||||
toReturn["ops:>"] = out_volume;
|
||||
|
||||
return toReturn;
|
||||
}
|
||||
|
||||
std::map<std::string, uint64_t> HardtanhBackwardOp::getStatistics() {
|
||||
std::map<std::string, uint64_t> toReturn;
|
||||
// FIXME: unimplemented
|
||||
return toReturn;
|
||||
}
|
||||
|
||||
// max_pool2d
|
||||
std::map<std::string, uint64_t> MaxPool2dOp::getStatistics() {
|
||||
|
||||
std::map<std::string, uint64_t> toReturn;
|
||||
|
||||
TensorType resultTy = getResult().getType().cast<TensorType>();
|
||||
TensorType inputType = getOperand(0).getType().cast<TensorType>();
|
||||
|
||||
uint64_t ofm_volume = getTensorVolume(resultTy);
|
||||
toReturn["result:0:activation_out"] = ofm_volume;
|
||||
|
||||
uint64_t ifm_volume = getTensorVolume(inputType);
|
||||
toReturn["input:0:activation_in"] = ifm_volume;
|
||||
|
||||
// To find the number of compares, we need the filter extent
|
||||
|
||||
std::vector<uint64_t> kernel_size = unpackListConstant(getOperand(1));
|
||||
|
||||
uint64_t aperture = kernel_size[0] * kernel_size[1];
|
||||
toReturn["ops:>"] = ofm_volume * (aperture - 1);
|
||||
|
||||
toReturn["reads"] = ifm_volume;
|
||||
toReturn["writes"] = ofm_volume;
|
||||
|
||||
return toReturn;
|
||||
}
|
||||
|
||||
// max_pool2d_with_indices
|
||||
std::map<std::string, uint64_t> MaxPool2dWithIndicesOp::getStatistics() {
|
||||
|
||||
std::map<std::string, uint64_t> toReturn;
|
||||
|
||||
uint64_t ofm_volume =
|
||||
getTensorVolume(getResult(0).getType().cast<TensorType>());
|
||||
uint64_t indices_volume =
|
||||
getTensorVolume(getResult(1).getType().cast<TensorType>());
|
||||
|
||||
toReturn["writes"] = ofm_volume + indices_volume;
|
||||
toReturn["result:0:activation_out"] = ofm_volume;
|
||||
toReturn["result:1:indices_out"] = indices_volume;
|
||||
|
||||
uint64_t ifm_volume =
|
||||
getTensorVolume(getOperand(0).getType().cast<TensorType>());
|
||||
toReturn["reads"] = ifm_volume;
|
||||
toReturn["operand:0:activation_in"] = ifm_volume;
|
||||
|
||||
// To find the number of compares, we need the filter extent
|
||||
|
||||
std::vector<uint64_t> kernel_size = unpackListConstant(getOperand(1));
|
||||
|
||||
uint64_t aperture = kernel_size[0] * kernel_size[1];
|
||||
toReturn["ops:>"] = ofm_volume * (aperture - 1);
|
||||
|
||||
return toReturn;
|
||||
}
|
||||
|
||||
// max_pool2d_with_indices_backward
|
||||
std::map<std::string, uint64_t>
|
||||
MaxPool2dWithIndicesBackwardOp::getStatistics() {
|
||||
|
||||
std::map<std::string, uint64_t> toReturn;
|
||||
|
||||
Type resultTy = getResult().getType();
|
||||
TensorType tensorResultTy = resultTy.cast<TensorType>();
|
||||
uint64_t loss_out_volume = getTensorVolume(tensorResultTy);
|
||||
toReturn["writes"] = loss_out_volume;
|
||||
|
||||
uint64_t loss_in_volume =
|
||||
getTensorVolume(getOperand(0).getType().cast<TensorType>());
|
||||
uint64_t act_in_volume = getTensorVolume(
|
||||
getOperand(1).getType().cast<TensorType>()); // TODO: Why is this needed?
|
||||
uint64_t indices_volume =
|
||||
getTensorVolume(getOperand(7).getType().cast<TensorType>());
|
||||
toReturn["reads"] = loss_in_volume + act_in_volume + indices_volume;
|
||||
toReturn["operand:0:activation_in"] = loss_in_volume;
|
||||
toReturn["operand:1:activation_in"] = act_in_volume;
|
||||
toReturn["operand:3:activation_in"] = indices_volume;
|
||||
toReturn["result:0:grad:dx"] = loss_out_volume;
|
||||
|
||||
return toReturn;
|
||||
}
|
||||
|
||||
// mean
|
||||
std::map<std::string, uint64_t> MeanOp::getStatistics() {
|
||||
|
||||
std::map<std::string, uint64_t> toReturn;
|
||||
|
||||
TensorType resultTy = getResult().getType().cast<TensorType>();
|
||||
TensorType aType = getOperand().getType().cast<TensorType>();
|
||||
|
||||
uint64_t ofm_volume = getTensorVolume(resultTy);
|
||||
toReturn["ops:+"] = ofm_volume;
|
||||
toReturn["result:0:activation_out"] = ofm_volume;
|
||||
|
||||
// Find the size of the A and B operands
|
||||
uint64_t a_volume = getTensorVolume(aType);
|
||||
|
||||
toReturn["operand:0:activation_in"] = a_volume;
|
||||
|
||||
toReturn["reads"] = a_volume;
|
||||
toReturn["writes"] = ofm_volume;
|
||||
|
||||
return toReturn;
|
||||
}
|
||||
|
||||
// mm
|
||||
// std::map<std::string, uint64_t> MMOp::getStatistics() {
|
||||
// getMMOpStatistics(*this);
|
||||
// }
|
||||
std::map<std::string, uint64_t> MmOp::getStatistics() {
|
||||
return getMMOpStatistics(*this );
|
||||
}
|
||||
|
||||
// mul
|
||||
std::map<std::string, uint64_t> MulOp::getStatistics() {
|
||||
|
||||
std::map<std::string, uint64_t> toReturn;
|
||||
|
||||
TensorType resultTy = getResult().getType().cast<TensorType>();
|
||||
TensorType aType = getOperand(0).getType().cast<TensorType>();
|
||||
Type bType = getOperand(1).getType();
|
||||
|
||||
uint64_t ofm_volume = getTensorVolume(resultTy);
|
||||
toReturn["ops:*"] = ofm_volume;
|
||||
toReturn["result:0:activation_out"] = ofm_volume;
|
||||
|
||||
// Find the size of the A and B operands
|
||||
uint64_t a_volume = getTensorVolume(aType);
|
||||
uint64_t b_volume = getTensorVolume(bType);
|
||||
|
||||
toReturn["operand:0:activation_in"] = a_volume;
|
||||
toReturn["operand:1:activation_in"] = b_volume;
|
||||
|
||||
toReturn["reads"] = a_volume + b_volume;
|
||||
toReturn["writes"] = ofm_volume;
|
||||
|
||||
return toReturn;
|
||||
}
|
||||
|
||||
// mul_
|
||||
std::map<std::string, uint64_t> MulUnderOp::getStatistics() {
|
||||
|
||||
std::map<std::string, uint64_t> toReturn;
|
||||
|
||||
TensorType resultTy = getResult().getType().cast<TensorType>();
|
||||
TensorType aType = getOperand(0).getType().cast<TensorType>();
|
||||
Type bType = getOperand(1).getType();
|
||||
|
||||
uint64_t ofm_volume = getTensorVolume(resultTy);
|
||||
toReturn["ops:*"] = ofm_volume;
|
||||
toReturn["result:0:activation_out"] = ofm_volume;
|
||||
|
||||
// Find the size of the A and B operands
|
||||
uint64_t a_volume = getTensorVolume(aType);
|
||||
uint64_t b_volume = getTensorVolume(bType);
|
||||
|
||||
toReturn["operand:0:activation_in"] = a_volume;
|
||||
toReturn["operand:1:activation_in"] = b_volume;
|
||||
|
||||
toReturn["reads"] = a_volume + b_volume;
|
||||
toReturn["writes"] = ofm_volume;
|
||||
|
||||
return toReturn;
|
||||
}
|
||||
|
||||
// native_batch_norm
|
||||
std::map<std::string, uint64_t> NativeBatchNormOp::getStatistics() {
|
||||
|
||||
std::map<std::string, uint64_t> toReturn;
|
||||
|
||||
TensorType resultTy = getResult(0).getType().cast<TensorType>();
|
||||
uint64_t op_volume = getTensorVolume(resultTy);
|
||||
uint64_t weight_volume = getTensorVolume(getOperand(1).getType());
|
||||
uint64_t bias_volume = getTensorVolume(getOperand(2).getType());
|
||||
toReturn["operand:0:activation_in"] = op_volume;
|
||||
toReturn["result:0:activation_out"] = op_volume;
|
||||
toReturn["operand:1:parameters_in:weight"] = weight_volume;
|
||||
toReturn["operand:2:parameters_in:bias"] = bias_volume;
|
||||
|
||||
// Now for the arithmetic. Assume variance is calculated as sum of squares
|
||||
uint64_t ifm_depth = resultTy.getShape()[1];
|
||||
|
||||
toReturn["ops:+"] = op_volume; // Add up for mean
|
||||
toReturn["ops:*"] = op_volume; // Square for variance
|
||||
toReturn["ops:+"] += op_volume; // Add up squares for variance
|
||||
|
||||
toReturn["ops:*"] += ifm_depth; // Calc channel means
|
||||
toReturn["ops:-"] += ifm_depth; // Calc channel vars
|
||||
toReturn["ops:*"] += ifm_depth; // Calc channel vars
|
||||
|
||||
toReturn["ops:sqrt"] = ifm_depth; // Convert to SD
|
||||
toReturn["ops:/"] = ifm_depth; // Get the reciprocal
|
||||
|
||||
toReturn["ops:+"] += op_volume; // Subtract mean off each pixel
|
||||
toReturn["ops:*"] += op_volume; // Multiply by 1/SD for each pixel
|
||||
|
||||
toReturn["ops:+"] += op_volume; // Bias
|
||||
toReturn["ops:*"] += op_volume; // Scale
|
||||
|
||||
toReturn["reads"] = op_volume + weight_volume + bias_volume;
|
||||
toReturn["writes"] = op_volume;
|
||||
|
||||
return toReturn;
|
||||
}
|
||||
|
||||
// batchnorm backward
|
||||
std::map<std::string, uint64_t> NativeBatchNormBackwardOp::getStatistics() {
|
||||
|
||||
std::map<std::string, uint64_t> toReturn;
|
||||
|
||||
ShapedType inputTy = getOperand(0).getType().cast<ShapedType>();
|
||||
uint64_t input_volume = getTensorVolume(inputTy);
|
||||
uint64_t input_channels = inputTy.getShape()[1];
|
||||
|
||||
// # 3 components make up the gradInput: 1 gradInput, 2 gradMean, 3 gradVar
|
||||
// # totalGradInput = gradInput + (dL / dMean * dMean / dInput) +
|
||||
// # (dL / dVar * dVar / dInput)
|
||||
|
||||
// # gradInput
|
||||
// total_ops["backward"]["*"] = in_c * (in_h*in_w*batch_size) # scale
|
||||
// # Bootstrap from previous
|
||||
// #total_ops["backward"]["sqrt"] = in_c # Convert to std_dev
|
||||
// #total_ops["backward"]["/"] = in_c # Calculate inverse sqrt first
|
||||
toReturn["ops:*"] = input_volume; // scale
|
||||
|
||||
// # dL / dGradVar
|
||||
// total_ops["backward"]["pow"] = in_c
|
||||
// total_ops["backward"]["*"] = total_ops["backward"]["*"] + in_c
|
||||
// #total_ops["backward"]["+"] = total_ops["backward"]["+"] + in_c *
|
||||
// in_h*in_w*batch_size # Subtract mean, bootstrap from previous calculation
|
||||
// total_ops["backward"]["*"] = total_ops["backward"]["*"] + in_c *
|
||||
// (in_h*in_w*batch_size)
|
||||
toReturn["ops:pow"] = input_channels;
|
||||
;
|
||||
toReturn["ops:*"] += input_channels;
|
||||
toReturn["ops:*"] += input_volume;
|
||||
|
||||
// # dL / dGradMean
|
||||
// #total_ops["backward"]["+"] = total_ops["backward"]["+"] + in_c *
|
||||
// (in_h*in_w*batch_size) # bootstrap from previous total_ops["backward"]["*"]
|
||||
// = total_ops["backward"]["*"] + in_c # scale gradMean
|
||||
// total_ops["backward"]["*"] = total_ops["backward"]["*"] + in_c # eltwise
|
||||
// with dL / dGradVar total_ops["backward"]["+"] = in_c *
|
||||
// (in_h*in_w*batch_size) # sum gradXhat total_ops["backward"]["*"] =
|
||||
// total_ops["backward"]["*"] + in_c # scale gradXhat
|
||||
toReturn["ops:*"] += input_channels; // scale gradMean
|
||||
toReturn["ops:*"] += input_channels; // eltwise with dL / dGradVar
|
||||
toReturn["ops:+"] = input_volume; // sum gradXhat
|
||||
toReturn["ops:*"] += input_channels; // scale gradXhat
|
||||
|
||||
// # totalGradInput
|
||||
// total_ops["backward"]["+"] = total_ops["backward"]["+"] + in_c *
|
||||
// (in_h*in_w*batch_size) # Subtract mean, can't bootstrap this one
|
||||
// total_ops["backward"]["*"] = total_ops["backward"]["*"] + in_c # scale dL /
|
||||
// dMean total_ops["backward"]["*"] = total_ops["backward"]["*"] + in_c #
|
||||
// scale dL / dVar total_ops["backward"]["*"] = total_ops["backward"]["*"] +
|
||||
// in_c * (in_h*in_w*batch_size) # Eltwise multiply by dL / dVar
|
||||
// total_ops["backward"]["+"] = total_ops["backward"]["+"] + 2 * in_c *
|
||||
// (in_h*in_w*batch_size) # Accumulate gradient terms
|
||||
toReturn["ops:+"] += input_volume; // Subtract mean, can't bootstrap this one
|
||||
toReturn["ops:*"] += input_channels; // scale dL / dMean
|
||||
toReturn["ops:*"] += input_channels; // scale dL / dVar
|
||||
toReturn["ops:*"] += input_volume; // Eltwise multiply by dL / dVar
|
||||
toReturn["OPS:+"] += 2 * input_volume; // Accumulate gradient terms
|
||||
|
||||
uint64_t reads = 0;
|
||||
for (int i = 0; i < 7; i++) {
|
||||
auto v = getTensorVolume(getOperand(i).getType());
|
||||
toReturn["operand:" + std::to_string(i) + ":activation_in"] = v;
|
||||
reads += v;
|
||||
}
|
||||
|
||||
uint64_t writes = 0;
|
||||
for (int i = 0; i < 3; i++) {
|
||||
auto v = getTensorVolume(getResult(i).getType());
|
||||
toReturn["result:" + std::to_string(i) + ":grad"] = v;
|
||||
writes += v;
|
||||
}
|
||||
|
||||
toReturn["reads"] = reads;
|
||||
toReturn["writes"] = writes;
|
||||
|
||||
return toReturn;
|
||||
}
|
||||
|
||||
std::map<std::string, uint64_t> NllLossForwardOp::getStatistics() {
|
||||
std::map<std::string, uint64_t> toReturn;
|
||||
// FIXME: unimplemented
|
||||
toReturn["reads"] = -1;
|
||||
toReturn["writes"] = -1;
|
||||
return toReturn;
|
||||
}
|
||||
std::map<std::string, uint64_t> NllLossBackwardOp::getStatistics() {
|
||||
std::map<std::string, uint64_t> toReturn;
|
||||
// FIXME: unimplemented
|
||||
toReturn["reads"] = -1;
|
||||
toReturn["writes"] = -1;
|
||||
return toReturn;
|
||||
}
|
||||
std::map<std::string, uint64_t> NllLoss2dForwardOp::getStatistics() {
|
||||
std::map<std::string, uint64_t> toReturn;
|
||||
// FIXME: unimplemented
|
||||
toReturn["reads"] = -1;
|
||||
toReturn["writes"] = -1;
|
||||
return toReturn;
|
||||
}
|
||||
std::map<std::string, uint64_t> NllLoss2dBackwardOp::getStatistics() {
|
||||
std::map<std::string, uint64_t> toReturn;
|
||||
// FIXME: unimplemented
|
||||
toReturn["reads"] = -1;
|
||||
toReturn["writes"] = -1;
|
||||
return toReturn;
|
||||
}
|
||||
|
||||
// neg op
|
||||
std::map<std::string, uint64_t> NegOp::getStatistics() {
|
||||
std::map<std::string, uint64_t> toReturn;
|
||||
auto insize = getTensorVolume(getOperand().getType());
|
||||
auto outsize = getTensorVolume(getResult().getType());
|
||||
toReturn["reads"] = toReturn["operand:0:activation_in"] = insize;
|
||||
toReturn["writes"] = toReturn["result:0:activation_out"] = outsize;
|
||||
return toReturn;
|
||||
}
|
||||
|
||||
// relu
|
||||
// std::map<std::string, uint64_t> ReLUOp::getStatistics() {
|
||||
// return getReLUOpStatistics(*this);
|
||||
// }
|
||||
std::map<std::string, uint64_t> ReluOp::getStatistics() {
|
||||
return getReLUOpStatistics(*this);
|
||||
}
|
||||
// std::map<std::string, uint64_t> ReLUUnderOp::getStatistics() {
|
||||
// return getReLUOpStatistics(*this);
|
||||
// }
|
||||
std::map<std::string, uint64_t> ReluUnderOp::getStatistics() {
|
||||
return getReLUOpStatistics(*this);
|
||||
}
|
||||
|
||||
// sub
|
||||
std::map<std::string, uint64_t> SubOp::getStatistics() {
|
||||
|
||||
std::map<std::string, uint64_t> toReturn;
|
||||
|
||||
TensorType resultTy = getResult().getType().cast<TensorType>();
|
||||
TensorType aType = getOperand(0).getType().cast<TensorType>();
|
||||
Type bType = getOperand(1).getType();
|
||||
|
||||
uint64_t ofm_volume = getTensorVolume(resultTy);
|
||||
|
||||
toReturn["ops:-"] = ofm_volume;
|
||||
toReturn["result:0:activation_out"] = ofm_volume;
|
||||
|
||||
// Find the size of the A and B operands
|
||||
uint64_t a_volume = getTensorVolume(aType);
|
||||
uint64_t b_volume = getTensorVolume(bType);
|
||||
|
||||
toReturn["operand:0:activation_in"] = a_volume;
|
||||
toReturn["operand:1:activation_in"] = b_volume;
|
||||
|
||||
toReturn["reads"] = a_volume + b_volume;
|
||||
toReturn["writes"] = ofm_volume;
|
||||
|
||||
return toReturn;
|
||||
}
|
||||
|
||||
// sub_
|
||||
std::map<std::string, uint64_t> SubUnderOp::getStatistics() {
|
||||
|
||||
std::map<std::string, uint64_t> toReturn;
|
||||
|
||||
TensorType resultTy = getResult().getType().cast<TensorType>();
|
||||
TensorType aType = getOperand(0).getType().cast<TensorType>();
|
||||
Type bType = getOperand(1).getType();
|
||||
|
||||
uint64_t ofm_volume = getTensorVolume(resultTy);
|
||||
|
||||
toReturn["ops:-"] = ofm_volume;
|
||||
toReturn["result:0:activation_out"] = ofm_volume;
|
||||
|
||||
// Find the size of the A and B operands
|
||||
uint64_t a_volume = getTensorVolume(aType);
|
||||
uint64_t b_volume = getTensorVolume(bType);
|
||||
|
||||
toReturn["operand:0:activation_in"] = a_volume;
|
||||
toReturn["operand:1:activation_in"] = b_volume;
|
||||
|
||||
toReturn["reads"] = a_volume + b_volume;
|
||||
toReturn["writes"] = ofm_volume;
|
||||
|
||||
return toReturn;
|
||||
}
|
||||
|
||||
// sum
|
||||
std::map<std::string, uint64_t> SumOp::getStatistics() {
|
||||
|
||||
std::map<std::string, uint64_t> toReturn;
|
||||
TensorType ty = getOperand(0).getType().cast<TensorType>();
|
||||
uint64_t volume = getTensorVolume(ty);
|
||||
|
||||
toReturn["ops:+"] = volume;
|
||||
|
||||
toReturn["operand:0:activation_in"] = volume;
|
||||
toReturn["result:0:activation_out"] = volume;
|
||||
|
||||
toReturn["reads"] = volume;
|
||||
toReturn["writes"] = volume;
|
||||
|
||||
return toReturn;
|
||||
}
|
||||
|
||||
// size op can be zero overhead
|
||||
std::map<std::string, uint64_t> SizeOp::getStatistics() {
|
||||
std::map<std::string, uint64_t> toReturn;
|
||||
toReturn["reads"] = toReturn["operand:0:activation_in"] = 0;
|
||||
toReturn["writes"] = toReturn["result:0:activation_out"] = 0;
|
||||
return toReturn;
|
||||
}
|
||||
|
||||
// squeeze can be zero overhead
|
||||
std::map<std::string, uint64_t> SqueezeOp::getStatistics() {
|
||||
std::map<std::string, uint64_t> toReturn;
|
||||
toReturn["reads"] = toReturn["operand:0:activation_in"] = 0;
|
||||
toReturn["writes"] = toReturn["result:0:activation_out"] = 0;
|
||||
return toReturn;
|
||||
}
|
||||
|
||||
// transpose can be zero overhead
|
||||
std::map<std::string, uint64_t> TOp::getStatistics() {
|
||||
std::map<std::string, uint64_t> toReturn;
|
||||
toReturn["reads"] = toReturn["operand:0:activation_in"] = 0;
|
||||
toReturn["writes"] = toReturn["result:0:activation_out"] = 0;
|
||||
return toReturn;
|
||||
}
|
||||
|
||||
// threshold_backward
|
||||
std::map<std::string, uint64_t> ThresholdBackwardOp::getStatistics() {
|
||||
|
||||
std::map<std::string, uint64_t> toReturn;
|
||||
uint64_t loss_in_volume =
|
||||
getTensorVolume(getOperand(0).getType().cast<TensorType>());
|
||||
uint64_t act_in_volume =
|
||||
getTensorVolume(getOperand(1).getType().cast<TensorType>());
|
||||
uint64_t loss_out_volume =
|
||||
getTensorVolume(getResult().getType().cast<TensorType>());
|
||||
|
||||
toReturn["reads"] = toReturn["operand:0:activation_in"] =
|
||||
loss_in_volume + act_in_volume;
|
||||
toReturn["writes"] = toReturn["result:0:grad:dx"] = loss_out_volume;
|
||||
|
||||
return toReturn;
|
||||
}
|
||||
|
||||
// unsqueeze can be zero overhead
|
||||
std::map<std::string, uint64_t> UnsqueezeOp::getStatistics() {
|
||||
std::map<std::string, uint64_t> toReturn;
|
||||
toReturn["reads"] = toReturn["operand:0:activation_in"] = 0;
|
||||
toReturn["writes"] = toReturn["result:0:activation_out"] = 0;
|
||||
return toReturn;
|
||||
}
|
||||
|
||||
// view can be zero overhead
|
||||
std::map<std::string, uint64_t> ViewOp::getStatistics() {
|
||||
std::map<std::string, uint64_t> toReturn;
|
||||
toReturn["reads"] = toReturn["operand:0:activation_in"] = 0;
|
||||
toReturn["writes"] = toReturn["result:0:activation_out"] = 0;
|
||||
return toReturn;
|
||||
}
|
||||
|
||||
} // namespace aten
|
||||
} // namespace NPCOMP
|
||||
} // namespace mlir
|
||||
|
|
@ -0,0 +1,98 @@
|
|||
//===- ATenLayerNamePass.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 "npcomp/Dialect/ATen/ATenLayerNamePass.h"
|
||||
#include "npcomp/Dialect/ATen/ATenDialect.h"
|
||||
|
||||
#include "llvm/Support/Debug.h"
|
||||
#include "llvm/Support/ErrorHandling.h"
|
||||
#include "llvm/Support/raw_ostream.h"
|
||||
|
||||
#include "mlir/Pass/Pass.h"
|
||||
|
||||
#include <iostream>
|
||||
#include <vector>
|
||||
|
||||
#define DEBUG_TYPE "aten-layer-name"
|
||||
|
||||
using namespace mlir;
|
||||
|
||||
namespace {
|
||||
|
||||
struct ATenLayerNamePass
|
||||
: public PassWrapper<ATenLayerNamePass, OperationPass<ModuleOp>> {
|
||||
|
||||
private:
|
||||
std::map<Operation *, std::string> opToName;
|
||||
|
||||
public:
|
||||
ATenLayerNamePass() {}
|
||||
|
||||
void runOnOperation() override {
|
||||
|
||||
markAllAnalysesPreserved();
|
||||
|
||||
auto module = getOperation();
|
||||
|
||||
// find the function called 'graph'
|
||||
auto graph = module.lookupSymbol<mlir::FuncOp>("graph");
|
||||
if (!graph) {
|
||||
emitError(mlir::UnknownLoc::get(module.getContext()),
|
||||
"OpReportPass failed: can't find a graph function\n");
|
||||
signalPassFailure();
|
||||
return;
|
||||
}
|
||||
|
||||
// Construct a name for each aten operation
|
||||
std::map<std::string, uint64_t> layerIDmap;
|
||||
unsigned currentLayer = 0;
|
||||
|
||||
graph.walk([&](Operation *op) {
|
||||
auto name = op->getName().getStringRef();
|
||||
|
||||
// if it's not an aten operation, continue
|
||||
// TODO: need an interface for this rather than just
|
||||
// doing string compare.
|
||||
if (!name.startswith("aten."))
|
||||
return;
|
||||
|
||||
// strip the aten prefix to get the operation type
|
||||
auto type = name.split("aten.").second;
|
||||
|
||||
// if it's an aten constant op, continue
|
||||
if (type.equals("constant"))
|
||||
return;
|
||||
|
||||
unsigned ID = 0;
|
||||
if (layerIDmap.count(type.str()) == 0)
|
||||
layerIDmap[type.str()] = 0;
|
||||
else
|
||||
ID = ++layerIDmap[type.str()];
|
||||
|
||||
std::string layerName = "L" + std::to_string(currentLayer++) + "-" +
|
||||
type.str() + "-" + std::to_string(ID);
|
||||
|
||||
LLVM_DEBUG(llvm::dbgs()
|
||||
<< "generated layer_name: '" << layerName << "'\n");
|
||||
|
||||
auto attr = StringAttr::get(layerName, module.getContext());
|
||||
op->setAttr(StringRef("layer_name"), attr);
|
||||
});
|
||||
}
|
||||
};
|
||||
|
||||
} // namespace
|
||||
|
||||
std::unique_ptr<mlir::Pass> mlir::NPCOMP::aten::createATenLayerNamePass() {
|
||||
return std::make_unique<ATenLayerNamePass>();
|
||||
}
|
||||
|
||||
void mlir::NPCOMP::aten::registerATenLayerNamePass() {
|
||||
PassRegistration<ATenLayerNamePass>("aten-layer-name",
|
||||
"Generate layer names for ATen Dialect");
|
||||
}
|
||||
|
|
@ -0,0 +1,962 @@
|
|||
//===- ATenLoweringPass.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 "npcomp/Dialect/ATen/ATenLoweringPass.h"
|
||||
#include "npcomp/Dialect/ATen/ATenDialect.h"
|
||||
#include "npcomp/Dialect/ATen/ATenToStd.h"
|
||||
|
||||
#include "mlir/Dialect/Affine/EDSC/Builders.h"
|
||||
#include "mlir/Dialect/Affine/IR/AffineOps.h"
|
||||
#include "mlir/Dialect/Affine/IR/AffineValueMap.h"
|
||||
#include "mlir/Dialect/LLVMIR/LLVMDialect.h"
|
||||
#include "mlir/Dialect/SCF/EDSC/Builders.h"
|
||||
#include "mlir/Dialect/SCF/SCF.h"
|
||||
#include "mlir/Dialect/StandardOps/EDSC/Builders.h"
|
||||
#include "mlir/Dialect/StandardOps/EDSC/Intrinsics.h"
|
||||
#include "mlir/Dialect/StandardOps/IR/Ops.h"
|
||||
#include "mlir/EDSC/Builders.h"
|
||||
#include "mlir/IR/Builders.h"
|
||||
#include "mlir/IR/OperationSupport.h"
|
||||
#include "mlir/IR/StandardTypes.h"
|
||||
#include "mlir/Parser.h"
|
||||
#include "mlir/Pass/Pass.h"
|
||||
#include "mlir/Transforms/DialectConversion.h"
|
||||
|
||||
#include "mlir/Conversion/StandardToLLVM/ConvertStandardToLLVM.h"
|
||||
#include "mlir/Conversion/StandardToLLVM/ConvertStandardToLLVMPass.h"
|
||||
|
||||
#include "llvm/ADT/DenseSet.h"
|
||||
#include "llvm/ADT/STLExtras.h"
|
||||
#include "llvm/IR/DerivedTypes.h"
|
||||
#include "llvm/IR/IRBuilder.h"
|
||||
#include "llvm/IR/Type.h"
|
||||
#include "llvm/Support/Debug.h"
|
||||
#include "llvm/Support/ErrorHandling.h"
|
||||
#include "llvm/Support/raw_ostream.h"
|
||||
|
||||
#include <algorithm>
|
||||
#include <sstream>
|
||||
|
||||
using namespace mlir;
|
||||
using namespace edsc::intrinsics;
|
||||
using namespace mlir::NPCOMP::aten;
|
||||
|
||||
using callOperation = edsc::OperationBuilder<mlir::CallOp>;
|
||||
using call = edsc::ValueBuilder<mlir::CallOp>;
|
||||
using constInt = edsc::intrinsics::std_constant_int;
|
||||
using constFloat = edsc::intrinsics::std_constant_float;
|
||||
|
||||
namespace {
|
||||
|
||||
/// Utility function for type casting: this is making the type checker happy,
|
||||
/// while delaying the actual work involved to convert the type. Most of the
|
||||
/// time both side of the cast (producer and consumer) will be lowered to a
|
||||
/// dialect like LLVM and end up with the same LLVM representation, at which
|
||||
/// point this becomes a no-op and is eliminated.
|
||||
static Value typeCast(PatternRewriter &builder, Value val, Type destTy) {
|
||||
if (val.getType() == destTy)
|
||||
return val;
|
||||
return builder.create<mlir::NPCOMP::aten::TypeCastOp>(val.getLoc(), destTy, val)
|
||||
.getResult();
|
||||
}
|
||||
|
||||
/// Given a MemRefType, return a new MemRefType with the same rank, but
|
||||
/// unknown shape.
|
||||
static MemRefType getShapeErasedMemRefType(MemRefType type) {
|
||||
std::vector<int64_t> shape = type.getShape();
|
||||
for(int i = 0; i < shape.size(); i++) {
|
||||
shape[i] = -1;
|
||||
}
|
||||
return MemRefType::get(shape, type.getElementType(),
|
||||
type.getAffineMaps(), type.getMemorySpace());
|
||||
}
|
||||
|
||||
/// Create a type cast to memref
|
||||
static Value memRefTypeCast(PatternRewriter &builder, Value val) {
|
||||
Type type = val.getType();
|
||||
|
||||
if (auto memrefTy = type.dyn_cast<MemRefType>()) {
|
||||
MemRefType newType = getShapeErasedMemRefType(memrefTy);
|
||||
return builder.create<MemRefCastOp>(val.getLoc(),
|
||||
val, newType)
|
||||
.getResult();
|
||||
}
|
||||
if (auto tensorTy = type.dyn_cast<TensorType>()) {
|
||||
auto memRefType = mlir::MemRefType::get(tensorTy.getShape(),
|
||||
tensorTy.getElementType(), {}, 0);
|
||||
return typeCast(builder, val, memRefType);
|
||||
}
|
||||
return val;
|
||||
}
|
||||
|
||||
// Mangle a type in a way that encodes the full shape information.
|
||||
// TODO: Currently only supports MemRef, Float, Integer, and AtenList (poorly)
|
||||
static std::string getFullyMangledType(const Type ty) {
|
||||
std::stringstream ret;
|
||||
|
||||
if (const MemRefType mrt = ty.dyn_cast<const MemRefType>()) {
|
||||
ret << "M";
|
||||
auto shape = mrt.getShape();
|
||||
const Type elem = mrt.getElementType();
|
||||
for (auto s : shape)
|
||||
ret << s << "x";
|
||||
ret << getFullyMangledType(elem);
|
||||
} else if (FloatType ft = ty.dyn_cast<FloatType>()) {
|
||||
ret << "F" << ft.getWidth();
|
||||
} else if (const IntegerType it = ty.dyn_cast<const IntegerType>()) {
|
||||
ret << "I" << it.getWidth();
|
||||
} else if (const mlir::NPCOMP::aten::ATenListType alt =
|
||||
ty.dyn_cast<const mlir::NPCOMP::aten::ATenListType>()) {
|
||||
|
||||
} else {
|
||||
Type t = ty;
|
||||
t.dump();
|
||||
assert(0 && "unhandled type in getFullyMangledType");
|
||||
}
|
||||
return ret.str();
|
||||
}
|
||||
|
||||
// Mangle the argument shapes into the function name. This is impractical for
|
||||
// a library-based implementation, since each different shape has to be
|
||||
// implemented by a different function. The function name is constructed
|
||||
// from the prefix, the mangled result types, the mangled operand types.
|
||||
// Types are mangled in a way that encodes the full shape information.
|
||||
static std::string getFullyMangledFuncName(std::string prefix,
|
||||
FunctionType fnTy) {
|
||||
std::string sep = "_";
|
||||
|
||||
ArrayRef<Type> resultTy = fnTy.getResults();
|
||||
ArrayRef<Type> operTy = fnTy.getInputs();
|
||||
|
||||
std::string ret = prefix + "_AtenAcapOp_";
|
||||
for (const Type t : resultTy)
|
||||
ret = ret + sep + getFullyMangledType(t);
|
||||
for (const Type t : operTy)
|
||||
ret = ret + sep + getFullyMangledType(t);
|
||||
|
||||
return ret;
|
||||
}
|
||||
|
||||
// Mangle the argument ranks into the function name.
|
||||
// TODO: Currently only supports MemRef, Float, Integer, and AtenList (poorly)
|
||||
static std::string getSimplyMangledType(const Type ty) {
|
||||
std::stringstream ret;
|
||||
|
||||
if (const MemRefType mrt = ty.dyn_cast<const MemRefType>()) {
|
||||
// ret << "M";
|
||||
ArrayRef<int64_t> shape = mrt.getShape();
|
||||
const Type elem = mrt.getElementType();
|
||||
ret << shape.size();
|
||||
ret << getFullyMangledType(elem);
|
||||
} else if (FloatType ft = ty.dyn_cast<FloatType>()) {
|
||||
// ret << "F" << ft.getWidth();
|
||||
} else if (const IntegerType it = ty.dyn_cast<const IntegerType>()) {
|
||||
// ret << "I" << it.getWidth();
|
||||
} else if (const mlir::NPCOMP::aten::ATenListType alt =
|
||||
ty.dyn_cast<const mlir::NPCOMP::aten::ATenListType>()) {
|
||||
|
||||
} else {
|
||||
Type t = ty;
|
||||
t.dump();
|
||||
assert(0 && "unhandled type in getSimplyMangledType");
|
||||
}
|
||||
return ret.str();
|
||||
}
|
||||
|
||||
// Return a simply mangled function name. The function name is constructed
|
||||
// from the prefix, the mangled result types, the mangled operand types.
|
||||
// Types are mangled in a way that encodes only the rank. Shape information
|
||||
// is passed runtime using the standard calling convention. This simpler
|
||||
// version of mangling allows us to implement most of the functions with only
|
||||
// a few variations. However, it means we need to convert from tensor types
|
||||
// with known size to tensor types with unknown size to have a consistent
|
||||
// runtime calling convention.
|
||||
static std::string getSimplyMangledFuncName(std::string prefix,
|
||||
ArrayRef<Type> operTy,
|
||||
ArrayRef<Type> resultTy) {
|
||||
std::string sep = "_";
|
||||
|
||||
std::string ret = prefix;
|
||||
for (const Type t : resultTy)
|
||||
ret = ret + sep + getSimplyMangledType(t);
|
||||
for (const Type t : operTy) {
|
||||
std::string s = getSimplyMangledType(t);
|
||||
if(s.size() > 0)
|
||||
ret = ret + sep + getSimplyMangledType(t);
|
||||
}
|
||||
ret += "_out";
|
||||
|
||||
return ret;
|
||||
}
|
||||
static std::string getSimplyMangledFuncName(std::string prefix,
|
||||
FunctionType fnTy) {
|
||||
|
||||
return getSimplyMangledFuncName(prefix, fnTy.getInputs(), fnTy.getResults());
|
||||
}
|
||||
|
||||
std::string getMangledFuncName(std::string prefix,
|
||||
FunctionType fnTy) {
|
||||
return getSimplyMangledFuncName(prefix, fnTy);
|
||||
}
|
||||
|
||||
std::string getMangledFuncName(std::string prefix,
|
||||
ArrayRef<Type> opTys,
|
||||
ArrayRef<Type> retTys) {
|
||||
return getSimplyMangledFuncName(prefix, opTys, retTys);
|
||||
}
|
||||
|
||||
static FuncOp getATenFn(ModuleOp module, std::string mangledFunctionName,
|
||||
ArrayRef<Value> operands,
|
||||
ArrayRef<Type> retTys) {
|
||||
Builder builder(module);
|
||||
|
||||
SmallVector<Type, 8> tys;
|
||||
for (Value o : operands) {
|
||||
Type t = o.getType();
|
||||
// Erase the dimensions of the memref.
|
||||
if (t.isa<MemRefType>()) {
|
||||
auto mt = t.cast<MemRefType>();
|
||||
tys.push_back(getShapeErasedMemRefType(mt));
|
||||
} else
|
||||
tys.push_back(t);
|
||||
}
|
||||
|
||||
auto fnTy = builder.getFunctionType(tys, retTys);
|
||||
|
||||
auto fn = module.lookupSymbol<FuncOp>(mangledFunctionName);
|
||||
|
||||
if (!fn) {
|
||||
fn = FuncOp::create(builder.getUnknownLoc(), mangledFunctionName, fnTy);
|
||||
module.push_back(fn);
|
||||
}
|
||||
|
||||
return fn;
|
||||
}
|
||||
|
||||
/// Lower an aten.add to an affine loop nest.
|
||||
class AddOpConversion_affine : public ConversionPattern {
|
||||
public:
|
||||
explicit AddOpConversion_affine(MLIRContext *context)
|
||||
: ConversionPattern(mlir::NPCOMP::aten::AddOp::getOperationName(), 1, context) {
|
||||
}
|
||||
|
||||
LogicalResult
|
||||
matchAndRewrite(Operation *op, ArrayRef<Value> operands,
|
||||
ConversionPatternRewriter &rewriter) const override {
|
||||
auto add = cast<mlir::NPCOMP::aten::AddOp>(op);
|
||||
auto loc = add.getLoc();
|
||||
Type resultTy = add.getResult().getType();
|
||||
TensorType tensorResultTy = resultTy.cast<TensorType>();
|
||||
MemRefType memRefResultTy = mlir::MemRefType::get(
|
||||
tensorResultTy.getShape(), tensorResultTy.getElementType(), {}, 0);
|
||||
|
||||
Value result = rewriter.create<AllocOp>(loc, memRefResultTy);
|
||||
Value lhs = memRefTypeCast(rewriter, operands[0]);
|
||||
Value rhs = memRefTypeCast(rewriter, operands[1]);
|
||||
auto indexType = IndexType::get(op->getContext());
|
||||
|
||||
using namespace edsc;
|
||||
|
||||
ScopedContext scope(rewriter, loc);
|
||||
Value zero = intrinsics::std_constant_index(0);
|
||||
Value one = intrinsics::std_constant_index(1);
|
||||
MemRefBoundsCapture vRes(result), vLHS(lhs), vRHS(rhs);
|
||||
StdIndexedValue iRes(result), iLHS(lhs), iRHS(rhs);
|
||||
Value M(vRes.ub(0));
|
||||
if (vRes.rank() == 1) {
|
||||
affineLoopNestBuilder({zero}, {M}, 1, [&](ValueRange ivs) {
|
||||
Value i = ivs[0];
|
||||
iRes(i) = iLHS(i) + iRHS(i);
|
||||
});
|
||||
} else if (vRes.rank() == 2) {
|
||||
Value N(vRes.ub(1));
|
||||
affineLoopNestBuilder({zero, zero}, {M, N}, {1, 1}, [&](ValueRange ivs) {
|
||||
Value i = ivs[0];
|
||||
Value j = ivs[1];
|
||||
iRes(i, j) = iLHS(i, j) + iRHS(i, j);
|
||||
});
|
||||
} else if (vRes.rank() == 3) {
|
||||
Value N(vRes.ub(1));
|
||||
Value O(vRes.ub(2));
|
||||
affineLoopNestBuilder({zero, zero, zero}, {M, N, O}, {1, 1, 1},
|
||||
[&](ValueRange ivs) {
|
||||
Value i = ivs[0];
|
||||
Value j = ivs[1];
|
||||
Value k = ivs[2];
|
||||
iRes(i, j, k) = iLHS(i, j, k) + iRHS(i, j, k);
|
||||
});
|
||||
} else {
|
||||
Value N(vRes.ub(1));
|
||||
Value O(vRes.ub(2));
|
||||
Value P(vRes.ub(3));
|
||||
affineLoopNestBuilder({zero, zero, zero, zero}, {M, N, O, P},
|
||||
{1, 1, 1, 1}, [&](ValueRange ivs) {
|
||||
Value i = ivs[0];
|
||||
Value j = ivs[1];
|
||||
Value k = ivs[2];
|
||||
Value l = ivs[3];
|
||||
iRes(i, j, k, l) =
|
||||
iLHS(i, j, k, l) + iRHS(i, j, k, l);
|
||||
});
|
||||
}
|
||||
// Return the newly allocated buffer.
|
||||
rewriter.replaceOp(op, {result});
|
||||
return success();
|
||||
}
|
||||
};
|
||||
|
||||
|
||||
// Replace the given operation with a call to the given function.
|
||||
// The function is assumed to accept memrefs and scalar types and return
|
||||
// Memrefs. Here the result types are converted back to the result types of op,
|
||||
// but operands are NOT converted. This allows non-standard mappings from
|
||||
// operand types to function types.
|
||||
LogicalResult
|
||||
rewriteWithVoidFunctionCallExplicit(Operation *op,
|
||||
ArrayRef<Value> callops,
|
||||
ArrayRef<Value> operands,
|
||||
ConversionPatternRewriter &rewriter,
|
||||
std::string functionName) {
|
||||
|
||||
auto loc = op->getLoc();
|
||||
edsc::ScopedContext scope(rewriter, loc);
|
||||
|
||||
// The original operation types.
|
||||
SmallVector<Type, 8> opTys;
|
||||
// Shape erased versions of the original operation types.
|
||||
SmallVector<Type, 8> erasedOpTys;
|
||||
for (const Value &o: callops) {
|
||||
Type t = o.getType();
|
||||
opTys.push_back(t);
|
||||
if (t.isa<MemRefType>())
|
||||
erasedOpTys.push_back(getShapeErasedMemRefType(t.cast<MemRefType>()));
|
||||
else
|
||||
erasedOpTys.push_back(t);
|
||||
}
|
||||
|
||||
std::vector<Value> newOps = callops;
|
||||
SmallVector<Value, 8> newResults;
|
||||
|
||||
// Result types of the original operation, converted to memrefs.
|
||||
SmallVector<Type, 8> retTys;
|
||||
// Erased version of the return type. This is the return types of the
|
||||
// generated function call.
|
||||
SmallVector<Type, 8> erasedRetTys;
|
||||
for (const auto &o: op->getResults()) {
|
||||
Type t = o.getType();
|
||||
if (t.isa<TensorType>()) {
|
||||
TensorType tensorResultTy = t.cast<TensorType>();
|
||||
MemRefType memRefResultTy =
|
||||
mlir::MemRefType::get(tensorResultTy.getShape(),
|
||||
tensorResultTy.getElementType(), {}, 0);
|
||||
MemRefType erasedMemRefResultTy = getShapeErasedMemRefType(memRefResultTy);
|
||||
retTys.push_back(memRefResultTy);
|
||||
|
||||
// assume memRefResultTy has known shape, so we don't need any
|
||||
// dynamic dimensions for the alloc.
|
||||
assert(memRefResultTy.hasStaticShape());
|
||||
Value allocVal = rewriter.create<AllocOp>(op->getLoc(),
|
||||
memRefResultTy);
|
||||
Value castVal = memRefTypeCast(rewriter, allocVal);
|
||||
newOps.push_back(castVal);
|
||||
newResults.push_back(allocVal);
|
||||
} else {
|
||||
return failure();
|
||||
}
|
||||
}
|
||||
|
||||
SmallVector<Type, 8> empty;
|
||||
std::string mangledFunctionName = getMangledFuncName(functionName, opTys, retTys);
|
||||
FuncOp funcOp = getATenFn(op->getParentOfType<ModuleOp>(),
|
||||
mangledFunctionName,
|
||||
newOps,
|
||||
empty);
|
||||
|
||||
auto new_call = callOperation(empty,
|
||||
rewriter.getSymbolRefAttr(funcOp), newOps);
|
||||
|
||||
rewriter.replaceOp(op, newResults);
|
||||
return success();
|
||||
}
|
||||
|
||||
// Replace the given operation with a call to the given function.
|
||||
// The function is assumed to accept memrefs and scalar types and return
|
||||
// Memrefs. Other operand types (e.g. aten.list and tensor<> are converted
|
||||
// appropriately. The called function passes results of the original function
|
||||
// as memref arguments at the end of the original set of operands.
|
||||
LogicalResult
|
||||
rewriteWithFunctionCall(Operation *op, ArrayRef<Value> operands,
|
||||
ConversionPatternRewriter &rewriter,
|
||||
std::string functionName) {
|
||||
auto loc = op->getLoc();
|
||||
edsc::ScopedContext scope(rewriter, loc);
|
||||
|
||||
// Convert the arguments to the original call.
|
||||
SmallVector<Value, 8> callops;
|
||||
for (auto &o: operands) {
|
||||
Type t = o.getType();
|
||||
if (t.isa<MemRefType>()) {
|
||||
// Cast it to some memref type that we accept
|
||||
callops.push_back(memRefTypeCast(rewriter, o));
|
||||
} else if (t.isa<IntegerType>() || t.isa<FloatType>()) {
|
||||
callops.push_back(o);
|
||||
} else if (t.isa<ATenListType>()) {
|
||||
// FIXME: lots of assumptions here.
|
||||
auto unpack = [](auto &op, auto &v) -> void {
|
||||
auto co = cast<mlir::NPCOMP::aten::ConstantOp>(op.getDefiningOp());
|
||||
DenseElementsAttr a =
|
||||
co.template getAttrOfType<DenseElementsAttr>("value");
|
||||
for (auto i : a.getIntValues())
|
||||
v.push_back(i.getSExtValue());
|
||||
};
|
||||
std::vector<uint64_t> values;
|
||||
unpack(o, values);
|
||||
callops.push_back(constInt(values[0], 32));
|
||||
} else {
|
||||
return failure();
|
||||
}
|
||||
}
|
||||
return rewriteWithVoidFunctionCallExplicit(op, callops, operands, rewriter, functionName);
|
||||
}
|
||||
|
||||
|
||||
/// Lower Add
|
||||
template<typename Op>
|
||||
class ATenFunctionCallConversion : public ConversionPattern {
|
||||
public:
|
||||
explicit ATenFunctionCallConversion(MLIRContext *context)
|
||||
: ConversionPattern(Op::getOperationName(), 1, context) {
|
||||
}
|
||||
|
||||
LogicalResult
|
||||
matchAndRewrite(Operation *op, ArrayRef<Value> operands,
|
||||
ConversionPatternRewriter &rewriter) const override {
|
||||
return rewriteWithFunctionCall(op, operands, rewriter, Op::getFunctionConversionName());
|
||||
}
|
||||
};
|
||||
|
||||
/// Lower aten.constant
|
||||
class ConstantOpConversion : public ConversionPattern {
|
||||
public:
|
||||
explicit ConstantOpConversion(MLIRContext *context)
|
||||
: ConversionPattern(mlir::NPCOMP::aten::ConstantOp::getOperationName(), 1, context) {
|
||||
}
|
||||
|
||||
LogicalResult
|
||||
matchAndRewrite(Operation *op, ArrayRef<Value> operands,
|
||||
ConversionPatternRewriter &rewriter) const override {
|
||||
auto loc = op->getLoc();
|
||||
edsc::ScopedContext scope(rewriter, loc);
|
||||
|
||||
auto constOp = cast<mlir::NPCOMP::aten::ConstantOp>(op);
|
||||
|
||||
Value result = op->getResult(0);
|
||||
Type t = result.getType();
|
||||
if (t.isa<IntegerType>()) {
|
||||
auto it = t.cast<IntegerType>();
|
||||
if(it.getWidth() > 1) {
|
||||
auto a = op->getAttrOfType<IntegerAttr>("value");
|
||||
SmallVector<Value, 8> newValues {rewriter.create<mlir::ConstantOp>(loc, a)};
|
||||
rewriter.replaceOp(op, newValues);
|
||||
return success();
|
||||
} else {
|
||||
auto a = op->getAttrOfType<BoolAttr>("value");
|
||||
SmallVector<Value, 8> newValues {constInt(a.getValue(), it.getWidth())};
|
||||
rewriter.replaceOp(op, newValues);
|
||||
return success();
|
||||
}
|
||||
}
|
||||
// FIXME: support float types
|
||||
// if(t.isa<FloatType>()) {
|
||||
// APFloat f = *(a.float_value_begin());
|
||||
// rewriter.replaceOp(op, constFloat(f));
|
||||
// return success();
|
||||
// }
|
||||
return failure();
|
||||
}
|
||||
};
|
||||
|
||||
/// Lower Add
|
||||
class AddOpConversion : public ConversionPattern {
|
||||
public:
|
||||
explicit AddOpConversion(MLIRContext *context)
|
||||
: ConversionPattern(mlir::NPCOMP::aten::AddOp::getOperationName(), 1, context) {
|
||||
}
|
||||
|
||||
LogicalResult
|
||||
matchAndRewrite(Operation *op, ArrayRef<Value> operands,
|
||||
ConversionPatternRewriter &rewriter) const override {
|
||||
return rewriteWithFunctionCall(op, operands, rewriter, "add");
|
||||
}
|
||||
};
|
||||
|
||||
/// Lower Addmm
|
||||
class AddmmOpConversion : public ConversionPattern {
|
||||
public:
|
||||
explicit AddmmOpConversion(MLIRContext *context)
|
||||
: ConversionPattern(mlir::NPCOMP::aten::AddmmOp::getOperationName(), 1,
|
||||
context) {}
|
||||
|
||||
LogicalResult
|
||||
matchAndRewrite(Operation *op, ArrayRef<Value> operands,
|
||||
ConversionPatternRewriter &rewriter) const override {
|
||||
return rewriteWithFunctionCall(op, operands, rewriter, "addmm");
|
||||
}
|
||||
};
|
||||
|
||||
/// Lower AsStrided
|
||||
class AsStridedOpConversion : public ConversionPattern {
|
||||
public:
|
||||
explicit AsStridedOpConversion(MLIRContext *context)
|
||||
: ConversionPattern(mlir::NPCOMP::aten::AsStridedOp::getOperationName(), 1,
|
||||
context) {}
|
||||
|
||||
LogicalResult
|
||||
matchAndRewrite(Operation *op, ArrayRef<Value> operands,
|
||||
ConversionPatternRewriter &rewriter) const override {
|
||||
auto loc = op->getLoc();
|
||||
edsc::ScopedContext scope(rewriter, loc);
|
||||
|
||||
Value xVal = memRefTypeCast(rewriter, operands[0]);
|
||||
|
||||
// construct the shape argument
|
||||
std::vector<Value> shape;
|
||||
std::vector<int64_t> result_shape;
|
||||
auto co0 = cast<mlir::NPCOMP::aten::ConstantOp>(operands[1].getDefiningOp());
|
||||
DenseElementsAttr a0 =
|
||||
co0.template getAttrOfType<DenseElementsAttr>("value");
|
||||
for (auto i : a0.getAttributeValues())
|
||||
shape.push_back(rewriter.create<mlir::ConstantOp>(co0.getLoc(), i));
|
||||
|
||||
// pad out the shape with -1 to make it 4d
|
||||
while (shape.size() < 4)
|
||||
shape.push_back(constInt(-1, 32));
|
||||
|
||||
// construct the stride argument
|
||||
std::vector<Value> stride;
|
||||
auto co1 = cast<mlir::NPCOMP::aten::ConstantOp>(operands[2].getDefiningOp());
|
||||
DenseElementsAttr a1 =
|
||||
co1.template getAttrOfType<DenseElementsAttr>("value");
|
||||
for (auto i : a1.getAttributeValues())
|
||||
stride.push_back(rewriter.create<mlir::ConstantOp>(co1.getLoc(), i));
|
||||
|
||||
// pad out the stride with -1 to make it 4d
|
||||
while (stride.size() < 4)
|
||||
stride.push_back(constInt(-1, 32));
|
||||
|
||||
APInt offset(32, 0);
|
||||
if (operands.size() > 3) {
|
||||
auto co2 = cast<mlir::NPCOMP::aten::ConstantOp>(operands[3].getDefiningOp());
|
||||
auto ia2 = co2.getAttrOfType<IntegerAttr>("value");
|
||||
offset = ia2.getValue();
|
||||
}
|
||||
|
||||
SmallVector<Value, 8> callops{xVal, shape[0],
|
||||
shape[1], shape[2],
|
||||
shape[3], stride[0],
|
||||
stride[1], stride[2],
|
||||
stride[3], constInt(offset.getSExtValue(), 32)};
|
||||
|
||||
|
||||
return rewriteWithVoidFunctionCallExplicit(op, callops, operands, rewriter, "as_strided");
|
||||
}
|
||||
};
|
||||
|
||||
/// Lower batchnorm
|
||||
class BatchNormOpConversion : public ConversionPattern {
|
||||
public:
|
||||
explicit BatchNormOpConversion(MLIRContext *context)
|
||||
: ConversionPattern(mlir::NPCOMP::aten::BatchNormOp::getOperationName(), 1,
|
||||
context) {}
|
||||
|
||||
LogicalResult
|
||||
matchAndRewrite(Operation *op, ArrayRef<Value> operands,
|
||||
ConversionPatternRewriter &rewriter) const override {
|
||||
return rewriteWithFunctionCall(op, operands, rewriter, "batch_norm");
|
||||
}
|
||||
};
|
||||
|
||||
/// Lower conv2d
|
||||
class ConvolutionOpConversion : public ConversionPattern {
|
||||
public:
|
||||
explicit ConvolutionOpConversion(MLIRContext *context)
|
||||
: ConversionPattern(mlir::NPCOMP::aten::ConvolutionOp::getOperationName(), 1,
|
||||
context) {}
|
||||
|
||||
LogicalResult
|
||||
matchAndRewrite(Operation *op, ArrayRef<Value> operands,
|
||||
ConversionPatternRewriter &rewriter) const override {
|
||||
return rewriteWithFunctionCall(op, operands, rewriter, "conv2d");
|
||||
}
|
||||
};
|
||||
|
||||
/// Lower conv2d backward
|
||||
class ConvolutionBackwardOpConversion : public ConversionPattern {
|
||||
public:
|
||||
explicit ConvolutionBackwardOpConversion(MLIRContext *context)
|
||||
: ConversionPattern(
|
||||
mlir::NPCOMP::aten::ConvolutionBackwardOp::getOperationName(), 1,
|
||||
context) {}
|
||||
|
||||
LogicalResult
|
||||
matchAndRewrite(Operation *op, ArrayRef<Value> operands,
|
||||
ConversionPatternRewriter &rewriter) const override {
|
||||
return rewriteWithFunctionCall(op, operands, rewriter, "conv2d_backward");
|
||||
}
|
||||
};
|
||||
|
||||
/// Lower Div
|
||||
class DivOpConversion : public ConversionPattern {
|
||||
public:
|
||||
explicit DivOpConversion(MLIRContext *context)
|
||||
: ConversionPattern(mlir::NPCOMP::aten::DivOp::getOperationName(), 1, context) {
|
||||
}
|
||||
|
||||
LogicalResult
|
||||
matchAndRewrite(Operation *op, ArrayRef<Value> operands,
|
||||
ConversionPatternRewriter &rewriter) const override {
|
||||
return rewriteWithFunctionCall(op, operands, rewriter, "div");
|
||||
}
|
||||
};
|
||||
/// Lower LogSoftmax
|
||||
class LogSoftmaxOpConversion : public ConversionPattern {
|
||||
public:
|
||||
explicit LogSoftmaxOpConversion(MLIRContext *context)
|
||||
: ConversionPattern(mlir::NPCOMP::aten::LogSoftmaxOp::getOperationName(), 1,
|
||||
context) {}
|
||||
|
||||
LogicalResult
|
||||
matchAndRewrite(Operation *op, ArrayRef<Value> operands,
|
||||
ConversionPatternRewriter &rewriter) const override {
|
||||
return rewriteWithFunctionCall(op, operands, rewriter, "log_softmax");
|
||||
}
|
||||
};
|
||||
|
||||
/// Lower LogSoftmaxBackwardData
|
||||
class LogSoftmaxBackwardDataOpConversion : public ConversionPattern {
|
||||
public:
|
||||
explicit LogSoftmaxBackwardDataOpConversion(MLIRContext *context)
|
||||
: ConversionPattern(
|
||||
mlir::NPCOMP::aten::LogSoftmaxBackwardDataOp::getOperationName(), 1,
|
||||
context) {}
|
||||
|
||||
LogicalResult
|
||||
matchAndRewrite(Operation *op, ArrayRef<Value> operands,
|
||||
ConversionPatternRewriter &rewriter) const override {
|
||||
return rewriteWithFunctionCall(op, operands, rewriter,
|
||||
"log_softmax_backward_data");
|
||||
}
|
||||
};
|
||||
|
||||
/// Lower maxpool2d
|
||||
class MaxPoolOpConversion : public ConversionPattern {
|
||||
public:
|
||||
explicit MaxPoolOpConversion(MLIRContext *context)
|
||||
: ConversionPattern(mlir::NPCOMP::aten::MaxPool2dOp::getOperationName(), 1,
|
||||
context) {}
|
||||
|
||||
LogicalResult
|
||||
matchAndRewrite(Operation *op, ArrayRef<Value> operands,
|
||||
ConversionPatternRewriter &rewriter) const override {
|
||||
return rewriteWithFunctionCall(op, operands, rewriter, "max_pool2d");
|
||||
}
|
||||
};
|
||||
|
||||
/// Lower maxpool2d
|
||||
class MaxPool2dWithIndicesOpConversion : public ConversionPattern {
|
||||
public:
|
||||
explicit MaxPool2dWithIndicesOpConversion(MLIRContext *context)
|
||||
: ConversionPattern(
|
||||
mlir::NPCOMP::aten::MaxPool2dWithIndicesOp::getOperationName(), 1,
|
||||
context) {}
|
||||
|
||||
LogicalResult
|
||||
matchAndRewrite(Operation *op, ArrayRef<Value> operands,
|
||||
ConversionPatternRewriter &rewriter) const override {
|
||||
return rewriteWithFunctionCall(op, operands, rewriter, "max_pool2d_with_indices");
|
||||
}
|
||||
};
|
||||
|
||||
/// Lower max_pool2d_with_indices_backward
|
||||
class MaxPool2dWithIndicesBackwardOpConversion : public ConversionPattern {
|
||||
public:
|
||||
explicit MaxPool2dWithIndicesBackwardOpConversion(MLIRContext *context)
|
||||
: ConversionPattern(
|
||||
mlir::NPCOMP::aten::MaxPool2dWithIndicesBackwardOp::getOperationName(), 1,
|
||||
context) {}
|
||||
|
||||
LogicalResult
|
||||
matchAndRewrite(Operation *op, ArrayRef<Value> operands,
|
||||
ConversionPatternRewriter &rewriter) const override {
|
||||
return rewriteWithFunctionCall(op, operands, rewriter, "max_pool2d_with_indices_backward");
|
||||
}
|
||||
};
|
||||
|
||||
/// Lower MM
|
||||
class MMOpConversion : public ConversionPattern {
|
||||
public:
|
||||
explicit MMOpConversion(MLIRContext *context)
|
||||
: ConversionPattern(mlir::NPCOMP::aten::MmOp::getOperationName(), 1, context) {}
|
||||
|
||||
LogicalResult
|
||||
matchAndRewrite(Operation *op, ArrayRef<Value> operands,
|
||||
ConversionPatternRewriter &rewriter) const override {
|
||||
return rewriteWithFunctionCall(op, operands, rewriter, "mm");
|
||||
}
|
||||
};
|
||||
|
||||
/// Lower Mul
|
||||
class MulOpConversion : public ConversionPattern {
|
||||
public:
|
||||
explicit MulOpConversion(MLIRContext *context)
|
||||
: ConversionPattern(mlir::NPCOMP::aten::MulOp::getOperationName(), 1, context) {
|
||||
}
|
||||
|
||||
LogicalResult
|
||||
matchAndRewrite(Operation *op, ArrayRef<Value> operands,
|
||||
ConversionPatternRewriter &rewriter) const override {
|
||||
return rewriteWithFunctionCall(op, operands, rewriter, "mul");
|
||||
}
|
||||
};
|
||||
|
||||
/// Lower batchnorm
|
||||
class NativeBatchNormOpConversion : public ConversionPattern {
|
||||
public:
|
||||
explicit NativeBatchNormOpConversion(MLIRContext *context)
|
||||
: ConversionPattern(mlir::NPCOMP::aten::NativeBatchNormOp::getOperationName(),
|
||||
1, context) {}
|
||||
|
||||
LogicalResult
|
||||
matchAndRewrite(Operation *op, ArrayRef<Value> operands,
|
||||
ConversionPatternRewriter &rewriter) const override {
|
||||
return rewriteWithFunctionCall(op, operands, rewriter, "native_batch_norm");
|
||||
}
|
||||
};
|
||||
|
||||
/// lower NLL Loss backward
|
||||
class NllLoss2dBackwardOpConversion : public ConversionPattern {
|
||||
public:
|
||||
explicit NllLoss2dBackwardOpConversion(MLIRContext *context)
|
||||
: ConversionPattern(mlir::NPCOMP::aten::NllLoss2dBackwardOp::getOperationName(),
|
||||
1, context) {}
|
||||
|
||||
LogicalResult
|
||||
matchAndRewrite(Operation *op, ArrayRef<Value> operands,
|
||||
ConversionPatternRewriter &rewriter) const override {
|
||||
return rewriteWithFunctionCall(op, operands, rewriter, "nll_loss2d_backward");
|
||||
}
|
||||
};
|
||||
|
||||
/// lower NLL Loss forward
|
||||
class NllLoss2dForwardOpConversion : public ConversionPattern {
|
||||
public:
|
||||
explicit NllLoss2dForwardOpConversion(MLIRContext *context)
|
||||
: ConversionPattern(mlir::NPCOMP::aten::NllLoss2dForwardOp::getOperationName(),
|
||||
1, context) {}
|
||||
|
||||
LogicalResult
|
||||
matchAndRewrite(Operation *op, ArrayRef<Value> operands,
|
||||
ConversionPatternRewriter &rewriter) const override {
|
||||
return rewriteWithFunctionCall(op, operands, rewriter, "nll_loss2d_forward");
|
||||
}
|
||||
};
|
||||
|
||||
/// lower NLL Loss backward
|
||||
class NllLossBackwardOpConversion : public ConversionPattern {
|
||||
public:
|
||||
explicit NllLossBackwardOpConversion(MLIRContext *context)
|
||||
: ConversionPattern(mlir::NPCOMP::aten::NllLossBackwardOp::getOperationName(),
|
||||
1, context) {}
|
||||
|
||||
LogicalResult
|
||||
matchAndRewrite(Operation *op, ArrayRef<Value> operands,
|
||||
ConversionPatternRewriter &rewriter) const override {
|
||||
return rewriteWithFunctionCall(op, operands, rewriter, "nll_loss_backward");
|
||||
}
|
||||
};
|
||||
|
||||
/// lower NLL Loss forward
|
||||
class NllLossForwardOpConversion : public ConversionPattern {
|
||||
public:
|
||||
explicit NllLossForwardOpConversion(MLIRContext *context)
|
||||
: ConversionPattern(mlir::NPCOMP::aten::NllLossForwardOp::getOperationName(), 1,
|
||||
context) {}
|
||||
|
||||
LogicalResult
|
||||
matchAndRewrite(Operation *op, ArrayRef<Value> operands,
|
||||
ConversionPatternRewriter &rewriter) const override {
|
||||
return rewriteWithFunctionCall(op, operands, rewriter, "nll_loss_forward"); }
|
||||
};
|
||||
|
||||
/// Lower ReLU
|
||||
class ReLUOpConversion : public ConversionPattern {
|
||||
public:
|
||||
explicit ReLUOpConversion(MLIRContext *context)
|
||||
: ConversionPattern(mlir::NPCOMP::aten::ReluOp::getOperationName(), 1,
|
||||
context) {}
|
||||
|
||||
LogicalResult
|
||||
matchAndRewrite(Operation *op, ArrayRef<Value> operands,
|
||||
ConversionPatternRewriter &rewriter) const override {
|
||||
return rewriteWithFunctionCall(op, operands, rewriter, "relu");
|
||||
}
|
||||
};
|
||||
|
||||
/// Lower ThresholdBackward
|
||||
class ThresholdBackwardOpConversion : public ConversionPattern {
|
||||
public:
|
||||
explicit ThresholdBackwardOpConversion(MLIRContext *context)
|
||||
: ConversionPattern(mlir::NPCOMP::aten::ThresholdBackwardOp::getOperationName(),
|
||||
1, context) {}
|
||||
|
||||
LogicalResult
|
||||
matchAndRewrite(Operation *op, ArrayRef<Value> operands,
|
||||
ConversionPatternRewriter &rewriter) const override {
|
||||
return rewriteWithFunctionCall(op, operands, rewriter, "threshold_backward");
|
||||
}
|
||||
};
|
||||
|
||||
/// Lower transpose
|
||||
class TransposeOpConversion : public ConversionPattern {
|
||||
public:
|
||||
explicit TransposeOpConversion(MLIRContext *context)
|
||||
: ConversionPattern(mlir::NPCOMP::aten::TOp::getOperationName(), 1, context) {}
|
||||
|
||||
LogicalResult
|
||||
matchAndRewrite(Operation *op, ArrayRef<Value> operands,
|
||||
ConversionPatternRewriter &rewriter) const override {
|
||||
return rewriteWithFunctionCall(op, operands, rewriter, "t");
|
||||
}
|
||||
};
|
||||
|
||||
/// Lower view
|
||||
class ViewOpConversion : public ConversionPattern {
|
||||
public:
|
||||
explicit ViewOpConversion(MLIRContext *context)
|
||||
: ConversionPattern(mlir::NPCOMP::aten::ViewOp::getOperationName(), 1,
|
||||
context) {}
|
||||
|
||||
LogicalResult
|
||||
matchAndRewrite(Operation *op, ArrayRef<Value> operands,
|
||||
ConversionPatternRewriter &rewriter) const override {
|
||||
auto loc = op->getLoc();
|
||||
edsc::ScopedContext scope(rewriter, loc);
|
||||
|
||||
Value xVal = memRefTypeCast(rewriter, operands[0]);
|
||||
|
||||
// construct the shape argument
|
||||
SmallVector<Value, 8> shape;
|
||||
auto co = dyn_cast<mlir::NPCOMP::aten::ConstantOp>(operands[1].getDefiningOp());
|
||||
DenseElementsAttr a = co.template getAttrOfType<DenseElementsAttr>("value");
|
||||
for (auto i : a.getAttributeValues())
|
||||
shape.push_back(rewriter.create<mlir::ConstantOp>(co.getLoc(), i));
|
||||
|
||||
|
||||
// pad out the shape with -1 to make it 4d
|
||||
while (shape.size() < 4)
|
||||
shape.push_back(constInt(-1, 32));
|
||||
|
||||
SmallVector<Value, 8> callops{xVal, shape[0], shape[1], shape[2], shape[3]};
|
||||
|
||||
return rewriteWithVoidFunctionCallExplicit(op, callops, operands, rewriter, "view");
|
||||
}
|
||||
};
|
||||
|
||||
/// Convert an ATen type, this gets called for block and region arguments, and
|
||||
/// attributes.
|
||||
MemRefType convertTensorType(TensorType tensor) {
|
||||
return mlir::MemRefType::get(tensor.getShape(), tensor.getElementType(), {},
|
||||
0);
|
||||
}
|
||||
|
||||
/// Lower ATen to Standard dialect. Currently most of the lowerings are done
|
||||
/// through function calls, which are expected to be implemented through an
|
||||
/// external library and linked into the resulting code. In the future, the
|
||||
/// expectation is that the preferred lowering path would go through TCP.
|
||||
/// FIXME: Audit this for completeness
|
||||
struct ATenLoweringPass
|
||||
: public PassWrapper<ATenLoweringPass, OperationPass<ModuleOp>> {
|
||||
|
||||
void runOnOperation() override {
|
||||
LLVMTypeConverter typeConverter(getOperation().getContext());
|
||||
typeConverter.addConversion([&](Type type) {
|
||||
if (auto tensor = type.dyn_cast<TensorType>())
|
||||
return convertTensorType(tensor).cast<Type>();
|
||||
return type;
|
||||
});
|
||||
|
||||
OwningRewritePatternList acapPatterns;
|
||||
auto module = getOperation();
|
||||
auto context = module.getContext();
|
||||
|
||||
// c++ patterns
|
||||
acapPatterns.insert<
|
||||
ConstantOpConversion,
|
||||
AddOpConversion, ConvolutionOpConversion, ReLUOpConversion,
|
||||
TransposeOpConversion, BatchNormOpConversion,
|
||||
NativeBatchNormOpConversion, MaxPoolOpConversion,
|
||||
MaxPool2dWithIndicesOpConversion, AddmmOpConversion, ViewOpConversion,
|
||||
MulOpConversion, MMOpConversion, AsStridedOpConversion,
|
||||
LogSoftmaxOpConversion, ThresholdBackwardOpConversion,
|
||||
MaxPool2dWithIndicesBackwardOpConversion,
|
||||
ConvolutionBackwardOpConversion, NllLossForwardOpConversion,
|
||||
NllLossBackwardOpConversion, NllLoss2dForwardOpConversion,
|
||||
NllLoss2dBackwardOpConversion, LogSoftmaxOpConversion,
|
||||
LogSoftmaxBackwardDataOpConversion, DivOpConversion>(context);
|
||||
|
||||
mlir::populateFuncOpTypeConversionPattern(acapPatterns, context,
|
||||
typeConverter);
|
||||
|
||||
// tablegen patterns
|
||||
populateATenToStdPatterns(context, acapPatterns);
|
||||
|
||||
// Perform acap specific lowering.
|
||||
ConversionTarget target(getContext());
|
||||
target.addLegalDialect<LLVM::LLVMDialect, StandardOpsDialect,
|
||||
scf::SCFDialect>();
|
||||
target.addLegalOp<AffineForOp, AffineApplyOp, AffineYieldOp>();
|
||||
target.addDynamicallyLegalOp<FuncOp>([&](FuncOp op) {
|
||||
return typeConverter.isSignatureLegal(op.getType());
|
||||
});
|
||||
|
||||
if (failed(applyPartialConversion(module, target, acapPatterns))) {
|
||||
emitError(UnknownLoc::get(context), "error lowering ATen\n");
|
||||
signalPassFailure();
|
||||
}
|
||||
|
||||
// remove dead constant ops
|
||||
for (auto function : getOperation().getOps<FuncOp>()) {
|
||||
function.walk([&](Operation *op) {
|
||||
auto constOp = dyn_cast<mlir::NPCOMP::aten::ConstantOp>(op);
|
||||
if (!constOp)
|
||||
return;
|
||||
if (op->use_empty())
|
||||
op->erase();
|
||||
});
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
} // namespace
|
||||
|
||||
namespace mlir {
|
||||
namespace NPCOMP {
|
||||
namespace aten {
|
||||
|
||||
std::unique_ptr<mlir::Pass> createATenLoweringPass() {
|
||||
return std::make_unique<ATenLoweringPass>();
|
||||
}
|
||||
|
||||
} // namespace aten
|
||||
} // namespace NPCOMP
|
||||
} // namespace mlir
|
||||
|
||||
void mlir::NPCOMP::aten::registerATenLoweringPass() {
|
||||
PassRegistration<ATenLoweringPass>("aten-to-std",
|
||||
"ATen dialect lowering to function calls");
|
||||
}
|
||||
|
|
@ -0,0 +1,155 @@
|
|||
//===- ATenOpReport.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 "npcomp/Dialect/ATen/ATenOpReport.h"
|
||||
#include "npcomp/Dialect/ATen/ATenDialect.h"
|
||||
|
||||
#include "llvm/Support/Debug.h"
|
||||
#include "llvm/Support/ErrorHandling.h"
|
||||
#include "llvm/Support/JSON.h"
|
||||
#include "llvm/Support/raw_ostream.h"
|
||||
|
||||
#include "mlir/Pass/Pass.h"
|
||||
|
||||
#include <iostream>
|
||||
#include <vector>
|
||||
|
||||
#define DEBUG_TYPE "aten-op-stats"
|
||||
|
||||
using namespace mlir;
|
||||
|
||||
namespace {
|
||||
|
||||
std::string getAsString(std::map<std::string, uint64_t> &m, std::string &e) {
|
||||
return m.count(e) ? std::to_string(m[e]) : " ";
|
||||
}
|
||||
|
||||
/// Query operations through the StatisticsOpInterface and print the result
|
||||
/// in a human-readable way. This replicates the functionality in various
|
||||
/// network analysis tools and is a stepping stone toward using the information
|
||||
/// as an analysis to drive optimization.
|
||||
struct ATenOpReportPass
|
||||
: public PassWrapper<ATenOpReportPass, OperationPass<ModuleOp>> {
|
||||
|
||||
private:
|
||||
std::string *output;
|
||||
std::vector<std::string> tableFields;
|
||||
std::map<Operation *, std::string> opToName;
|
||||
|
||||
public:
|
||||
ATenOpReportPass() :
|
||||
output(nullptr),
|
||||
tableFields({"reads", "writes", "activation_in", "activation_out",
|
||||
"parameters_in", "ops:MAC", "ops:==", "ops:>", "ops:*",
|
||||
"ops:+", "ops:/", "ops:sqrt", "ops:-", "grad"}) {}
|
||||
|
||||
ATenOpReportPass(std::string *output) :
|
||||
output(output),
|
||||
tableFields({"reads", "writes", "activation_in", "activation_out",
|
||||
"parameters_in", "ops:MAC", "ops:==", "ops:>", "ops:*",
|
||||
"ops:+", "ops:/", "ops:sqrt", "ops:-", "grad"}) {}
|
||||
|
||||
|
||||
std::string emitJSONReport() {
|
||||
|
||||
llvm::json::Object top;
|
||||
|
||||
auto graph = getOperation().lookupSymbol<mlir::FuncOp>("graph");
|
||||
graph.walk([&](Operation *op) {
|
||||
if (auto stats =
|
||||
mlir::dyn_cast<mlir::NPCOMP::StatisticsOpInterface>(op)) {
|
||||
|
||||
// name for this layer
|
||||
std::string layerName = opToName[op];
|
||||
|
||||
// raw stats for this layer
|
||||
std::map<std::string, uint64_t> layerStatsMap = stats.getStatistics();
|
||||
|
||||
// JSON version of the stats we are building
|
||||
llvm::json::Object layerStatsJSON;
|
||||
|
||||
// foreach string f in tableField,
|
||||
// get the sum of all entries in layerStatsMap containing f
|
||||
for (auto &f : tableFields) {
|
||||
for (auto &p : layerStatsMap) {
|
||||
if (p.first.find(f) != std::string::npos) {
|
||||
if (auto count = layerStatsJSON[f].getAsInteger())
|
||||
layerStatsJSON[f] = (int64_t)p.second + *count;
|
||||
else
|
||||
layerStatsJSON[f] = (int64_t)p.second;
|
||||
}
|
||||
}
|
||||
}
|
||||
top[layerName] = llvm::json::Value(std::move(layerStatsJSON));
|
||||
}
|
||||
});
|
||||
|
||||
llvm::json::Value topv(std::move(top));
|
||||
std::string ret;
|
||||
llvm::raw_string_ostream ss(ret);
|
||||
ss << llvm::formatv("{0:2}", topv) << "\n";
|
||||
return ss.str();
|
||||
}
|
||||
|
||||
void runOnOperation() override {
|
||||
|
||||
// I don't change anything
|
||||
markAllAnalysesPreserved();
|
||||
|
||||
auto module = getOperation();
|
||||
|
||||
// check that a function called "graph" exists
|
||||
auto graph = module.lookupSymbol<mlir::FuncOp>("graph");
|
||||
if (!graph) {
|
||||
emitError(mlir::UnknownLoc::get(module.getContext()),
|
||||
"OpReportPass failed: can't find a graph function\n");
|
||||
signalPassFailure();
|
||||
return;
|
||||
}
|
||||
|
||||
unsigned currentLayer = 0;
|
||||
opToName.clear();
|
||||
graph.walk([&](Operation *op) {
|
||||
auto attr = op->getAttrOfType<StringAttr>("layer_name");
|
||||
if (attr)
|
||||
opToName[op] = attr.getValue().str();
|
||||
else
|
||||
opToName[op] = "unknown-layer-" + std::to_string(currentLayer);
|
||||
currentLayer++;
|
||||
});
|
||||
|
||||
std::string report = emitJSONReport();
|
||||
if(output) {
|
||||
*output = report;
|
||||
} else {
|
||||
graph.emitWarning(report);
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
} // namespace
|
||||
|
||||
namespace mlir {
|
||||
namespace NPCOMP {
|
||||
namespace aten {
|
||||
|
||||
std::unique_ptr<mlir::Pass> createATenOpReportPass() {
|
||||
return std::make_unique<ATenOpReportPass>();
|
||||
}
|
||||
std::unique_ptr<mlir::Pass> createATenOpReportPass(std::string &report) {
|
||||
return std::make_unique<ATenOpReportPass>(&report);
|
||||
}
|
||||
|
||||
void mlir::NPCOMP::aten::registerATenOpReportPass() {
|
||||
PassRegistration<ATenOpReportPass>("aten-op-report",
|
||||
"Generate ATen operation report");
|
||||
}
|
||||
|
||||
} // namespace aten
|
||||
} // namespace NPCOMP
|
||||
} // namespace mlir
|
||||
|
|
@ -0,0 +1,21 @@
|
|||
//===- ATenPasses.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 "npcomp/Dialect/ATen/ATenPasses.h"
|
||||
|
||||
using namespace mlir::NPCOMP::aten;
|
||||
|
||||
void mlir::NPCOMP::aten::registerATenPasses() {
|
||||
// TODO: Use the automatically generated pass registration.
|
||||
// #define GEN_PASS_REGISTRATION
|
||||
// #include "npcomp/Dialect/ATen/ATenPasses.h.inc"
|
||||
registerATenLayerNamePass();
|
||||
registerATenOpReportPass();
|
||||
registerATenLoweringPass();
|
||||
registerReturnEliminationPass();
|
||||
}
|
||||
|
|
@ -0,0 +1,26 @@
|
|||
//===- ATenToStd.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 "npcomp/Dialect/ATen/ATenToStd.h"
|
||||
#include "npcomp/Dialect/ATen/ATenDialect.h"
|
||||
#include "mlir/Dialect/StandardOps/IR/Ops.h"
|
||||
|
||||
using namespace mlir;
|
||||
using namespace mlir::NPCOMP;
|
||||
|
||||
namespace {
|
||||
// import patterns
|
||||
#include "npcomp/Dialect/ATen/ATenToStd.cpp.inc"
|
||||
} // namespace
|
||||
|
||||
namespace mlir {
|
||||
void populateATenToStdPatterns(MLIRContext *context,
|
||||
OwningRewritePatternList &patterns) {
|
||||
populateWithGenerated(context, &patterns);
|
||||
}
|
||||
} // namespace mlir
|
||||
|
|
@ -0,0 +1,25 @@
|
|||
add_mlir_dialect_library(NPCOMPATenDialect
|
||||
ATenDialect.cpp
|
||||
ATenDialectOpStats.cpp
|
||||
ATenPasses.cpp
|
||||
ATenLayerNamePass.cpp
|
||||
ATenLoweringPass.cpp
|
||||
ATenOpReport.cpp
|
||||
ATenToStd.cpp
|
||||
LivenessReport.cpp
|
||||
ReturnEliminationPass.cpp
|
||||
|
||||
ADDITIONAL_HEADER_DIRS
|
||||
${PROJECT_SOURCE_DIR}/dialect/include
|
||||
${PROJECT_BINARY_DIR}/dialect/include
|
||||
|
||||
DEPENDS
|
||||
MLIRATenIncGen
|
||||
MLIRATenEnumsIncGen
|
||||
MLIRATenOpInterfacesIncGen
|
||||
MLIRATenToStdIncGen
|
||||
|
||||
LINK_LIBS PUBLIC
|
||||
MLIRPass
|
||||
MLIRTransformUtils
|
||||
)
|
||||
|
|
@ -0,0 +1,253 @@
|
|||
//===- LivenessReport.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 "npcomp/Dialect/ATen/ATenDialect.h"
|
||||
|
||||
#include "mlir/Analysis/Liveness.h"
|
||||
#include "mlir/Dialect/StandardOps/IR/Ops.h"
|
||||
#include "mlir/IR/Builders.h"
|
||||
#include "mlir/IR/Module.h"
|
||||
#include "llvm/Support/Debug.h"
|
||||
#include "llvm/Support/JSON.h"
|
||||
|
||||
#include "npcomp/Dialect/ATen/LivenessReport.h"
|
||||
|
||||
#include <iostream>
|
||||
#include <sstream>
|
||||
#include <vector>
|
||||
|
||||
#define DEBUG_TYPE "liveness-report"
|
||||
|
||||
using namespace mlir;
|
||||
|
||||
namespace {
|
||||
|
||||
uint64_t getTensorVolume(const ShapedType ty) {
|
||||
|
||||
if (!ty.hasRank())
|
||||
return 1;
|
||||
|
||||
uint64_t volume = 1;
|
||||
for (auto &d : ty.getShape())
|
||||
volume *= d;
|
||||
return volume;
|
||||
}
|
||||
|
||||
uint64_t getTensorVolume(const Type ty) {
|
||||
if (auto t = ty.dyn_cast<ShapedType>()) {
|
||||
return getTensorVolume(t);
|
||||
} else {
|
||||
return 1;
|
||||
}
|
||||
}
|
||||
} // namespace
|
||||
|
||||
namespace mlir {
|
||||
namespace NPCOMP {
|
||||
namespace aten {
|
||||
|
||||
std::string LivenessReport::generateTextReport() {
|
||||
resolveLiveness();
|
||||
std::string output;
|
||||
for (auto &p : livenessIntervals) {
|
||||
Value v = p.first;
|
||||
std::vector<Operation *> &oplist = p.second;
|
||||
llvm::outs() << "// begin\n";
|
||||
v.print(llvm::outs());
|
||||
llvm::outs() << "\n";
|
||||
for (auto *o : oplist) {
|
||||
o->print(llvm::outs());
|
||||
llvm::outs() << "\n";
|
||||
}
|
||||
llvm::outs() << "// end \n";
|
||||
}
|
||||
return output;
|
||||
}
|
||||
|
||||
std::string LivenessReport::emitJSONReport() {
|
||||
resolveLiveness();
|
||||
llvm::json::Object top;
|
||||
auto context = module.getContext();
|
||||
auto loc = mlir::UnknownLoc::get(context);
|
||||
|
||||
auto graph = module.lookupSymbol<mlir::FuncOp>("graph");
|
||||
|
||||
std::map<Operation *, std::vector<Value>> liveAt;
|
||||
|
||||
graph.walk([&](Operation *op) {
|
||||
for (Value result : op->getResults()) {
|
||||
for (auto &p : livenessIntervals) {
|
||||
Value v = p.first;
|
||||
if (v == result) {
|
||||
std::vector<Operation *> &oplist = p.second;
|
||||
for (auto *o : oplist)
|
||||
liveAt[o].push_back(result);
|
||||
}
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
auto argList = graph.getBody().getBlocks().front().getArguments();
|
||||
for (Value &arg : argList) {
|
||||
for (auto &p : livenessIntervals) {
|
||||
Value v = p.first;
|
||||
if (v == arg) {
|
||||
std::vector<Operation *> &oplist = p.second;
|
||||
for (auto *o : oplist)
|
||||
liveAt[o].push_back(arg);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
graph.walk([&](Operation *op) {
|
||||
llvm::json::Object layerDetail;
|
||||
auto attr = op->getAttrOfType<StringAttr>("layer_name");
|
||||
if (!attr)
|
||||
return;
|
||||
std::vector<Value> &vlist = liveAt[op];
|
||||
int64_t parameterVol = 0;
|
||||
int64_t returnVol = 0;
|
||||
for (auto v : vlist) {
|
||||
int64_t vol = getTensorVolume(v.getType());
|
||||
if (v.getDefiningOp()) {
|
||||
if (auto a = v.getDefiningOp()->getAttrOfType<StringAttr>(
|
||||
"layer_name")) {
|
||||
auto definingOp = v.getDefiningOp();
|
||||
auto ld = layerDetail.getInteger(a.getValue().str());
|
||||
if (ld)
|
||||
layerDetail[a.getValue().str()] = *ld + vol;
|
||||
else
|
||||
layerDetail[a.getValue().str()] = vol;
|
||||
} else {
|
||||
llvm_unreachable("unknown type");
|
||||
}
|
||||
} else if (std::find(argList.begin(), argList.end(), v) !=
|
||||
argList.end()) {
|
||||
parameterVol += vol;
|
||||
} else {
|
||||
llvm_unreachable("unknown type");
|
||||
}
|
||||
auto ret = cast<ReturnOp>(op->getBlock()->getTerminator());
|
||||
for (auto oper : ret.getOperands()) {
|
||||
if (oper == v) {
|
||||
returnVol += vol;
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
if (parameterVol) {
|
||||
layerDetail["parameters"] = parameterVol;
|
||||
}
|
||||
if (returnVol) {
|
||||
layerDetail["returns"] = returnVol;
|
||||
}
|
||||
top[attr.getValue().str()] = llvm::json::Value(std::move(layerDetail));
|
||||
});
|
||||
|
||||
llvm::json::Value topv(std::move(top));
|
||||
std::string ret;
|
||||
llvm::raw_string_ostream ss(ret);
|
||||
ss << llvm::formatv("{0:2}", topv) << "\n";
|
||||
return ss.str();
|
||||
}
|
||||
|
||||
void LivenessReport::resolveLiveness() {
|
||||
|
||||
auto context = module.getContext();
|
||||
auto loc = mlir::UnknownLoc::get(context);
|
||||
|
||||
// check that a function called "graph" exists
|
||||
auto graph = module.lookupSymbol<mlir::FuncOp>("graph");
|
||||
if (!graph) {
|
||||
emitError(mlir::UnknownLoc::get(module.getContext()),
|
||||
"LivenessReport failed: can't find a graph function\n");
|
||||
return;
|
||||
}
|
||||
|
||||
// put each aten operation into its own basic block,
|
||||
// so that we can use standard liveness
|
||||
Region &bodyRegion = graph.getBody();
|
||||
Block *entryBB = &bodyRegion.getBlocks().front();
|
||||
Block *BB = entryBB;
|
||||
std::vector<Block *> new_blocks;
|
||||
while (true) {
|
||||
std::vector<Operation *> ops;
|
||||
for (Operation &op : BB->getOperations())
|
||||
ops.push_back(&op);
|
||||
|
||||
// skip over constant ops
|
||||
int idx = 0;
|
||||
while (dyn_cast<mlir::NPCOMP::aten::ConstantOp>(ops[idx]))
|
||||
idx++;
|
||||
|
||||
if (dyn_cast<ReturnOp>(ops[idx]))
|
||||
break;
|
||||
|
||||
Block *newBB = BB->splitBlock(ops[idx + 1]);
|
||||
new_blocks.push_back(newBB);
|
||||
|
||||
mlir::OpBuilder builder = mlir::OpBuilder::atBlockBegin(BB);
|
||||
builder.create<BranchOp>(loc, newBB);
|
||||
BB = newBB;
|
||||
}
|
||||
|
||||
// dump transformed function
|
||||
// graph.dump();
|
||||
|
||||
// run MLIR Liveness analysis
|
||||
auto liveness = Liveness(graph);
|
||||
|
||||
for (BlockArgument &arg :
|
||||
graph.getBody().getBlocks().front().getArguments()) {
|
||||
auto liveOps = liveness.resolveLiveness(arg);
|
||||
for (Operation *o : liveOps) {
|
||||
livenessIntervals[arg].push_back(o);
|
||||
}
|
||||
}
|
||||
|
||||
graph.walk([&](Operation *op) {
|
||||
auto attr = op->getAttrOfType<StringAttr>("layer_name");
|
||||
if (!attr)
|
||||
return;
|
||||
|
||||
for (Value v : op->getResults()) {
|
||||
auto liveOps = liveness.resolveLiveness(v);
|
||||
for (Operation *o : liveOps)
|
||||
if (auto a = o->getAttrOfType<StringAttr>("layer_name"))
|
||||
livenessIntervals[v].push_back(o);
|
||||
}
|
||||
});
|
||||
|
||||
// undo the BB insert
|
||||
auto *deadBr = bodyRegion.getBlocks().front().getTerminator();
|
||||
for (Block *b : new_blocks) {
|
||||
|
||||
auto *br = b->getTerminator();
|
||||
std::vector<Operation *> ops;
|
||||
for (auto &op : b->getOperations())
|
||||
ops.push_back(&op);
|
||||
|
||||
for (auto *op : ops) {
|
||||
if (op == br && !dyn_cast<ReturnOp>(op)) {
|
||||
op->erase();
|
||||
continue;
|
||||
}
|
||||
op->moveBefore(deadBr);
|
||||
}
|
||||
}
|
||||
deadBr->erase();
|
||||
|
||||
for (Block *b : new_blocks)
|
||||
b->erase();
|
||||
|
||||
// graph.dump();
|
||||
}
|
||||
|
||||
} // namespace aten
|
||||
} // namespace NPCOMP
|
||||
} // namespace mlir
|
||||
|
|
@ -0,0 +1,185 @@
|
|||
//===- ReturnEliminationPass.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 "npcomp/Dialect/ATen/ReturnEliminationPass.h"
|
||||
#include "npcomp/Dialect/ATen/ATenDialect.h"
|
||||
|
||||
#include "llvm/Support/Debug.h"
|
||||
#include "llvm/Support/ErrorHandling.h"
|
||||
#include "llvm/Support/raw_ostream.h"
|
||||
|
||||
#include "mlir/Dialect/StandardOps/IR/Ops.h"
|
||||
#include "mlir/Pass/Pass.h"
|
||||
|
||||
#include <set>
|
||||
#include <vector>
|
||||
|
||||
#define DEBUG_TYPE "return-elimination"
|
||||
|
||||
using namespace mlir;
|
||||
|
||||
namespace {
|
||||
|
||||
/// In the process of lowering the ATenLoweringPass generates function calls
|
||||
/// that take and return memrefs. However, this makes memory managment
|
||||
/// somewhat awkward, since we need to take care of allocating and
|
||||
/// deallocating memory. It also forces a copy to pytorch, which wants to
|
||||
/// pass in a pre-allocated buffer for the return type. To simplify the
|
||||
/// library, we convert the signature of function calls (particularly the
|
||||
/// toplevel) to pass return values by reference.
|
||||
class ReturnEliminationPass
|
||||
: public PassWrapper<ReturnEliminationPass, OperationPass<ModuleOp>> {
|
||||
|
||||
public:
|
||||
ReturnEliminationPass() {}
|
||||
|
||||
void runOn(Operation *op) {
|
||||
auto module = getOperation();
|
||||
|
||||
if (visitedOps.count(op))
|
||||
return;
|
||||
visitedOps.insert(op);
|
||||
|
||||
if (auto callOp = dyn_cast<CallOp>(op)) {
|
||||
|
||||
auto builder = std::make_unique<mlir::OpBuilder>(op);
|
||||
|
||||
std::vector<Type> tys;
|
||||
for (auto t : callOp.getCalleeType().getInputs())
|
||||
tys.push_back(t);
|
||||
for (auto t : callOp.getCalleeType().getResults())
|
||||
tys.push_back(t);
|
||||
|
||||
auto newFnTy = FunctionType::get(tys, {}, op->getContext());
|
||||
// FIXME: possible name collision
|
||||
std::string newFnName = callOp.callee().str() + "_out";
|
||||
|
||||
if (!module.lookupSymbol<FuncOp>(newFnName)) {
|
||||
auto fn = FuncOp::create(op->getLoc(), newFnName, newFnTy);
|
||||
module.push_back(fn);
|
||||
}
|
||||
|
||||
std::vector<Value> newCallArgs{callOp.arg_operand_begin(),
|
||||
callOp.arg_operand_end()};
|
||||
|
||||
for (auto v : callOp.getResults()) {
|
||||
if (!v.getType().isa<MemRefType>())
|
||||
llvm_unreachable("function returns non-memref");
|
||||
if (!valueMap.count(v)) {
|
||||
valueMap[v] = builder->create<AllocOp>(
|
||||
op->getLoc(), v.getType().cast<MemRefType>());
|
||||
}
|
||||
v.replaceAllUsesWith(valueMap[v]);
|
||||
newCallArgs.push_back(valueMap[v]);
|
||||
}
|
||||
|
||||
auto newCallOp = builder->create<CallOp>(op->getLoc(), newFnName,
|
||||
ArrayRef<Type>{}, newCallArgs);
|
||||
erasedOps.insert(op);
|
||||
auto fn = module.lookupSymbol<FuncOp>(callOp.callee());
|
||||
if (fn && fn.use_empty())
|
||||
erasedOps.insert(fn);
|
||||
} else if (isa<AllocOp>(op)) {
|
||||
Value v = op->getResult(0);
|
||||
if (valueMap.count(v)) {
|
||||
v.replaceAllUsesWith(valueMap[v]);
|
||||
erasedOps.insert(op);
|
||||
}
|
||||
}
|
||||
|
||||
for (Value v : op->getOperands()) {
|
||||
if (!v.getType().isa<MemRefType>())
|
||||
continue;
|
||||
if (v.isa<BlockArgument>())
|
||||
continue;
|
||||
if (v.getDefiningOp())
|
||||
runOn(v.getDefiningOp());
|
||||
}
|
||||
}
|
||||
|
||||
void runOnOperation() override {
|
||||
|
||||
auto module = getOperation();
|
||||
auto context = module.getContext();
|
||||
|
||||
// check that a function called "graph" exists
|
||||
auto graph = module.lookupSymbol<mlir::FuncOp>("graph");
|
||||
if (!graph) {
|
||||
emitError(mlir::UnknownLoc::get(module.getContext()),
|
||||
"OpReportPass failed: can't find a graph function\n");
|
||||
signalPassFailure();
|
||||
return;
|
||||
}
|
||||
|
||||
// assume a single bb with a single return statement
|
||||
Block &BB = graph.front();
|
||||
|
||||
FunctionType funcTy = graph.getType();
|
||||
std::vector<Type> newFuncInputTys;
|
||||
|
||||
for (auto ty : funcTy.getInputs())
|
||||
newFuncInputTys.push_back(ty);
|
||||
|
||||
for (auto ty : funcTy.getResults())
|
||||
newFuncInputTys.push_back(ty);
|
||||
|
||||
FunctionType newFuncTy =
|
||||
FunctionType::get(newFuncInputTys, {}, module.getContext());
|
||||
graph.setType(newFuncTy);
|
||||
|
||||
Operation *retOp = BB.getTerminator();
|
||||
auto builder = std::make_unique<mlir::OpBuilder>(retOp);
|
||||
|
||||
builder->create<ReturnOp>(retOp->getLoc());
|
||||
|
||||
std::vector<Value> operands{retOp->getOperands().begin(),
|
||||
retOp->getOperands().end()};
|
||||
|
||||
retOp->dropAllReferences();
|
||||
erasedOps.insert(retOp);
|
||||
|
||||
for (Value v : operands)
|
||||
valueMap[v] = BB.addArgument(v.getType());
|
||||
|
||||
for (Value v : operands) {
|
||||
if (!v.getType().isa<MemRefType>())
|
||||
llvm_unreachable("graph function returns non-memref");
|
||||
if (v.getDefiningOp())
|
||||
runOn(v.getDefiningOp());
|
||||
}
|
||||
|
||||
for (auto oi = BB.rbegin(), oe = BB.rend(); oi != oe; ++oi) {
|
||||
Operation *o = &*oi;
|
||||
for (Value v : o->getResults()) {
|
||||
if (v.getType().isa<MemRefType>()) {
|
||||
runOn(o);
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
for (Operation *o : erasedOps)
|
||||
o->erase();
|
||||
}
|
||||
|
||||
private:
|
||||
llvm::DenseMap<Value, Value> valueMap;
|
||||
std::set<Operation *> visitedOps;
|
||||
std::set<Operation *> erasedOps;
|
||||
};
|
||||
|
||||
} // namespace
|
||||
|
||||
std::unique_ptr<mlir::Pass> mlir::NPCOMP::aten::createReturnEliminationPass() {
|
||||
return std::make_unique<ReturnEliminationPass>();
|
||||
}
|
||||
|
||||
void mlir::NPCOMP::aten::registerReturnEliminationPass() {
|
||||
PassRegistration<ReturnEliminationPass>("return-elimination",
|
||||
"eliminate returns");
|
||||
}
|
||||
|
|
@ -1,3 +1,4 @@
|
|||
add_subdirectory(ATen)
|
||||
add_subdirectory(Basicpy)
|
||||
add_subdirectory(Npcomprt)
|
||||
add_subdirectory(Numpy)
|
||||
|
|
|
|||
|
|
@ -8,6 +8,8 @@
|
|||
|
||||
#include "npcomp/InitAll.h"
|
||||
|
||||
#include "npcomp/Dialect/ATen/ATenDialect.h"
|
||||
#include "npcomp/Dialect/ATen/ATenPasses.h"
|
||||
#include "npcomp/Dialect/Basicpy/IR/BasicpyDialect.h"
|
||||
#include "npcomp/Dialect/Basicpy/Transforms/Passes.h"
|
||||
#include "npcomp/Dialect/Npcomprt/IR/NpcomprtDialect.h"
|
||||
|
|
@ -68,6 +70,7 @@ static void registerDependencyPasses() {
|
|||
}
|
||||
|
||||
void mlir::NPCOMP::registerAllDialects() {
|
||||
registerDialect<mlir::NPCOMP::aten::ATenDialect>();
|
||||
registerDialect<Basicpy::BasicpyDialect>();
|
||||
registerDialect<Numpy::NumpyDialect>();
|
||||
registerDialect<npcomprt::NpcomprtDialect>();
|
||||
|
|
@ -100,5 +103,6 @@ void mlir::NPCOMP::registerAllPasses() {
|
|||
#include "npcomp/Dialect/TCF/Transforms/Passes.h.inc"
|
||||
#define GEN_PASS_REGISTRATION
|
||||
#include "npcomp/Typing/Transforms/Passes.h.inc"
|
||||
mlir::NPCOMP::aten::registerATenPasses();
|
||||
registerDependencyPasses();
|
||||
}
|
||||
|
|
|
|||
|
|
@ -0,0 +1,15 @@
|
|||
// RUN: npcomp-opt %s -aten-layer-name -aten-op-report |& FileCheck %s
|
||||
// CHECK-LABEL: "L0-add-0": {
|
||||
// CHECK-NEXT: "activation_in": 12,
|
||||
// CHECK-NEXT: "activation_out": 6,
|
||||
// CHECK-NEXT: "ops:+": 6,
|
||||
// CHECK-NEXT: "reads": 12,
|
||||
// CHECK-NEXT: "writes": 6
|
||||
|
||||
// RUN: npcomp-opt %s -aten-to-std |& FileCheck %s --check-prefix=CHECK-CONVERSION
|
||||
// CHECK-CONVERSION-LABEL: @graph
|
||||
func @graph(%arg0: tensor<1x2x3xf32>, %arg1: tensor<1x2x3xf32>) -> tensor<1x2x3xf32> {
|
||||
%1 = "aten.constant"() {type = "i32", value = 1 : i32} : () -> i32
|
||||
%2 = "aten.add"(%arg0, %arg1, %1) : (tensor<1x2x3xf32>, tensor<1x2x3xf32>, i32) -> tensor<1x2x3xf32>
|
||||
"std.return"(%2) : (tensor<1x2x3xf32>) -> ()
|
||||
}
|
||||
|
|
@ -0,0 +1,20 @@
|
|||
// RUN: npcomp-opt %s -aten-layer-name -aten-op-report |& FileCheck %s
|
||||
// CHECK-LABEL: "L1-addmm-0": {
|
||||
// CHECK-NEXT: "activation_in": 1024,
|
||||
// CHECK-NEXT: "activation_out": 16,
|
||||
// CHECK-NEXT: "ops:+": 16,
|
||||
// CHECK-NEXT: "ops:MAC": 16384,
|
||||
// CHECK-NEXT: "parameters_in": 16400,
|
||||
// CHECK-NEXT: "reads": 17424,
|
||||
// CHECK-NEXT: "writes": 16
|
||||
//
|
||||
|
||||
module {
|
||||
func @graph(%arg0: tensor<1x1024xf32>, %arg1: tensor<16x1024xf32>, %arg2: tensor<16xf32>) -> tensor<1x16xf32> {
|
||||
%0 = "aten.t"(%arg1) : (tensor<16x1024xf32>) -> tensor<1024x16xf32>
|
||||
%1 = "aten.constant"() {type = "i32", value = 1 : i32} : () -> i32
|
||||
%2 = "aten.constant"() {type = "i32", value = 1 : i32} : () -> i32
|
||||
%3 = "aten.addmm"(%arg2, %arg0, %0, %1, %2) : (tensor<16xf32>, tensor<1x1024xf32>, tensor<1024x16xf32>, i32, i32) -> tensor<1x16xf32>
|
||||
"std.return"(%3) : (tensor<1x16xf32>) -> ()
|
||||
}
|
||||
}
|
||||
|
|
@ -0,0 +1,10 @@
|
|||
// RUN: npcomp-opt %s -aten-layer-name -aten-to-std |& FileCheck %s
|
||||
// CHECK: @graph
|
||||
module {
|
||||
func @graph(%arg0: tensor<64x36864xf32>) -> tensor<64x64x24x24xf32> {
|
||||
%0 = "aten.constant"() {type = "List[i32]", value = dense<[64, 64, 24, 24]> : vector<4xi32>} : () -> !aten.list<i32>
|
||||
%1 = "aten.constant"() {type = "List[i32]", value = dense<[1, 576, 24, 1]> : vector<4xi32>} : () -> !aten.list<i32>
|
||||
%2 = "aten.as_strided"(%arg0, %0, %1) {layer_name = "L0-as_strided-0"} : (tensor<64x36864xf32>, !aten.list<i32>, !aten.list<i32>) -> tensor<64x64x24x24xf32>
|
||||
return %2 : tensor<64x64x24x24xf32>
|
||||
}
|
||||
}
|
||||
|
|
@ -0,0 +1,24 @@
|
|||
// RUN: npcomp-opt %s -aten-layer-name -aten-op-report |& FileCheck %s
|
||||
// CHECK-LABEL: "L0-batch_norm-0": {
|
||||
// CHECK-NEXT: "activation_in": 103320,
|
||||
// CHECK-NEXT: "activation_out": 103320,
|
||||
// CHECK-NEXT: "ops:*": 310206,
|
||||
// CHECK-NEXT: "ops:+": 413280,
|
||||
// CHECK-NEXT: "ops:-": 123,
|
||||
// CHECK-NEXT: "ops:/": 123,
|
||||
// CHECK-NEXT: "ops:sqrt": 123,
|
||||
// CHECK-NEXT: "parameters_in": 246,
|
||||
// CHECK-NEXT: "reads": 103566,
|
||||
// CHECK-NEXT: "writes": 103320
|
||||
|
||||
module {
|
||||
func @graph(%arg0: tensor<42x123x4x5xf32>, %arg1: tensor<123xf32>, %arg2: tensor<123xf32>, %arg3: tensor<123xf32>, %arg4: tensor<123xf32>, %arg5: tensor<?xi64>) -> tensor<42x123x4x5xf32> {
|
||||
%0 = "aten.constant"() {type = "bool", value = 0 : i1} : () -> i1
|
||||
%1 = "aten.constant"() {type = "f32", value = 1.000000e-01 : f32} : () -> f32
|
||||
%2 = "aten.constant"() {type = "f32", value = 9.99999974E-6 : f32} : () -> f32
|
||||
%3 = "aten.constant"() {type = "bool", value = 1 : i1} : () -> i1
|
||||
%4:3 = "aten.batch_norm"(%arg0, %arg1, %arg2, %arg3, %arg4, %0, %1, %2, %3) : (tensor<42x123x4x5xf32>, tensor<123xf32>, tensor
|
||||
<123xf32>, tensor<123xf32>, tensor<123xf32>, i1, f32, f32, i1) -> (tensor<42x123x4x5xf32>, tensor<123xf32>, tensor<123xf32>)
|
||||
return %4#0 : tensor<42x123x4x5xf32>
|
||||
}
|
||||
}
|
||||
|
|
@ -0,0 +1,25 @@
|
|||
// RUN: npcomp-opt %s -aten-layer-name -aten-op-report |& FileCheck %s
|
||||
// CHECK-LABEL: "L0-_convolution-0": {
|
||||
// CHECK-NEXT: "activation_in": 32768,
|
||||
// CHECK-NEXT: "activation_out": 65536,
|
||||
// CHECK-NEXT: "ops:+": 65536,
|
||||
// CHECK-NEXT: "ops:MAC": 6422528,
|
||||
// CHECK-NEXT: "parameters_in": 1584,
|
||||
// CHECK-NEXT: "reads": 34352,
|
||||
// CHECK-NEXT: "writes": 65536
|
||||
|
||||
module {
|
||||
func @graph(%arg0: tensor<1x2x128x128xf32>, %arg1: tensor<16x2x7x7xf32>, %arg2: tensor<16xf32>) -> tensor<1x16x64x64xf32> {
|
||||
%0 = "aten.constant"() {type = "List[i32]", value = dense<2> : vector<2xi64>} : () -> !aten.list<i32>
|
||||
%1 = "aten.constant"() {type = "List[i32]", value = dense<3> : vector<2xi64>} : () -> !aten.list<i32>
|
||||
%2 = "aten.constant"() {type = "List[i32]", value = dense<1> : vector<2xi64>} : () -> !aten.list<i32>
|
||||
%3 = "aten.constant"() {type = "bool", value = 0 : i1} : () -> i1
|
||||
%4 = "aten.constant"() {type = "List[i32]", value = dense<0> : vector<2xi64>} : () -> !aten.list<i32>
|
||||
%5 = "aten.constant"() {type = "i32", value = 1 : i32} : () -> i32
|
||||
%6 = "aten.constant"() {type = "bool", value = 0 : i1} : () -> i1
|
||||
%7 = "aten.constant"() {type = "bool", value = 0 : i1} : () -> i1
|
||||
%8 = "aten.constant"() {type = "bool", value = 1 : i1} : () -> i1
|
||||
%9 = "aten._convolution"(%arg0, %arg1, %arg2, %0, %1, %2) : (tensor<1x2x128x128xf32>, tensor<16x2x7x7xf32>, tensor<16xf32>, !aten.list<i32>, !aten.list<i32>, !aten.list<i32>) -> tensor<1x16x64x64xf32>
|
||||
"std.return"(%9) : (tensor<1x16x64x64xf32>) -> ()
|
||||
}
|
||||
}
|
||||
|
|
@ -0,0 +1,23 @@
|
|||
// RUN: npcomp-opt %s -aten-layer-name -aten-op-report |& FileCheck %s
|
||||
// CHECK-LABEL: "L0-convolution_backward_overrideable-0": {
|
||||
// CHECK-NEXT: "activation_in": 5568,
|
||||
// CHECK-NEXT: "grad": 5380,
|
||||
// CHECK-NEXT: "ops:+": 768,
|
||||
// CHECK-NEXT: "ops:MAC": 345600,
|
||||
// CHECK-NEXT: "parameters_in": 576,
|
||||
// CHECK-NEXT: "reads": 6144,
|
||||
// CHECK-NEXT: "writes": 5380
|
||||
// CHECK-NEXT: }
|
||||
|
||||
// RUN: npcomp-opt %s -aten-to-std |& FileCheck %s --check-prefix=CHECK-CONVERSION
|
||||
// CHECK-CONVERSION-LABEL: @graph
|
||||
module {
|
||||
func @graph(%arg0: tensor<3x4x8x8xf32>, %arg1: tensor<3x16x10x10xf32>, %arg2: tensor<4x16x3x3xf32>) -> tensor<4x16x3x3xf32> {
|
||||
%0 = "aten.constant"() {type = "List[i32]", value = dense<1> : vector<2xi32>} : () -> !aten.list<i32>
|
||||
%1 = "aten.constant"() {type = "List[i32]", value = dense<0> : vector<2xi32>} : () -> !aten.list<i32>
|
||||
%2 = "aten.constant"() {type = "bool", value = false} : () -> i1
|
||||
%3 = "aten.constant"() {type = "i32", value = 1 : i32} : () -> i32
|
||||
%10:3 = "aten.convolution_backward_overrideable"(%arg0, %arg1, %arg2, %0, %1, %0, %2, %1, %3) {layer_name = "L5-convolution_backward_overrideable-0"} : (tensor<3x4x8x8xf32>, tensor<3x16x10x10xf32>, tensor<4x16x3x3xf32>, !aten.list<i32>, !aten.list<i32>, !aten.list<i32>, i1, !aten.list<i32>, i32) -> (tensor<3x16x10x10xf32>, tensor<4x16x3x3xf32>, tensor<4xf32>)
|
||||
return %10#1 : tensor<4x16x3x3xf32>
|
||||
}
|
||||
}
|
||||
|
|
@ -0,0 +1,19 @@
|
|||
// RUN: npcomp-opt %s -aten-layer-name -aten-op-report |& FileCheck %s
|
||||
// CHECK-LABEL: "L0-max_pool2d-0": {
|
||||
// CHECK-NEXT: "activation_in": 8192,
|
||||
// CHECK-NEXT: "activation_out": 2048,
|
||||
// CHECK-NEXT: "ops:>": 16384,
|
||||
// CHECK-NEXT: "reads": 8192,
|
||||
// CHECK-NEXT: "writes": 2048
|
||||
|
||||
module {
|
||||
func @graph(%arg0: tensor<1x32x16x16xf32>) -> tensor<1x32x8x8xf32> {
|
||||
%0 = "aten.constant"() {type = "List[i32]", value = dense<3> : vector<2xi64>} : () -> !aten.list<i32>
|
||||
%1 = "aten.constant"() {type = "List[i32]", value = dense<2> : vector<2xi64>} : () -> !aten.list<i32>
|
||||
%2 = "aten.constant"() {type = "List[i32]", value = dense<1> : vector<2xi64>} : () -> !aten.list<i32>
|
||||
%3 = "aten.constant"() {type = "List[i32]", value = dense<1> : vector<2xi64>} : () -> !aten.list<i32>
|
||||
%4 = "aten.constant"() {type = "bool", value = 0 : i1} : () -> i1
|
||||
%5 = "aten.max_pool2d"(%arg0, %0, %1, %2, %3, %4) : (tensor<1x32x16x16xf32>, !aten.list<i32>, !aten.list<i32>, !aten.list<i32>, !aten.list<i32>, i1) -> tensor<1x32x8x8xf32>
|
||||
"std.return"(%5) : (tensor<1x32x8x8xf32>) -> ()
|
||||
}
|
||||
}
|
||||
|
|
@ -0,0 +1,15 @@
|
|||
// RUN: npcomp-opt %s -aten-layer-name -aten-op-report |& FileCheck %s
|
||||
// CHECK-LABEL: "L0-relu-0": {
|
||||
// CHECK-NEXT: "activation_in": 6,
|
||||
// CHECK-NEXT: "activation_out": 6,
|
||||
// CHECK-NEXT: "ops:>": 6,
|
||||
// CHECK-NEXT: "reads": 6,
|
||||
// CHECK-NEXT: "writes": 6
|
||||
|
||||
module {
|
||||
func @graph(%arg0: tensor<1x2x3xf32>) -> tensor<1x2x3xf32> {
|
||||
%0 = "aten.relu"(%arg0) : (tensor<1x2x3xf32>) -> tensor<1x2x3xf32>
|
||||
"std.return"(%0) : (tensor<1x2x3xf32>) -> ()
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -0,0 +1,61 @@
|
|||
// RUN: npcomp-opt %s -aten-layer-name -aten-op-report |& FileCheck %s
|
||||
// CHECK-LABEL: "L0-native_batch_norm-0": {
|
||||
// CHECK-LABEL: "L1-relu-0": {
|
||||
// CHECK-LABEL: "L2-_convolution-0": {
|
||||
// CHECK-LABEL: "L3-native_batch_norm-1": {
|
||||
// CHECK-LABEL: "L4-relu-1": {
|
||||
// CHECK-LABEL: "L5-_convolution-1": {
|
||||
// CHECK-LABEL: "L6-native_batch_norm-2": {
|
||||
// CHECK-LABEL: "L7-relu-2": {
|
||||
// CHECK-LABEL: "L8-_convolution-2": {
|
||||
// CHECK-LABEL: "L9-add-0": {
|
||||
|
||||
module {
|
||||
func @graph(%arg0: tensor<1x16x128x128xf32>, %arg1: tensor<1x16x128x128xf32>, %arg2: tensor<16xf32>, %arg3: tensor<16xf32>, %arg4: tensor<16xf32>, %arg5: tensor<16xf32>, %arg6: tensor<8x16x1x1xf32>, %arg7: tensor<8xf32>, %arg8: tensor<8xf32>, %arg9: tensor<8xf32>, %arg10: tensor<8xf32>, %arg11: tensor<8xf32>, %arg12: tensor<8x8x3x3xf32>, %arg13: tensor<8xf32>, %arg14: tensor<8xf32>, %arg15: tensor<8xf32>, %arg16: tensor<8xf32>, %arg17: tensor<8xf32>, %arg18: tensor<16x8x1x1xf32>, %arg19: tensor<16xf32>) -> tensor<1x16x128x128xf32> {
|
||||
%0 = "aten.constant"() {type = "bool", value = 1 : i1} : () -> i1
|
||||
%1 = "aten.constant"() {type = "f32", value = 1.000000e-01 : f32} : () -> f32
|
||||
%2 = "aten.constant"() {type = "f32", value = 9.99999974E-6 : f32} : () -> f32
|
||||
%3:3 = "aten.native_batch_norm"(%arg1, %arg2, %arg3, %arg4, %arg5, %0, %1, %2) : (tensor<1x16x128x128xf32>, tensor<16xf32>, tensor<16xf32>, tensor<16xf32>, tensor<16xf32>, i1, f32, f32) -> (tensor<1x16x128x128xf32>, tensor<16xf32>, tensor<16xf32>)
|
||||
%4 = "aten.relu"(%3#0) : (tensor<1x16x128x128xf32>) -> tensor<1x16x128x128xf32>
|
||||
%5 = "aten.constant"() {type = "List[i32]", value = dense<1> : vector<2xi32>} : () -> !aten.list<i32>
|
||||
%6 = "aten.constant"() {type = "List[i32]", value = dense<0> : vector<2xi32>} : () -> !aten.list<i32>
|
||||
%7 = "aten.constant"() {type = "List[i32]", value = dense<1> : vector<2xi32>} : () -> !aten.list<i32>
|
||||
%8 = "aten.constant"() {type = "bool", value = 0 : i1} : () -> i1
|
||||
%9 = "aten.constant"() {type = "List[i32]", value = dense<0> : vector<2xi32>} : () -> !aten.list<i32>
|
||||
%10 = "aten.constant"() {type = "i32", value = 1 : i32} : () -> i32
|
||||
%11 = "aten.constant"() {type = "bool", value = 0 : i1} : () -> i1
|
||||
%12 = "aten.constant"() {type = "bool", value = 1 : i1} : () -> i1
|
||||
%13 = "aten._convolution"(%4, %arg6, %arg7, %5, %6, %7) : (tensor<1x16x128x128xf32>, tensor<8x16x1x1xf32>, tensor<8xf32>, !aten.list<i32>, !aten.list<i32>, !aten.list<i32>) -> tensor<1x8x128x128xf32>
|
||||
%14 = "aten.constant"() {type = "bool", value = 1 : i1} : () -> i1
|
||||
%15 = "aten.constant"() {type = "f32", value = 1.000000e-01 : f32} : () -> f32
|
||||
%16 = "aten.constant"() {type = "f32", value = 9.99999974E-6 : f32} : () -> f32
|
||||
%17:3 = "aten.native_batch_norm"(%13, %arg8, %arg9, %arg10, %arg11, %14, %15, %16) : (tensor<1x8x128x128xf32>, tensor<8xf32>, tensor<8xf32>, tensor<8xf32>, tensor<8xf32>, i1, f32, f32) -> (tensor<1x8x128x128xf32>, tensor<8xf32>, tensor<8xf32>)
|
||||
%18 = "aten.relu"(%17#0) : (tensor<1x8x128x128xf32>) -> tensor<1x8x128x128xf32>
|
||||
%19 = "aten.constant"() {type = "List[i32]", value = dense<1> : vector<2xi32>} : () -> !aten.list<i32>
|
||||
%20 = "aten.constant"() {type = "List[i32]", value = dense<1> : vector<2xi32>} : () -> !aten.list<i32>
|
||||
%21 = "aten.constant"() {type = "List[i32]", value = dense<1> : vector<2xi32>} : () -> !aten.list<i32>
|
||||
%22 = "aten.constant"() {type = "bool", value = 0 : i1} : () -> i1
|
||||
%23 = "aten.constant"() {type = "List[i32]", value = dense<0> : vector<2xi32>} : () -> !aten.list<i32>
|
||||
%24 = "aten.constant"() {type = "i32", value = 1 : i32} : () -> i32
|
||||
%25 = "aten.constant"() {type = "bool", value = 0 : i1} : () -> i1
|
||||
%26 = "aten.constant"() {type = "bool", value = 1 : i1} : () -> i1
|
||||
%27 = "aten._convolution"(%18, %arg12, %arg13, %19, %20, %21) : (tensor<1x8x128x128xf32>, tensor<8x8x3x3xf32>, tensor<8xf32>, !aten.list<i32>, !aten.list<i32>, !aten.list<i32>) -> tensor<1x8x128x128xf32>
|
||||
%28 = "aten.constant"() {type = "bool", value = 1 : i1} : () -> i1
|
||||
%29 = "aten.constant"() {type = "f32", value = 1.000000e-01 : f32} : () -> f32
|
||||
%30 = "aten.constant"() {type = "f32", value = 9.99999974E-6 : f32} : () -> f32
|
||||
%31:3 = "aten.native_batch_norm"(%27, %arg14, %arg15, %arg16, %arg17, %28, %29, %30) : (tensor<1x8x128x128xf32>, tensor<8xf32>, tensor<8xf32>, tensor<8xf32>, tensor<8xf32>, i1, f32, f32) -> (tensor<1x8x128x128xf32>, tensor<8xf32>, tensor<8xf32>)
|
||||
%32 = "aten.relu"(%31#0) : (tensor<1x8x128x128xf32>) -> tensor<1x8x128x128xf32>
|
||||
%33 = "aten.constant"() {type = "List[i32]", value = dense<1> : vector<2xi32>} : () -> !aten.list<i32>
|
||||
%34 = "aten.constant"() {type = "List[i32]", value = dense<0> : vector<2xi32>} : () -> !aten.list<i32>
|
||||
%35 = "aten.constant"() {type = "List[i32]", value = dense<1> : vector<2xi32>} : () -> !aten.list<i32>
|
||||
%36 = "aten.constant"() {type = "bool", value = 0 : i1} : () -> i1
|
||||
%37 = "aten.constant"() {type = "List[i32]", value = dense<0> : vector<2xi32>} : () -> !aten.list<i32>
|
||||
%38 = "aten.constant"() {type = "i32", value = 1 : i32} : () -> i32
|
||||
%39 = "aten.constant"() {type = "bool", value = 0 : i1} : () -> i1
|
||||
%40 = "aten.constant"() {type = "bool", value = 1 : i1} : () -> i1
|
||||
%41 = "aten._convolution"(%32, %arg18, %arg19, %33, %34, %35) : (tensor<1x8x128x128xf32>, tensor<16x8x1x1xf32>, tensor<16xf32>, !aten.list<i32>, !aten.list<i32>, !aten.list<i32>) -> tensor<1x16x128x128xf32>
|
||||
%42 = "aten.constant"() {type = "i32", value = 1 : i32} : () -> i32
|
||||
%43 = "aten.add"(%arg0, %41, %42) : (tensor<1x16x128x128xf32>, tensor<1x16x128x128xf32>, i32) -> tensor<1x16x128x128xf32>
|
||||
return %43 : tensor<1x16x128x128xf32>
|
||||
}
|
||||
}
|
||||
|
|
@ -0,0 +1,26 @@
|
|||
// RUN: npcomp-opt %s -aten-to-std |& FileCheck %s --check-prefix=CHECK-CONVERSION
|
||||
// CHECK-CONVERSION-LABEL: @graph
|
||||
|
||||
module {
|
||||
func @graph(%arg0: tensor<10xf32>, %arg1: tensor<128xf32>, %arg2: tensor<4x1x28x28xf32>, %arg3: tensor<32x1x3x3xf32>, %arg4: tensor<32xf32>, %arg5: tensor<64x32x3x3xf32>, %arg6: tensor<64xf32>, %arg7: tensor<128x9216xf32>, %arg8: tensor<10x128xf32>) -> tensor<4x10xf32> {
|
||||
%0 = "aten.constant"() {type = "List[i32]", value = dense<1> : vector<2xi32>} : () -> !aten.list<i32>
|
||||
%1 = "aten.constant"() {type = "List[i32]", value = dense<0> : vector<2xi32>} : () -> !aten.list<i32>
|
||||
%2 = "aten.constant"() {type = "bool", value = false} : () -> i1
|
||||
%3 = "aten.constant"() {type = "i32", value = 1 : i32} : () -> i32
|
||||
%4 = "aten.constant"() {type = "bool", value = true} : () -> i1
|
||||
%5 = "aten._convolution"(%arg2, %arg3, %arg4, %0, %1, %0) {layer_name = "L0-_convolution-0"} : (tensor<4x1x28x28xf32>, tensor<32x1x3x3xf32>, tensor<32xf32>, !aten.list<i32>, !aten.list<i32>, !aten.list<i32>) -> tensor<4x32x26x26xf32>
|
||||
%6 = "aten.relu"(%5) {layer_name = "L1-relu-0"} : (tensor<4x32x26x26xf32>) -> tensor<4x32x26x26xf32>
|
||||
%7 = "aten._convolution"(%6, %arg5, %arg6, %0, %1, %0) {layer_name = "L2-_convolution-1"} : (tensor<4x32x26x26xf32>, tensor<64x32x3x3xf32>, tensor<64xf32>, !aten.list<i32>, !aten.list<i32>, !aten.list<i32>) -> tensor<4x64x24x24xf32>
|
||||
%8 = "aten.constant"() {type = "List[i32]", value = dense<2> : vector<2xi32>} : () -> !aten.list<i32>
|
||||
%9:2 = "aten.max_pool2d_with_indices"(%7, %8, %8, %1, %0, %2) {layer_name = "L3-max_pool2d_with_indices-0"} : (tensor<4x64x24x24xf32>, !aten.list<i32>, !aten.list<i32>, !aten.list<i32>, !aten.list<i32>, i1) -> (tensor<4x64x12x12xf32>, tensor<4x64x12x12xi64>)
|
||||
%10 = "aten.constant"() {type = "List[i32]", value = dense<[4, 9216]> : vector<2xi32>} : () -> !aten.list<i32>
|
||||
%11 = "aten.view"(%9#0, %10) {layer_name = "L4-view-0"} : (tensor<4x64x12x12xf32>, !aten.list<i32>) -> tensor<4x9216xf32>
|
||||
%12 = "aten.t"(%arg7) {layer_name = "L5-t-0"} : (tensor<128x9216xf32>) -> tensor<9216x128xf32>
|
||||
%13 = "aten.addmm"(%arg1, %11, %12, %3, %3) {layer_name = "L6-addmm-0"} : (tensor<128xf32>, tensor<4x9216xf32>, tensor<9216x128xf32>, i32, i32) -> tensor<4x128xf32>
|
||||
%14 = "aten.relu"(%13) {layer_name = "L7-relu-1"} : (tensor<4x128xf32>) -> tensor<4x128xf32>
|
||||
%15 = "aten.t"(%arg8) {layer_name = "L8-t-1"} : (tensor<10x128xf32>) -> tensor<128x10xf32>
|
||||
%16 = "aten.addmm"(%arg0, %14, %15, %3, %3) {layer_name = "L9-addmm-1"} : (tensor<10xf32>, tensor<4x128xf32>, tensor<128x10xf32>, i32, i32) -> tensor<4x10xf32>
|
||||
%17 = "aten._log_softmax"(%16, %3, %2) {layer_name = "L10-_log_softmax-0"} : (tensor<4x10xf32>, i32, i1) -> tensor<4x10xf32>
|
||||
return %17 : tensor<4x10xf32>
|
||||
}
|
||||
}
|
||||
Loading…
Reference in New Issue