Files
clang-p2996/mlir/lib/Dialect/Vector/Transforms/LowerVectorScan.cpp
Tres Popp 5550c82189 [mlir] Move casting calls from methods to function calls
The MLIR classes Type/Attribute/Operation/Op/Value support
cast/dyn_cast/isa/dyn_cast_or_null functionality through llvm's doCast
functionality in addition to defining methods with the same name.
This change begins the migration of uses of the method to the
corresponding function call as has been decided as more consistent.

Note that there still exist classes that only define methods directly,
such as AffineExpr, and this does not include work currently to support
a functional cast/isa call.

Caveats include:
- This clang-tidy script probably has more problems.
- This only touches C++ code, so nothing that is being generated.

Context:
- https://mlir.llvm.org/deprecation/ at "Use the free function variants
  for dyn_cast/cast/isa/…"
- Original discussion at https://discourse.llvm.org/t/preferred-casting-style-going-forward/68443

Implementation:
This first patch was created with the following steps. The intention is
to only do automated changes at first, so I waste less time if it's
reverted, and so the first mass change is more clear as an example to
other teams that will need to follow similar steps.

Steps are described per line, as comments are removed by git:
0. Retrieve the change from the following to build clang-tidy with an
   additional check:
   https://github.com/llvm/llvm-project/compare/main...tpopp:llvm-project:tidy-cast-check
1. Build clang-tidy
2. Run clang-tidy over your entire codebase while disabling all checks
   and enabling the one relevant one. Run on all header files also.
3. Delete .inc files that were also modified, so the next build rebuilds
   them to a pure state.
4. Some changes have been deleted for the following reasons:
   - Some files had a variable also named cast
   - Some files had not included a header file that defines the cast
     functions
   - Some files are definitions of the classes that have the casting
     methods, so the code still refers to the method instead of the
     function without adding a prefix or removing the method declaration
     at the same time.

```
ninja -C $BUILD_DIR clang-tidy

run-clang-tidy -clang-tidy-binary=$BUILD_DIR/bin/clang-tidy -checks='-*,misc-cast-functions'\
               -header-filter=mlir/ mlir/* -fix

rm -rf $BUILD_DIR/tools/mlir/**/*.inc

git restore mlir/lib/IR mlir/lib/Dialect/DLTI/DLTI.cpp\
            mlir/lib/Dialect/Complex/IR/ComplexDialect.cpp\
            mlir/lib/**/IR/\
            mlir/lib/Dialect/SparseTensor/Transforms/SparseVectorization.cpp\
            mlir/lib/Dialect/Vector/Transforms/LowerVectorMultiReduction.cpp\
            mlir/test/lib/Dialect/Test/TestTypes.cpp\
            mlir/test/lib/Dialect/Transform/TestTransformDialectExtension.cpp\
            mlir/test/lib/Dialect/Test/TestAttributes.cpp\
            mlir/unittests/TableGen/EnumsGenTest.cpp\
            mlir/test/python/lib/PythonTestCAPI.cpp\
            mlir/include/mlir/IR/
```

Differential Revision: https://reviews.llvm.org/D150123
2023-05-12 11:21:25 +02:00

252 lines
9.1 KiB
C++

