mirror of https://github.com/llvm/torch-mlir
85 lines
3.8 KiB
C++
85 lines
3.8 KiB
C++
//===------------------------------------------------------------*- C++ -*-===//
|
|
//
|
|
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
|
|
// See https://llvm.org/LICENSE.txt for license information.
|
|
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
|
|
// Also available under a BSD-style license. See LICENSE.
|
|
//
|
|
//===----------------------------------------------------------------------===//
|
|
|
|
#include "mlir/Transforms/DialectConversion.h"
|
|
|
|
namespace mlir {
|
|
namespace torch {
|
|
namespace torch_to_linalg {
|
|
|
|
struct ReductionOpInfo {
|
|
bool keepDim;
|
|
Value tensorOperand;
|
|
DenseSet<int64_t> dimSet;
|
|
};
|
|
|
|
// Helper function to get the padding tensor given the padding int values.
|
|
Value getPaddedTensor(Operation *op, OpBuilder &b, Value &input,
|
|
SmallVectorImpl<int64_t> &lowPaddingInts,
|
|
SmallVectorImpl<int64_t> &highPaddingInts, Value pad);
|
|
|
|
// Helper function to get the padding tensor given the padding int values.
|
|
// It's assumed that the padding on the low end and high end are the same,
|
|
// and that zero padding is required.
|
|
Value getZeroPaddedTensor(Operation *op, OpBuilder &b, Value &input,
|
|
SmallVectorImpl<int64_t> &paddingInts);
|
|
|
|
// Helper function that adds dynamic padding to a tensor, ignoring unpaddedDims
|
|
// dimensions at the beginning. The high and low padding are the same, and the
|
|
// padding value is zero.
|
|
Value getDynamicZeroPaddedTensor(Operation *op, OpBuilder &b, Value &input,
|
|
SmallVectorImpl<Value> &padding,
|
|
int unpaddedDims = 0);
|
|
|
|
// Helper function to caculate the output tensor dims for convolution-like ops.
|
|
// Along each dim:
|
|
// dim_out =
|
|
// floor((dim_in + 2 * padding - dilation * (kernelSize - 1) - 1) / stride) + 1
|
|
Value getOutputDimForConvOps(OpBuilder &b, Location loc, Value in,
|
|
Value paddingInt, Value dilationInt,
|
|
Value kernelSizeInt, Value strideInt,
|
|
bool ceilMode = false);
|
|
|
|
// As above but for transposed convolution ops
|
|
// Along each dim:
|
|
// dim_out =
|
|
// (dim_in - 1) * stride - 2 * padding + dilation * (kernelSize - 1) + 1
|
|
Value getOutputDimForConvTransposeOps(OpBuilder &b, Location loc, Value in,
|
|
Value paddingInt, Value dilationInt,
|
|
Value kernelSizeInt, Value strideInt);
|
|
|
|
// Create a reduction of `opInfo.tensorOperand`, reducing along the dimensions
|
|
// in `opInfo.dimSet`. If `opInfo.keepDim` is true, the output tensor is the
|
|
// same rank as the `opInfo.tensorOperand` and reduced dimensions are set to
|
|
// size 1. `initElem` is the element used to initialize the output tensor where
|
|
// the reduction will be stored.
|
|
Value createReductionLinalgGeneric(
|
|
OpBuilder &b, Location loc, const ReductionOpInfo &opInfo, Value initElem,
|
|
function_ref<void(OpBuilder &, Location, ValueRange)> bodyBuild);
|
|
|
|
// Create a pointwise operation that uses values in `tensorOperands`, such that
|
|
// the element type of the resulting tensor is `resultElementType`.
|
|
Value createElementwiseLinalgGeneric(
|
|
OpBuilder &b, Location loc, ValueRange tensorOperands,
|
|
Type resultElementType,
|
|
function_ref<void(OpBuilder &, Location, ValueRange)> bodyBuild);
|
|
|
|
// Broadcasts input tensor based on the broadcastToShape.
|
|
LogicalResult broadcastToGivenShape(Operation *op, PatternRewriter &rewriter,
|
|
Value input,
|
|
SmallVector<Value> broadcastToShape,
|
|
Value &result);
|
|
|
|
// Cast a tensor to a rank-equivalent tensor of unknown size, i.e. <1x2xf32> ->
|
|
// <?x?xf32>
|
|
Value removeSizeInformation(OpBuilder &b, Location loc, Value tensor);
|
|
} // namespace torch_to_linalg
|
|
} // namespace torch
|
|
} // namespace mlir
|