//===- LowerVectorScam.cpp - Lower 'vector.scan' operation ----------------===//
//
// 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
//
//===----------------------------------------------------------------------===//
//
// This file implements target-independent rewrites and utilities to lower the
// 'vector.scan' operation.
//
//===----------------------------------------------------------------------===//
#include "mlir/Dialect/Affine/IR/AffineOps.h"
#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/Dialect/Arith/Utils/Utils.h"
#include "mlir/Dialect/Linalg/IR/Linalg.h"
#include "mlir/Dialect/MemRef/IR/MemRef.h"
#include "mlir/Dialect/SCF/IR/SCF.h"
#include "mlir/Dialect/Tensor/IR/Tensor.h"
#include "mlir/Dialect/Utils/IndexingUtils.h"
#include "mlir/Dialect/Utils/StructuredOpsUtils.h"
#include "mlir/Dialect/Vector/IR/VectorOps.h"
#include "mlir/Dialect/Vector/Transforms/LoweringPatterns.h"
#include "mlir/Dialect/Vector/Utils/VectorUtils.h"
#include "mlir/IR/BuiltinAttributeInterfaces.h"
#include "mlir/IR/BuiltinTypes.h"
#include "mlir/IR/ImplicitLocOpBuilder.h"
#include "mlir/IR/Location.h"
#include "mlir/IR/Matchers.h"
#include "mlir/IR/PatternMatch.h"
#include "mlir/IR/TypeUtilities.h"
#include "mlir/Interfaces/VectorInterfaces.h"
#include "mlir/Support/LogicalResult.h"
#define DEBUG_TYPE "vector-broadcast-lowering"
using namespace mlir;
using namespace mlir::vector;
/// This function constructs the appropriate integer or float
/// operation given the vector combining kind and operands. The
/// supported int operations are : add, mul, min (signed/unsigned),
/// max(signed/unsigned), and, or, xor. The supported float
/// operations are : add, mul, min and max.
static Value genOperator(Location loc, Value x, Value y,
vector::CombiningKind kind,
PatternRewriter &rewriter) {
using vector::CombiningKind;
auto elType = cast<VectorType>(x.getType()).getElementType();
bool isInt = elType.isIntOrIndex();
Value combinedResult{nullptr};
switch (kind) {
case CombiningKind::ADD:
if (isInt)
combinedResult = rewriter.create<arith::AddIOp>(loc, x, y);
else
combinedResult = rewriter.create<arith::AddFOp>(loc, x, y);
break;
case CombiningKind::MUL:
if (isInt)
combinedResult = rewriter.create<arith::MulIOp>(loc, x, y);
else
combinedResult = rewriter.create<arith::MulFOp>(loc, x, y);
break;
case CombiningKind::MINUI:
combinedResult = rewriter.create<arith::MinUIOp>(loc, x, y);
break;
case CombiningKind::MINSI:
combinedResult = rewriter.create<arith::MinSIOp>(loc, x, y);
break;
case CombiningKind::MAXUI:
combinedResult = rewriter.create<arith::MaxUIOp>(loc, x, y);
break;
case CombiningKind::MAXSI:
combinedResult = rewriter.create<arith::MaxSIOp>(loc, x, y);
break;
case CombiningKind::AND:
combinedResult = rewriter.create<arith::AndIOp>(loc, x, y);
break;
case CombiningKind::OR:
combinedResult = rewriter.create<arith::OrIOp>(loc, x, y);
break;
case CombiningKind::XOR:
combinedResult = rewriter.create<arith::XOrIOp>(loc, x, y);
break;
case CombiningKind::MINF:
combinedResult = rewriter.create<arith::MinFOp>(loc, x, y);
break;
case CombiningKind::MAXF:
combinedResult = rewriter.create<arith::MaxFOp>(loc, x, y);
break;
}
return combinedResult;
}
/// This function checks to see if the vector combining kind
/// is consistent with the integer or float element type.
static bool isValidKind(bool isInt, vector::CombiningKind kind) {
using vector::CombiningKind;
enum class KindType { FLOAT, INT, INVALID };
KindType type{KindType::INVALID};
switch (kind) {
case CombiningKind::MINF:
case CombiningKind::MAXF:
type = KindType::FLOAT;
break;
case CombiningKind::MINUI:
case CombiningKind::MINSI:
case CombiningKind::MAXUI:
case CombiningKind::MAXSI:
case CombiningKind::AND:
case CombiningKind::OR:
case CombiningKind::XOR:
type = KindType::INT;
break;
case CombiningKind::ADD:
case CombiningKind::MUL:
type = isInt ? KindType::INT : KindType::FLOAT;
break;
}
bool isValidIntKind = (type == KindType::INT) && isInt;
bool isValidFloatKind = (type == KindType::FLOAT) && (!isInt);
return (isValidIntKind || isValidFloatKind);
}
namespace {
/// Convert vector.scan op into arith ops and vector.insert_strided_slice /
/// vector.extract_strided_slice.
///
/// Example:
///
/// ```
/// %0:2 = vector.scan <add>, %arg0, %arg1
/// {inclusive = true, reduction_dim = 1} :
/// (vector<2x3xi32>, vector<2xi32>) to (vector<2x3xi32>, vector<2xi32>)
/// ```
///
/// is converted to:
///
/// ```
/// %cst = arith.constant dense<0> : vector<2x3xi32>
/// %0 = vector.extract_strided_slice %arg0
/// {offsets = [0, 0], sizes = [2, 1], strides = [1, 1]}
/// : vector<2x3xi32> to vector<2x1xi32>
/// %1 = vector.insert_strided_slice %0, %cst
/// {offsets = [0, 0], strides = [1, 1]}
/// : vector<2x1xi32> into vector<2x3xi32>
/// %2 = vector.extract_strided_slice %arg0
/// {offsets = [0, 1], sizes = [2, 1], strides = [1, 1]}
/// : vector<2x3xi32> to vector<2x1xi32>
/// %3 = arith.muli %0, %2 : vector<2x1xi32>
/// %4 = vector.insert_strided_slice %3, %1
/// {offsets = [0, 1], strides = [1, 1]}
/// : vector<2x1xi32> into vector<2x3xi32>
/// %5 = vector.extract_strided_slice %arg0
/// {offsets = [0, 2], sizes = [2, 1], strides = [1, 1]}
/// : vector<2x3xi32> to vector<2x1xi32>
/// %6 = arith.muli %3, %5 : vector<2x1xi32>
/// %7 = vector.insert_strided_slice %6, %4
/// {offsets = [0, 2], strides = [1, 1]}
/// : vector<2x1xi32> into vector<2x3xi32>
/// %8 = vector.shape_cast %6 : vector<2x1xi32> to vector<2xi32>
/// return %7, %8 : vector<2x3xi32>, vector<2xi32>
/// ```
struct ScanToArithOps : public OpRewritePattern<vector::ScanOp> {
using OpRewritePattern::OpRewritePattern;
LogicalResult matchAndRewrite(vector::ScanOp scanOp,
PatternRewriter &rewriter) const override {
auto loc = scanOp.getLoc();
VectorType destType = scanOp.getDestType();
ArrayRef<int64_t> destShape = destType.getShape();
auto elType = destType.getElementType();
bool isInt = elType.isIntOrIndex();
if (!isValidKind(isInt, scanOp.getKind()))
return failure();
VectorType resType = VectorType::get(destShape, elType);
Value result = rewriter.create<arith::ConstantOp>(
loc, resType, rewriter.getZeroAttr(resType));
int64_t reductionDim = scanOp.getReductionDim();
bool inclusive = scanOp.getInclusive();
int64_t destRank = destType.getRank();
VectorType initialValueType = scanOp.getInitialValueType();
int64_t initialValueRank = initialValueType.getRank();
SmallVector<int64_t> reductionShape(destShape.begin(), destShape.end());
reductionShape[reductionDim] = 1;
VectorType reductionType = VectorType::get(reductionShape, elType);
SmallVector<int64_t> offsets(destRank, 0);
SmallVector<int64_t> strides(destRank, 1);
SmallVector<int64_t> sizes(destShape.begin(), destShape.end());
sizes[reductionDim] = 1;
ArrayAttr scanSizes = rewriter.getI64ArrayAttr(sizes);
ArrayAttr scanStrides = rewriter.getI64ArrayAttr(strides);
Value lastOutput, lastInput;
for (int i = 0; i < destShape[reductionDim]; i++) {
offsets[reductionDim] = i;
ArrayAttr scanOffsets = rewriter.getI64ArrayAttr(offsets);
Value input = rewriter.create<vector::ExtractStridedSliceOp>(
loc, reductionType, scanOp.getSource(), scanOffsets, scanSizes,
scanStrides);
Value output;
if (i == 0) {
if (inclusive) {
output = input;
} else {
if (initialValueRank == 0) {
// ShapeCastOp cannot handle 0-D vectors
output = rewriter.create<vector::BroadcastOp>(
loc, input.getType(), scanOp.getInitialValue());
} else {
output = rewriter.create<vector::ShapeCastOp>(
loc, input.getType(), scanOp.getInitialValue());
}
}
} else {
Value y = inclusive ? input : lastInput;
output = genOperator(loc, lastOutput, y, scanOp.getKind(), rewriter);
assert(output != nullptr);
}
result = rewriter.create<vector::InsertStridedSliceOp>(
loc, output, result, offsets, strides);
lastOutput = output;
lastInput = input;
}
Value reduction;
if (initialValueRank == 0) {
Value v = rewriter.create<vector::ExtractOp>(loc, lastOutput, 0);
reduction =
rewriter.create<vector::BroadcastOp>(loc, initialValueType, v);
} else {
reduction = rewriter.create<vector::ShapeCastOp>(loc, initialValueType,
lastOutput);
}
rewriter.replaceOp(scanOp, {result, reduction});
return success();
}
};
} // namespace
void mlir::vector::populateVectorScanLoweringPatterns(
RewritePatternSet &patterns, PatternBenefit benefit) {
patterns.add<ScanToArithOps>(patterns.getContext(), benefit);
}