Google Git Google Git
Sign in
opensecura/3p/openxla/iree/a5c3879bdd0779de2858f8095771dcc7e353b544/./compiler/plugins/input/StableHLO/Conversion/Preprocessing/StableHLOToStableHLO.cpp
blob: 6cc62aa3896a5773900d7c105b088917c49f78ec [file]
// Copyright 2020 The IREE Authors
//
// 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
// Implements StableHLO-to-StableHLO IREE-specific preprocessing.
#include <numeric>
#include <random>
#include "compiler/plugins/input/StableHLO/Conversion/Preprocessing/Passes.h"
#include "compiler/plugins/input/StableHLO/Conversion/Preprocessing/Rewriters.h"
#include "llvm/ADT/STLExtras.h"
#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/Dialect/Func/IR/FuncOps.h"
#include "mlir/Dialect/Math/IR/Math.h"
#include "mlir/Dialect/Tensor/IR/Tensor.h"
#include "mlir/IR/BuiltinAttributes.h"
#include "mlir/IR/BuiltinTypes.h"
#include "mlir/IR/ImplicitLocOpBuilder.h"
#include "mlir/IR/PatternMatch.h"
#include "mlir/IR/TypeUtilities.h"
#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
#include "stablehlo/dialect/ChloOps.h"
#include "stablehlo/dialect/StablehloOps.h"
namespace mlir::iree_compiler::stablehlo {
#define GEN_PASS_DEF_STABLEHLOTOSTABLEHLOPREPROCESSING
#include "compiler/plugins/input/StableHLO/Conversion/Preprocessing/Passes.h.inc"
namespace {
bool isIota(ArrayRef<int64_t> array) {
for (auto [idx, value] : llvm::enumerate(array)) {
if (static_cast<int64_t>(idx) != value)
return false;
}
return true;
}
DenseI64ArrayAttr make1DElementsAttr(OpBuilder &b, ArrayRef<int64_t> integers) {
return b.getDenseI64ArrayAttr(integers);
}
Value getF32Const(ImplicitLocOpBuilder b, ArrayRef<int64_t> shapes,
ArrayRef<float> values) {
RankedTensorType ty = RankedTensorType::get(shapes, b.getF32Type());
return b
.create<mlir::stablehlo::ConstantOp>(DenseFPElementsAttr::get(ty, values))
.getResult();
}
// Guarantee that the input dimensions are ordered batch, spatial_dims, feature
// dim.
struct ReorderConvOpInputDimensions final
: OpRewritePattern<mlir::stablehlo::ConvolutionOp> {
using OpRewritePattern::OpRewritePattern;
LogicalResult matchAndRewrite(mlir::stablehlo::ConvolutionOp op,
PatternRewriter &rewriter) const override {
auto lhsType = cast<ShapedType>(op.getLhs().getType());
auto lhsShape = lhsType.getShape();
if (!lhsType.hasRank()) {
return failure();
}
auto dimensionNumbers = op.getDimensionNumbers();
auto spatialDims = dimensionNumbers.getInputSpatialDimensions();
// Compute the permutation required to create a standard order.
llvm::SmallVector<int64_t> permutations;
permutations.push_back(dimensionNumbers.getInputBatchDimension());
llvm::append_range(permutations, spatialDims);
permutations.push_back(dimensionNumbers.getInputFeatureDimension());
// If the permutation is iota then no reordering is required.
if (isIota(permutations)) {
return failure();
}
llvm::SmallVector<int64_t> transposeShape;
for (int64_t p : permutations) {
transposeShape.push_back(lhsShape[p]);
}
auto transposed = rewriter.create<mlir::stablehlo::TransposeOp>(
op.getLoc(),
RankedTensorType::get(transposeShape, lhsType.getElementType()),
op.getLhs(), rewriter.getDenseI64ArrayAttr(permutations));
llvm::SmallVector<int64_t> newSpatialDimensions(spatialDims.size());
std::iota(newSpatialDimensions.begin(), newSpatialDimensions.end(), 1);
auto newDimensionNumbers = mlir::stablehlo::ConvDimensionNumbersAttr::get(
op.getContext(),
/*input_batch_dimension=*/0,
/*input_feature_dimension=*/newSpatialDimensions.size() + 1,
/*input_spatial_dimensions=*/newSpatialDimensions,
dimensionNumbers.getKernelInputFeatureDimension(),
dimensionNumbers.getKernelOutputFeatureDimension(),
dimensionNumbers.getKernelSpatialDimensions(),
dimensionNumbers.getOutputBatchDimension(),
dimensionNumbers.getOutputFeatureDimension(),
dimensionNumbers.getOutputSpatialDimensions());
SmallVector<Value, 2> operands = {transposed, op.getRhs()};
auto newConv = rewriter.create<mlir::stablehlo::ConvolutionOp>(
op.getLoc(), op.getType(), operands, op->getAttrs());
newConv.setDimensionNumbersAttr(newDimensionNumbers);
rewriter.replaceOp(op, newConv.getResult());
return success();
}
};
struct ReorderConvOpKernelDimensions final
: OpRewritePattern<mlir::stablehlo::ConvolutionOp> {
using OpRewritePattern::OpRewritePattern;
LogicalResult matchAndRewrite(mlir::stablehlo::ConvolutionOp op,
PatternRewriter &rewriter) const override {
auto kernel = op.getRhs();
auto kernelType = cast<ShapedType>(kernel.getType());
if (!kernelType.hasRank())
return failure();
auto kernelShape = kernelType.getShape();
auto dimensionNumbers = op.getDimensionNumbers();
auto spatialDims = dimensionNumbers.getKernelSpatialDimensions();
auto inputFeatureDimension =
dimensionNumbers.getKernelInputFeatureDimension();
auto outputFeatureDimension =
dimensionNumbers.getKernelOutputFeatureDimension();
// Compute the permutation for the transpose.
llvm::SmallVector<int64_t> permutation(spatialDims.begin(),
spatialDims.end());
permutation.push_back(inputFeatureDimension);
permutation.push_back(outputFeatureDimension);
// If the permutation is iota, then no transpose is required.
if (isIota(permutation))
return failure();
llvm::SmallVector<int64_t> transposeShape;
for (int64_t perm : permutation) {
transposeShape.push_back(kernelShape[perm]);
}
llvm::SmallVector<int64_t> newSpatialDimensions(spatialDims.size());
std::iota(newSpatialDimensions.begin(), newSpatialDimensions.end(), 0);
auto transposeKernel = rewriter.create<mlir::stablehlo::TransposeOp>(
op.getLoc(),
RankedTensorType::get(transposeShape, kernelType.getElementType()),
kernel, rewriter.getDenseI64ArrayAttr(permutation));
auto newDimensionNumbers = mlir::stablehlo::ConvDimensionNumbersAttr::get(
op.getContext(), dimensionNumbers.getInputBatchDimension(),
dimensionNumbers.getInputFeatureDimension(),
dimensionNumbers.getInputSpatialDimensions(),
/*kernel_input_feature_dimension=*/
newSpatialDimensions.size(),
/*kernel_output_feature_dimension=*/
newSpatialDimensions.size() + 1, newSpatialDimensions,
dimensionNumbers.getOutputBatchDimension(),
dimensionNumbers.getOutputFeatureDimension(),
dimensionNumbers.getOutputSpatialDimensions());
SmallVector<Value, 2> operands = {op.getLhs(), transposeKernel};
mlir::stablehlo::ConvolutionOp newConv =
rewriter.create<mlir::stablehlo::ConvolutionOp>(
op.getLoc(), op.getType(), operands, op->getAttrs());
newConv.setDimensionNumbersAttr(newDimensionNumbers);
rewriter.replaceOp(op, {newConv.getResult()});
return success();
}
};
// Guarantee that the output dimensions are ordered batch, spatial_dims, feature
// dim.
struct ReorderConvOpOutputDimensions final
: OpRewritePattern<mlir::stablehlo::ConvolutionOp> {
using OpRewritePattern::OpRewritePattern;
LogicalResult matchAndRewrite(mlir::stablehlo::ConvolutionOp op,
PatternRewriter &rewriter) const override {
auto resultType = llvm::cast<ShapedType>(op.getType());
auto resultShape = resultType.getShape();
if (!resultType.hasRank()) {
return failure();
}
auto dimensionNumbers = op.getDimensionNumbers();
auto spatialDims = dimensionNumbers.getOutputSpatialDimensions();
// Compute the permutation to transpose to an ordered output.
llvm::SmallVector<int64_t> permutation;
permutation.push_back(dimensionNumbers.getOutputBatchDimension());
permutation.append(spatialDims.begin(), spatialDims.end());
permutation.push_back(dimensionNumbers.getOutputFeatureDimension());
// If the permutation is iota then no reordering is required.
if (isIota(permutation)) {
return failure();
}
// Compute what the new conv shape should be.
llvm::SmallVector<int64_t> convShape;
for (auto p : permutation) {
convShape.push_back(resultShape[p]);
}
// Compute the inverse transpose to unordered and ordered output.
llvm::SmallVector<int64_t> invertPermutation(permutation.size());
for (auto it : llvm::enumerate(permutation)) {
invertPermutation[it.value()] = it.index();
}
llvm::SmallVector<int64_t> newSpatialDimensions(spatialDims.size());
std::iota(newSpatialDimensions.begin(), newSpatialDimensions.end(), 1);
auto newDimensionNumbers = mlir::stablehlo::ConvDimensionNumbersAttr::get(
op.getContext(), dimensionNumbers.getInputBatchDimension(),
dimensionNumbers.getInputFeatureDimension(),
dimensionNumbers.getInputSpatialDimensions(),
dimensionNumbers.getKernelInputFeatureDimension(),
dimensionNumbers.getKernelOutputFeatureDimension(),
dimensionNumbers.getKernelSpatialDimensions(),
/*output_batch_dimension=*/0,
/*output_feature_dimension=*/newSpatialDimensions.size() + 1,
/*output_spatial_dimensions=*/newSpatialDimensions);
SmallVector<Value, 2> operands = {op.getLhs(), op.getRhs()};
auto newConv = rewriter.create<mlir::stablehlo::ConvolutionOp>(
op.getLoc(),
RankedTensorType::get(convShape, resultType.getElementType()), operands,
op->getAttrs());
newConv.setDimensionNumbersAttr(newDimensionNumbers);
auto transposed = rewriter.create<mlir::stablehlo::TransposeOp>(
op.getLoc(), resultType, newConv,
rewriter.getDenseI64ArrayAttr(invertPermutation));
rewriter.replaceOp(op, transposed.getResult());
return success();
}
};
bool isConsecutive(ArrayRef<int64_t> array) {
for (size_t i = 1, e = array.size(); i < e; ++i) {
if (array[i] - array[i - 1] != 1)
return false;
}
return true;
}
/// Rewrites stablehlo.dot_general so lhs contraction dimensions are innermost
/// and rhs contraction dimensions are dims right after batch dimension. The
/// pattern inserts transposes so the dot_general always has the form:
/// {batch_dims, parallel_dims, contraction_dims}.
/// {batch_dims, contraction_dims, parallel_dims}
/// After that, batch_dims, contraction_dims, parallel_dims are
/// in consecutive order and not spliting the domain. This pattern inserts
/// reshapes to collapse consecutive reduction and parallel dims to always
/// generate a rank-3 dot_general op.
struct TransposeReshapeGenericDotGeneral final
: OpRewritePattern<mlir::stablehlo::DotGeneralOp> {
using OpRewritePattern::OpRewritePattern;
Value TransposeIfNonConsecutive(OpBuilder &b, Location loc, Value src,
ArrayRef<int64_t> targetOrder) const {
if (isConsecutive(targetOrder))
return src;
auto type = cast<RankedTensorType>(src.getType());
SmallVector<int64_t> transposeShape;
for (int64_t i : targetOrder) {
transposeShape.push_back(type.getDimSize(i));
}
return b.create<mlir::stablehlo::TransposeOp>(
loc, RankedTensorType::get(transposeShape, type.getElementType()), src,
b.getDenseI64ArrayAttr(targetOrder));
}
Value ReshapeIfNonStandard(OpBuilder &b, Location loc, Value src,
size_t dimsBorder0, size_t dimsBorder1) const {
auto type = cast<RankedTensorType>(src.getType());
ArrayRef<int64_t> shape = type.getShape();
if (dimsBorder0 <= 1 && dimsBorder1 - dimsBorder0 <= 1 &&
shape.size() - dimsBorder1 <= 1)
return src;
SmallVector<int64_t> result_shape = {
std::accumulate(shape.begin(), shape.begin() + dimsBorder0, 1,
std::multiplies<int64_t>()),
std::accumulate(shape.begin() + dimsBorder0,
shape.begin() + dimsBorder1, 1,
std::multiplies<int64_t>()),
std::accumulate(shape.begin() + dimsBorder1, shape.end(), 1,
std::multiplies<int64_t>())};
return b.create<mlir::stablehlo::ReshapeOp>(
loc, RankedTensorType::get(result_shape, type.getElementType()), src);
}
LogicalResult matchAndRewrite(mlir::stablehlo::DotGeneralOp op,
PatternRewriter &rewriter) const override {
auto lhsShapeType = dyn_cast<RankedTensorType>(op.getLhs().getType());
auto rhsShapeType = dyn_cast<RankedTensorType>(op.getRhs().getType());
auto resultType = dyn_cast<RankedTensorType>(op.getResult().getType());
if (!lhsShapeType || !rhsShapeType || !resultType)
return failure();
// TODO(jpienaar): This pattern is not safe for dynamic shapes and seems to
// be (now) redundant with later pass that does handle them. To decouple
// fixing and verifying redundant, this just limits to static shapes and
// then will remove this in follow up.
if (!lhsShapeType.hasStaticShape() || !rhsShapeType.hasStaticShape())
return failure();
SmallVector<int64_t> lhsTargetOrder, rhsTargetOrder;
mlir::stablehlo::DotDimensionNumbersAttr dimNumbers =
op.getDotDimensionNumbers();
ArrayRef<int64_t> lhsBatchingDims = dimNumbers.getLhsBatchingDimensions();
ArrayRef<int64_t> lhsContractingDims =
dimNumbers.getLhsContractingDimensions();
ArrayRef<int64_t> rhsBatchingDims = dimNumbers.getRhsBatchingDimensions();
ArrayRef<int64_t> rhsContractingDims =
dimNumbers.getRhsContractingDimensions();
// No contraction dims means this can be represented as a mul.
if (lhsContractingDims.empty() || rhsContractingDims.empty()) {
return rewriter.notifyMatchFailure(op,
"can be represented as stablehlo.mul");
}
// No batching dimensions means this can be represented a dot.
if (lhsBatchingDims.empty() || rhsBatchingDims.empty()) {
return rewriter.notifyMatchFailure(op,
"can be represented as stablehlo.dot");
}
SmallVector<bool> isLhsParallel(lhsShapeType.getRank(), true);
for (int64_t i : lhsBatchingDims) {
lhsTargetOrder.push_back(i);
isLhsParallel[i] = false;
}
for (int64_t i : lhsContractingDims) {
isLhsParallel[i] = false;
}
for (int64_t i = 0, e = lhsShapeType.getRank(); i < e; ++i) {
if (isLhsParallel[i]) {
lhsTargetOrder.push_back(i);
}
}
for (int64_t i : lhsContractingDims) {
lhsTargetOrder.push_back(i);
}
SmallVector<bool> isRhsParallel(rhsShapeType.getRank(), true);
for (int64_t i : rhsBatchingDims) {
rhsTargetOrder.push_back(i);
isRhsParallel[i] = false;
}
for (int64_t i : rhsContractingDims) {
rhsTargetOrder.push_back(i);
isRhsParallel[i] = false;
}
for (int64_t i = 0, e = rhsShapeType.getRank(); i < e; ++i) {
if (isRhsParallel[i]) {
rhsTargetOrder.push_back(i);
}
}
Value lhs = TransposeIfNonConsecutive(rewriter, op.getLoc(), op.getLhs(),
lhsTargetOrder);
Value rhs = TransposeIfNonConsecutive(rewriter, op.getLoc(), op.getRhs(),
rhsTargetOrder);
// The dimensions of this will always be transposed into {batch_dims,
// parallel_dims, contraction_dims}, and the
// following logic is based on this assumption.
// TODO(#7443): If we consider transpose performance, the above assumptions
// may not be true.
int64_t numLhsContractionDims = lhsContractingDims.size();
int64_t lhsContractionBase = lhsShapeType.getRank() - numLhsContractionDims;
int64_t rhsContractionBase = rhsBatchingDims.size();
int64_t numRhsContractionDims =
rhsContractionBase + rhsContractingDims.size();
lhs = ReshapeIfNonStandard(rewriter, op.getLoc(), lhs,
lhsBatchingDims.size(), lhsContractionBase);
rhs = ReshapeIfNonStandard(rewriter, op.getLoc(), rhs,
rhsBatchingDims.size(), numRhsContractionDims);
if (lhs == op.getLhs() && rhs == op.getRhs())
return rewriter.notifyMatchFailure(op, "already in canonical form");
auto dimensionNumbers = mlir::stablehlo::DotDimensionNumbersAttr::get(
rewriter.getContext(), /*lhsBatchingDimensions=*/0,
/*rhsBatchingDimensions=*/0,
/*lhsContractingDimensions=*/
cast<ShapedType>(lhs.getType()).getRank() - 1,
/*rhsContractingDimensions=*/1);
auto lhsNewType = cast<RankedTensorType>(lhs.getType());
auto rhsNewType = cast<RankedTensorType>(rhs.getType());
// batching、lhs parallel、rhs parallel this order is a conversion
SmallVector<int64_t, 3> newShape = {lhsNewType.getShape()[0]};
if (lhsNewType.getRank() > 2)
newShape.push_back(lhsNewType.getDimSize(1));
if (rhsNewType.getRank() > 2)
newShape.push_back(rhsNewType.getDimSize(2));
TensorType newResultType =
RankedTensorType::get(newShape, resultType.getElementType());
auto newOp = rewriter.create<mlir::stablehlo::DotGeneralOp>(
op.getLoc(), newResultType, lhs, rhs, dimensionNumbers,
op.getPrecisionConfigAttr(), op.getAlgorithmAttr());
// Copy over unknown attributes as we currently rely on it to let user tune
// lowering parameters.
ArrayRef<StringRef> odsAttrs = op.getAttributeNames();
for (NamedAttribute kv : op->getAttrs()) {
if (!llvm::is_contained(odsAttrs, kv.getName().getValue())) {
newOp->setAttr(kv.getName(), kv.getValue());
}
}
Value result = newOp.getResult();
if (op.getType() != newResultType) {
result = rewriter.create<mlir::stablehlo::ReshapeOp>(
op.getLoc(), op.getType(), newOp.getResult());
}
rewriter.replaceOp(op, result);
return success();
}
};
struct ScatterInt64Indices final
: OpRewritePattern<mlir::stablehlo::ScatterOp> {
using OpRewritePattern::OpRewritePattern;
LogicalResult matchAndRewrite(mlir::stablehlo::ScatterOp op,
PatternRewriter &rewriter) const override {
auto indices = op.getScatterIndices();
auto indicesTy = indices.getType();
auto indicesETy = indicesTy.getElementType();
if (indicesETy.isInteger(32)) {
return rewriter.notifyMatchFailure(op, "already has i32 index type");
}
if (!indicesTy.hasStaticShape()) {
return rewriter.notifyMatchFailure(op, "cannot validate legal size");
}
uint64_t maxSize = std::numeric_limits<int32_t>::max();
if (indicesETy.getIntOrFloatBitWidth() > 32) {
for (int i = 0, s = indicesTy.getRank(); i < s; ++i) {
if (indicesTy.getDimSize(i) > maxSize) {
return rewriter.notifyMatchFailure(op, "index may exceed i32 max");
}
}
}
indices = rewriter.create<mlir::stablehlo::ConvertOp>(
op.getLoc(), indicesTy.clone(rewriter.getI32Type()), indices);
auto newScatter = rewriter.create<mlir::stablehlo::ScatterOp>(
op.getLoc(), op.getResultTypes(), op.getInputs(), indices,
op.getUpdates(), op.getScatterDimensionNumbers(),
op.getIndicesAreSorted(), op.getUniqueIndices());
Region &region = newScatter.getUpdateComputation();
rewriter.cloneRegionBefore(op.getUpdateComputation(), region, region.end());
rewriter.replaceOp(op, newScatter.getResults());
return success();
}
};
// If the indices tensor has an implicit index vector dim we expand and make it
// an explicit dim.
struct ScatterImplicitIndex final
: OpRewritePattern<mlir::stablehlo::ScatterOp> {
using OpRewritePattern::OpRewritePattern;
LogicalResult matchAndRewrite(mlir::stablehlo::ScatterOp op,
PatternRewriter &rewriter) const override {
auto dimNumbers = op.getScatterDimensionNumbers();
auto indexVectorDim = dimNumbers.getIndexVectorDim();
Value indices = op.getScatterIndices();
auto indicesTy = llvm::cast<ShapedType>(indices.getType());
// Check indices vector has an implicit dim.
if (indexVectorDim != indicesTy.getRank()) {
return rewriter.notifyMatchFailure(op, "no implicit index dim");
}
// Materialize the implicit indices dim.
SmallVector<ReassociationExprs> reassociationMap;
reassociationMap.resize(indicesTy.getRank());
SmallVector<int64_t> newShape;
for (int i = 0, s = indicesTy.getRank(); i < s; ++i) {
reassociationMap[i].push_back(rewriter.getAffineDimExpr(i));
newShape.push_back(indicesTy.getDimSize(i));
}
if (!reassociationMap.empty()) {
reassociationMap.back().push_back(
rewriter.getAffineDimExpr(indicesTy.getRank()));
}
newShape.push_back(1);
indicesTy = RankedTensorType::get(newShape, indicesTy.getElementType());
indices = rewriter.create<tensor::ExpandShapeOp>(op.getLoc(), indicesTy,
indices, reassociationMap);
auto newScatter = rewriter.create<mlir::stablehlo::ScatterOp>(
op.getLoc(), op.getResultTypes(), op.getInputs(), indices,
op.getUpdates(), dimNumbers, op.getIndicesAreSorted(),
op.getUniqueIndices());
Region &region = newScatter.getUpdateComputation();
rewriter.cloneRegionBefore(op.getUpdateComputation(), region, region.end());
rewriter.replaceOp(op, newScatter.getResults());
return success();
}
};
struct ScatterImplicitBatch final
: OpRewritePattern<mlir::stablehlo::ScatterOp> {
using OpRewritePattern::OpRewritePattern;
static Value addUnitBatchDim(Location loc, Value value,
PatternRewriter &rewriter) {
auto valueTy = cast<ShapedType>(value.getType());
if (!valueTy.hasRank())
return nullptr;
// Materialize the implicit indices dim.
SmallVector<ReassociationExprs> reassociationMap(valueTy.getRank());
if (!reassociationMap.empty()) {
reassociationMap.front().push_back(rewriter.getAffineDimExpr(0));
}
SmallVector<int64_t> newShape = {1};
for (int i = 0, s = valueTy.getRank(); i < s; i++) {
reassociationMap[i].push_back(rewriter.getAffineDimExpr(i + 1));
newShape.push_back(valueTy.getDimSize(i));
}
valueTy = RankedTensorType::get(newShape, valueTy.getElementType());
return rewriter.create<tensor::ExpandShapeOp>(loc, valueTy, value,
reassociationMap);
}
LogicalResult matchAndRewrite(mlir::stablehlo::ScatterOp op,
PatternRewriter &rewriter) const override {
auto dimNumbers = op.getScatterDimensionNumbers();
auto indexVectorDim = dimNumbers.getIndexVectorDim();
auto indices = llvm::cast<Value>(op.getScatterIndices());
auto indicesTy = llvm::dyn_cast<RankedTensorType>(indices.getType());
// Check whether indices has no batch dimension.
if (!indicesTy)
return failure();
if (indicesTy.getRank() != 1 || indexVectorDim != 0) {
return rewriter.notifyMatchFailure(op,
"no implicit batch dimension to add.");
}
indices = addUnitBatchDim(op.getLoc(), indices, rewriter);
if (!indices) {
return rewriter.notifyMatchFailure(
op, "Unable to add implicit batch dim to indices.");
}
llvm::SmallVector<int64_t> newUpdateWindowDims;
for (auto dim : dimNumbers.getUpdateWindowDims()) {
// Batch dimension is inserted at the start so window dimensions are shift
// forwards.
newUpdateWindowDims.push_back(dim + 1);
}
llvm::SmallVector<Value> updates;
for (Value update : op.getUpdates()) {
update = addUnitBatchDim(op.getLoc(), update, rewriter);
if (!update) {
return rewriter.notifyMatchFailure(
op, "Unable to add implicit batch dim to update.");
}
updates.push_back(update);
}
auto newDimNumbers = mlir::stablehlo::ScatterDimensionNumbersAttr::get(
op.getContext(), newUpdateWindowDims,
dimNumbers.getInsertedWindowDims(), dimNumbers.getInputBatchingDims(),
dimNumbers.getScatterIndicesBatchingDims(),
dimNumbers.getScatterDimsToOperandDims(),
dimNumbers.getIndexVectorDim() + 1);
auto newScatter = rewriter.create<mlir::stablehlo::ScatterOp>(
op.getLoc(), op.getResultTypes(), op.getInputs(), indices, updates,
newDimNumbers, op.getIndicesAreSorted(), op.getUniqueIndices());
Region &region = newScatter.getUpdateComputation();
rewriter.cloneRegionBefore(op.getUpdateComputation(), region, region.end());
rewriter.replaceOp(op, newScatter.getResults());
return success();
}
};
struct ScatterCollapseBatch final
: OpRewritePattern<mlir::stablehlo::ScatterOp> {
using OpRewritePattern::OpRewritePattern;
static Value collapseBatchDims(Location loc, Value value, int64_t batchCount,
PatternRewriter &rewriter) {
auto valueTy = dyn_cast<ShapedType>(value.getType());
if (!valueTy)
return nullptr;
SmallVector<ReassociationExprs> reassociationMap(1);
reassociationMap.reserve(valueTy.getRank() - batchCount + 1);
int64_t batchSize = 1;
for (int i = 0, s = batchCount; i < s; i++) {
reassociationMap.front().push_back(rewriter.getAffineDimExpr(i));
bool isDynamic =
valueTy.isDynamicDim(i) || ShapedType::isDynamic(batchSize);
batchSize =
isDynamic ? ShapedType::kDynamic : valueTy.getDimSize(i) * batchSize;
}
SmallVector<int64_t> newShape = {batchSize};
for (int i = batchCount, s = valueTy.getRank(); i < s; i++) {
reassociationMap.push_back({rewriter.getAffineDimExpr(i)});
newShape.push_back(valueTy.getDimSize(i));
}
valueTy = RankedTensorType::get(newShape, valueTy.getElementType());
return rewriter.create<tensor::CollapseShapeOp>(loc, valueTy, value,
reassociationMap);
}
LogicalResult matchAndRewrite(mlir::stablehlo::ScatterOp op,
PatternRewriter &rewriter) const override {
auto dimNumbers = op.getScatterDimensionNumbers();
auto indexVectorDim = dimNumbers.getIndexVectorDim();
auto indices = llvm::cast<Value>(op.getScatterIndices());
auto indicesTy = llvm::cast<ShapedType>(indices.getType());
auto updatedWindowDims = dimNumbers.getUpdateWindowDims();
if (!indicesTy.hasRank()) {
return rewriter.notifyMatchFailure(op, "indices has unknown rank");
}
// Check for an explicit index dimension.
if (indexVectorDim != indicesTy.getRank() - 1) {
return rewriter.notifyMatchFailure(op, "no explicit indices dimension");
}
// Check that there are multiple batch dimensions.
if (indicesTy.getRank() < 3) {
return rewriter.notifyMatchFailure(op, "no multiple batch dimensions");
}
const int64_t batchCount = indicesTy.getRank() - 1;
for (auto it : llvm::enumerate(updatedWindowDims)) {
if (it.index() != it.value() - batchCount) {
return rewriter.notifyMatchFailure(
op, "update windows should be at the end.");
}
}
indices = collapseBatchDims(op.getLoc(), indices, batchCount, rewriter);
if (!indices) {
return rewriter.notifyMatchFailure(op,
"cannot collapse indices batch dims");
}
llvm::SmallVector<Value> updates;
for (Value update : op.getUpdates()) {
update = collapseBatchDims(op.getLoc(), update, batchCount, rewriter);
if (!update) {
return rewriter.notifyMatchFailure(op,
"cannot collapse update batch dims");
}
updates.push_back(update);
}
llvm::SmallVector<int64_t> newUpdatedWindowDims;
for (auto dim : updatedWindowDims) {
newUpdatedWindowDims.push_back(dim - batchCount + 1);
}
auto newDimNumbers = mlir::stablehlo::ScatterDimensionNumbersAttr::get(
op.getContext(), newUpdatedWindowDims,
dimNumbers.getInsertedWindowDims(), dimNumbers.getInputBatchingDims(),
dimNumbers.getScatterIndicesBatchingDims(),
dimNumbers.getScatterDimsToOperandDims(),
/*indexVectorDim=*/1);
auto newScatter = rewriter.create<mlir::stablehlo::ScatterOp>(
op.getLoc(), op.getResultTypes(), op.getInputs(), indices, updates,
newDimNumbers, op.getIndicesAreSorted(), op.getUniqueIndices());
Region &region = newScatter.getUpdateComputation();
rewriter.cloneRegionBefore(op.getUpdateComputation(), region, region.end());
rewriter.replaceOp(op, newScatter.getResults());
return success();
}
};
// Ensure the batch dimensions of both the indices and updates are the first
// dimensions. If they are not, transpose them to the start.
struct ScatterBatchFirst final : OpRewritePattern<mlir::stablehlo::ScatterOp> {
using OpRewritePattern::OpRewritePattern;
LogicalResult matchAndRewrite(mlir::stablehlo::ScatterOp op,
PatternRewriter &rewriter) const override {
ImplicitLocOpBuilder builder(op.getLoc(), rewriter);
auto dimNumbers = op.getScatterDimensionNumbers();
// If the index vector dim is not implicitly or explicitly at the end
// we need to transpose the batch dimensions to the start.
Value indices = op.getScatterIndices();
auto indicesTy = llvm::cast<ShapedType>(indices.getType());
auto indexVectorDim = dimNumbers.getIndexVectorDim();
if (indexVectorDim < indicesTy.getRank() - 1) {
llvm::SmallVector<int64_t> perm;
perm.reserve(indicesTy.getRank());
for (int i = 0, s = indicesTy.getRank(); i < s; ++i) {
if (i != indexVectorDim)
perm.push_back(i);
}
if (perm.size() < indicesTy.getRank())
perm.push_back(indexVectorDim);
llvm::SmallVector<int64_t> newShape;
for (int i = 0, s = perm.size(); i < s; ++i) {
newShape.push_back(indicesTy.getDimSize(perm[i]));
}
indices = builder.create<mlir::stablehlo::TransposeOp>(
indicesTy.clone(newShape), indices,
builder.getDenseI64ArrayAttr(perm));
indicesTy = llvm::cast<RankedTensorType>(indices.getType());
indexVectorDim = indicesTy.getRank() - 1;
}
// Compute the permutation require to transpose the batch dimensions to
// the beginning.
auto updates = op.getUpdates();
auto updates0 = updates.front();
auto updates0Ty = llvm::cast<ShapedType>(updates0.getType());
auto updatedWindowDims = dimNumbers.getUpdateWindowDims();
// Determine which dimensions are batch dimensions.
llvm::SmallVector<bool> isBatch(updates0Ty.getRank(), true);
for (int i = 0, s = updatedWindowDims.size(); i < s; ++i)
isBatch[updatedWindowDims[i]] = false;
// Permute batch dimensions to the start of the update tensor.
llvm::SmallVector<int64_t> updatePerm;
updatePerm.reserve(updates0Ty.getRank());
for (int i = 0, s = isBatch.size(); i < s; ++i)
if (isBatch[i])
updatePerm.push_back(i);
updatePerm.append(updatedWindowDims.begin(), updatedWindowDims.end());
llvm::SmallVector<int64_t> newUpdatedWindowDims;
int64_t batchCount = updates0Ty.getRank() - updatedWindowDims.size();
for (int i = batchCount, s = updates0Ty.getRank(); i < s; i++)
newUpdatedWindowDims.push_back(i);
bool indicesChanged = indices != op.getScatterIndices();
bool updatesChanged =
llvm::any_of(llvm::enumerate(updatePerm),
[](auto it) { return it.index() != it.value(); });
llvm::SmallVector<Value> newUpdates(updates.begin(), updates.end());
if (updatesChanged) {
for (Value &update : newUpdates) {
auto updateTy = llvm::cast<ShapedType>(update.getType());
llvm::SmallVector<int64_t> newShape;
newShape.reserve(updateTy.getRank());
for (int i = 0, s = updatePerm.size(); i < s; i++)
newShape.push_back(updateTy.getDimSize(updatePerm[i]));
update = builder.create<mlir::stablehlo::TransposeOp>(
updateTy.clone(newShape), update,
builder.getDenseI64ArrayAttr(updatePerm));
}
}
if (!indicesChanged && !updatesChanged)
return rewriter.notifyMatchFailure(
op, "batch dimensions are already leading");
auto newDimNumbers = mlir::stablehlo::ScatterDimensionNumbersAttr::get(
op.getContext(), newUpdatedWindowDims,
dimNumbers.getInsertedWindowDims(), dimNumbers.getInputBatchingDims(),
dimNumbers.getScatterIndicesBatchingDims(),
dimNumbers.getScatterDimsToOperandDims(),
/*indexVectorDim=*/indexVectorDim);
auto newScatter = rewriter.create<mlir::stablehlo::ScatterOp>(
op.getLoc(), op.getResultTypes(), op.getInputs(), indices, newUpdates,
newDimNumbers, op.getIndicesAreSorted(), op.getUniqueIndices());
Region &region = newScatter.getUpdateComputation();
rewriter.cloneRegionBefore(op.getUpdateComputation(), region, region.end());
rewriter.replaceOp(op, newScatter.getResults());
return success();
}
};
// stablehlo.scatter can materialize a unit dimension at both indexed dimensions
// or at unary dimensions in the destination matrix. linalg_ext.scatter only
// allows unit dimensions at indexed dimensions. This pattern inserts all
// unary dimensions that are not index dimensions to be compatible with
// linalg_ext.scatter.
//
// If converts an stablehlo.scatter as below:
// %result = "stablehlo.scatter"(...) ({
// indices_are_sorted = true,
// scatter_dimension_numbers = #stablehlo.scatter<
// update_window_dims = [1],
// inserted_window_dims = [0, 2],
// scatter_dims_to_operand_dims = [0],
// index_vector_dim = 1>,
// unique_indices = true} :
// (tensor<5x4x1xi32>, tensor<1x1xi32>, tensor<1x4xi32>)
//
// To:
// %result = "stablehlo.scatter"(...) ({
// indices_are_sorted = true,
// scatter_dimension_numbers = #stablehlo.scatter<
// update_window_dims = [1, 2],
// inserted_window_dims = [0],
// scatter_dims_to_operand_dims = [0],
// index_vector_dim = 1>,
// unique_indices = true} :
// (tensor<5x4x1xi32>, tensor<1x1xi32>, tensor<1x4x1xi32>)
// return %0 : tensor<5x4x1xi32>
struct ScatterMaterializeInsertedDim final
: OpRewritePattern<mlir::stablehlo::ScatterOp> {
using OpRewritePattern::OpRewritePattern;
LogicalResult matchAndRewrite(mlir::stablehlo::ScatterOp op,
PatternRewriter &rewriter) const override {
auto indices = op.getScatterIndices();
auto operand = op.getInputs().front();
auto indicesTy = llvm::cast<ShapedType>(indices.getType());
auto operandTy = llvm::cast<ShapedType>(operand.getType());
if (!operandTy.hasRank() || !indicesTy.hasRank()) {
return rewriter.notifyMatchFailure(op, "operand/indices have no rank");
}
auto dimNumbers = op.getScatterDimensionNumbers();
auto updateDims = dimNumbers.getUpdateWindowDims();
if (indicesTy.getRank() != 2 || dimNumbers.getIndexVectorDim() != 1) {
return rewriter.notifyMatchFailure(
op, "indices is not of shape [batch, indices]");
}
if (!updateDims.empty() && updateDims.front() == 0) {
return rewriter.notifyMatchFailure(
op, "updates is not of shape [batch, ...]");
}
auto scatterDimsToOperandDims = dimNumbers.getScatterDimsToOperandDims();
llvm::SmallVector<bool> isIndexDim(operandTy.getRank(), false);
for (auto val : scatterDimsToOperandDims) {
isIndexDim[val] = true;
}
int64_t firstNonIndex = 0;
for (int64_t s = scatterDimsToOperandDims.size(); firstNonIndex < s;
++firstNonIndex) {
if (!isIndexDim[firstNonIndex])
break;
}
llvm::SmallVector<bool> isInsertDims(operandTy.getRank(), false);
for (auto val : dimNumbers.getInsertedWindowDims()) {
isInsertDims[val] = true;
}
int64_t frontInsertedDims = 0;
for (; frontInsertedDims < firstNonIndex; ++frontInsertedDims) {
if (!isInsertDims[frontInsertedDims]) {
break;
}
}
llvm::ArrayRef<bool> toInsertDims =
llvm::ArrayRef<bool>(isInsertDims).drop_front(frontInsertedDims);
if (!llvm::any_of(toInsertDims, [](auto d) { return d; })) {
return rewriter.notifyMatchFailure(op, "no dimensions to insert");
}
// Create a reassociation map that starts with the batch dims.
SmallVector<ReassociationExprs> reassociationMap;
reassociationMap.push_back({rewriter.getAffineDimExpr(0)});
for (auto it : llvm::enumerate(llvm::ArrayRef<bool>(toInsertDims))) {
if (!it.value())
reassociationMap.push_back({});
reassociationMap.back().push_back(
rewriter.getAffineDimExpr(it.index() + 1));
}
llvm::SmallVector<Value> expandedUpdates;
for (auto update : op.getUpdates()) {
auto updatesTy = llvm::cast<ShapedType>(update.getType());
llvm::SmallVector<int64_t> newShape;
for (int i = 0, s = reassociationMap.size(); i < s; ++i) {
newShape.push_back(updatesTy.getDimSize(i));
for (int j = 1, s = reassociationMap[i].size(); j < s; ++j) {
newShape.push_back(1);
}
}
Value expandUpdate = rewriter.create<tensor::ExpandShapeOp>(
op.getLoc(),
RankedTensorType::get(newShape, updatesTy.getElementType()), update,
reassociationMap);
expandedUpdates.push_back(expandUpdate);
}
llvm::SmallVector<int64_t> newUpdatedWindowDims(toInsertDims.size());
llvm::SmallVector<int64_t> newInsertedWindowDims(frontInsertedDims);
std::iota(newUpdatedWindowDims.begin(), newUpdatedWindowDims.end(), 1);
std::iota(newInsertedWindowDims.begin(), newInsertedWindowDims.end(), 0);
auto newDimNumbers = mlir::stablehlo::ScatterDimensionNumbersAttr::get(
op.getContext(), newUpdatedWindowDims, newInsertedWindowDims,
dimNumbers.getInputBatchingDims(),
dimNumbers.getScatterIndicesBatchingDims(),
dimNumbers.getScatterDimsToOperandDims(),
/*indexVectorDim=*/1);
auto newScatter = rewriter.create<mlir::stablehlo::ScatterOp>(
op.getLoc(), op.getResultTypes(), op.getInputs(),
op.getScatterIndices(), expandedUpdates, newDimNumbers,
op.getIndicesAreSorted(), op.getUniqueIndices());
Region &region = newScatter.getUpdateComputation();
rewriter.cloneRegionBefore(op.getUpdateComputation(), region, region.end());
rewriter.replaceOp(op, newScatter.getResults());
return success();
}
};
// Traverse upward past common operations to see if the value came from a
// boolean tensor.
bool isFromBool(Value val) {
while (true) {
Operation *op = val.getDefiningOp();
if (!op)
return false;
if (auto convertOp = dyn_cast<mlir::stablehlo::ConvertOp>(op)) {
auto inTy = llvm::cast<ShapedType>(convertOp.getOperand().getType());
if (inTy.getElementType().isInteger(1)) {
return true;
}
val = convertOp.getOperand();
continue;
}
if (isa<mlir::stablehlo::DynamicBroadcastInDimOp>(op) ||
isa<mlir::stablehlo::BroadcastInDimOp>(op) ||
isa<mlir::stablehlo::BroadcastOp>(op)) {
val = op->getOperand(0);
continue;
}
return false;
}
}
// Mul of non-finite values (e.g. NaN, inf) and 0.0 produce 0.0 in StableHLO.
// For linalg we need to convert these to select operations.
struct MulCastOfBool final : OpRewritePattern<mlir::stablehlo::MulOp> {
using OpRewritePattern::OpRewritePattern;
LogicalResult matchAndRewrite(mlir::stablehlo::MulOp op,
PatternRewriter &rewriter) const override {
auto resultTy = cast<ShapedType>(op.getType());
if (!isa<FloatType>(resultTy.getElementType()))
return failure();
Value lhs = op.getLhs();
Value rhs = op.getRhs();
bool lhsIsBool = isFromBool(lhs);
bool rhsIsBool = isFromBool(rhs);
if (lhsIsBool == rhsIsBool)
return failure();
if (rhsIsBool)
std::swap(lhs, rhs);
Type eType = resultTy.getElementType();
auto lhsTy = cast<ShapedType>(lhs.getType());
Value lhsBool = rewriter.create<mlir::stablehlo::ConvertOp>(
op.getLoc(), lhsTy.clone(rewriter.getIntegerType(1)), lhs);
Value zero = rewriter.create<mlir::stablehlo::ConstantOp>(
op.getLoc(), DenseElementsAttr::get(RankedTensorType::get({}, eType),
rewriter.getZeroAttr(eType)));
auto lhsShape = rewriter.create<shape::ShapeOfOp>(
op.getLoc(),
RankedTensorType::get({lhsTy.getRank()}, rewriter.getIndexType()), lhs);
int64_t resultRank = resultTy.getRank();
auto broadcast = [&](Value value) -> Value {
auto valueTy = cast<ShapedType>(value.getType());
auto newTy =
RankedTensorType::get(resultTy.getShape(), valueTy.getElementType());
if (valueTy == newTy)
return value;
auto dimensions = llvm::to_vector(
llvm::seq<int64_t>(resultRank - valueTy.getRank(), resultRank));
return rewriter.create<mlir::stablehlo::DynamicBroadcastInDimOp>(
op.getLoc(), newTy, value, lhsShape,
rewriter.getDenseI64ArrayAttr(dimensions));
};
zero = broadcast(zero);
rewriter.replaceOpWithNewOp<mlir::stablehlo::SelectOp>(op, resultTy,
lhsBool, rhs, zero);
return success();
}
};
// Generates Gaussian noise with uniform random generator based on Box-Muller
// transform.
struct ExpandRngNormal final : OpRewritePattern<mlir::stablehlo::RngOp> {
using OpRewritePattern::OpRewritePattern;
LogicalResult matchAndRewrite(mlir::stablehlo::RngOp op,
PatternRewriter &rewriter) const override {
if (op.getRngDistribution() != mlir::stablehlo::RngDistribution::NORMAL)
return failure();
auto resTy = dyn_cast<RankedTensorType>(op.getType());
// We can support static shapes, but it's easier to implement Box-Muller
// transform if we know the number of elements.
if (!resTy || !resTy.hasStaticShape())
return failure();
// The algorithm requires even numbers and will generate pairs.
auto numElems = resTy.getNumElements();
if (numElems & 1)
numElems++;
auto halfNumElems = numElems / 2;
ImplicitLocOpBuilder b(op.getLoc(), rewriter);
// Explicitly set the seed to 0, so we have stateless generator. This is not
// a hard limit. Random generator is still a new topic, and we start with
// stateless random generator.
std::mt19937 rng{0};
std::uniform_real_distribution<> runif(0.0, 1.0);
SmallVector<float> sqrtValues(halfNumElems), cosValues(halfNumElems),
sinValues(halfNumElems);
for (auto i : llvm::seq<unsigned>(0, numElems / 2)) {
static constexpr float kEpsilon = std::numeric_limits<float>::epsilon();
static constexpr float kTwoPi = static_cast<float>(2.0 * M_PI);
float u1, u2;
do {
u1 = runif(rng);
u2 = runif(rng);
} while (u1 <= kEpsilon);
sqrtValues[i] = -2.0 * log(u1);
cosValues[i] = cos(kTwoPi * u2);
sinValues[i] = sin(kTwoPi * u2);
}
// mag = sigma * sqrt(-2.0 * log(u1));
Value mag = getF32Const(b, /*shapes=*/{halfNumElems}, sqrtValues);
Value sigma = b.create<mlir::stablehlo::BroadcastOp>(
mag.getType(), op.getB(), make1DElementsAttr(b, halfNumElems));
mag = b.create<mlir::stablehlo::MulOp>(
sigma, b.create<mlir::stablehlo::SqrtOp>(mag));
// z0 = mag * cos(two_pi * u2) + mu;
// z1 = mag * sin(two_pi * u2) + mu;
Value mu = b.create<mlir::stablehlo::BroadcastOp>(
mag.getType(), op.getA(), make1DElementsAttr(b, halfNumElems));
Value z0 = getF32Const(b, /*shapes=*/{halfNumElems}, cosValues);
z0 = b.create<mlir::stablehlo::MulOp>(mag, z0);
z0 = b.create<mlir::stablehlo::AddOp>(z0, mu);
Value z1 = getF32Const(b, /*shapes=*/{halfNumElems}, sinValues);
z1 = b.create<mlir::stablehlo::MulOp>(mag, z1);
z1 = b.create<mlir::stablehlo::AddOp>(z1, mu);
Value res = b.create<mlir::stablehlo::ConcatenateOp>(
ValueRange{z0, z1}, b.getI64IntegerAttr(0));
if (numElems != resTy.getNumElements()) {
OpFoldResult zero = b.getIndexAttr(0);
OpFoldResult one = b.getIndexAttr(1);
OpFoldResult size = b.getIndexAttr(resTy.getNumElements());
res = b.create<tensor::ExtractSliceOp>(res, zero, size, one);
}
if (resTy.getRank() != 1) {
res = b.create<mlir::stablehlo::ReshapeOp>(resTy, res);
}
rewriter.replaceOp(op, res);
return success();
}
};
// clang-format off
//
// Reorder BroadcastInDimOp and N-ary elementwise op.
//
// Rewrites the following pattern (take binary elementwise op as example)
//
// %bcastx = "stablehlo.broadcast_in_dim"(%x) {broadcast_dimensions = %[[BCAST_DIMS]]} : (%[[SHAPE_BEFORE_BCAST]]) -> %[[SHAPE_AFTER_BCAST]]
// %bcasty = "stablehlo.broadcast_in_dim"(%y) {broadcast_dimensions = %[[BCAST_DIMS]]} : (%[[SHAPE_BEFORE_BCAST]]) -> %[[SHAPE_AFTER_BCAST]]
// %result = "BinaryElementwiseOpT"(%bcastx, %bcasty) : (%[[SHAPE_AFTER_BCAST]], %[[SHAPE_AFTER_BCAST]]) -> %[[SHAPE_AFTER_BCAST]]
//
// into
//
// %z = "BinaryElementwiseOpT"(%x, %y) : (%[[SHAPE_BEFORE_BCAST]], %[[SHAPE_BEFORE_BCAST]]) -> %[[SHAPE_BEFORE_BCAST]]
// %result = "stablehlo.broadcast_in_dim"(%z) {broadcast_dimensions = %[[BCAST_DIMS]]} : (%[[SHAPE_BEFORE_BCAST]]) -> %[[SHAPE_AFTER_BCAST]]
//
// clang-format on
template <typename ElementwiseOpT>
struct ReorderBroadcastInDimOpAndElementwiseOp final
: OpRewritePattern<ElementwiseOpT> {
using OpRewritePattern<ElementwiseOpT>::OpRewritePattern;
LogicalResult matchAndRewrite(ElementwiseOpT op,
PatternRewriter &rewriter) const override {
Operation *operation = op.getOperation();
assert(operation->getNumOperands() >= 1 && operation->getNumResults() == 1);
// Verify if all operands are from BroadcastInDimOp and its
// broadcast_dimensions is the same.
llvm::SmallVector<mlir::stablehlo::BroadcastInDimOp, 2> bcastOps;
for (auto operand : operation->getOperands()) {
if (auto bcastOp =
operand.getDefiningOp<mlir::stablehlo::BroadcastInDimOp>()) {
bcastOps.push_back(bcastOp);
} else {
return failure();
}
}
if (llvm::any_of(bcastOps,
[&bcastOps](mlir::stablehlo::BroadcastInDimOp bcastOp) {
return bcastOp.getBroadcastDimensions() !=
bcastOps[0].getBroadcastDimensions();
})) {
return failure();
}
// Verify if all operands of BroadcastInDimOp are of same type and have
// static shape.
auto bcastOperandType =
dyn_cast<ShapedType>(bcastOps[0].getOperand().getType());
llvm::SmallVector<Value, 2> bcastOperands;
for (auto bcastOp : bcastOps) {
auto bcastOperand = bcastOp.getOperand();
auto type = dyn_cast<ShapedType>(bcastOperand.getType());
if (!type || !type.hasStaticShape() || type != bcastOperandType) {
return failure();
}
bcastOperands.push_back(bcastOperand);
}
// Some elementwise ops, mlir::stablehlo::RealOp for example, do not have
// SameOperandsAndResultType trait, so resultType might be different
// from bcastOperandType.
auto elementType = getElementTypeOrSelf(op.getResult());
auto resultShape = bcastOperandType.getShape();
auto resultType = RankedTensorType::get(resultShape, elementType);
Value result =
rewriter.create<ElementwiseOpT>(op.getLoc(), resultType, bcastOperands);
rewriter.replaceOpWithNewOp<mlir::stablehlo::BroadcastInDimOp>(
op, op.getType(), result, bcastOps[0].getBroadcastDimensionsAttr());
// NOTE: bcastOps may contain duplicates.
SetVector<Operation *> deadOps;
for (auto bcastOp : bcastOps) {
if (bcastOp.getOperation()->use_empty())
deadOps.insert(bcastOp);
}
for (auto *deadOp : deadOps)
rewriter.eraseOp(deadOp);
return success();
}
};
// Identifies cases where a dense operation has inputs that come from widening
// operations. For instance, a dot product widening from FP16 to FP32 is better
// to have the casting operation fused into the dot operation. This decreases
// the loading required during a dense computation.
template <class Op>
struct FuseWidenOperands final : OpRewritePattern<Op> {
using OpRewritePattern<Op>::OpRewritePattern;
LogicalResult matchAndRewrite(Op op,
PatternRewriter &rewriter) const override {
llvm::SmallVector<Value> operands;
for (Value operand : op->getOperands()) {
auto convertOp =
dyn_cast_or_null<mlir::stablehlo::ConvertOp>(operand.getDefiningOp());
if (convertOp) {
auto inputType = getElementTypeOrSelf(convertOp.getOperand().getType());
auto castedType = getElementTypeOrSelf(convertOp.getResult().getType());
if (!isa<FloatType, IntegerType>(inputType) ||
!isa<FloatType, IntegerType>(castedType)) {
return rewriter.notifyMatchFailure(
op, "non-integer or floating point type");
;
}
if (inputType.getIntOrFloatBitWidth() <
castedType.getIntOrFloatBitWidth()) {
operands.push_back(convertOp.getOperand());
continue;
}
}
operands.push_back(operand);
}
if (llvm::all_of(
llvm::zip_equal(operands, op->getOperands()),
[](auto pair) { return std::get<0>(pair) == std::get<1>(pair); }))
return failure();
rewriter.replaceOpWithNewOp<Op>(op, op->getResultTypes(), operands,
op->getAttrs());
return success();
}
};
struct DotToMul final : OpRewritePattern<mlir::stablehlo::DotOp> {
using OpRewritePattern::OpRewritePattern;
LogicalResult matchAndRewrite(mlir::stablehlo::DotOp op,
PatternRewriter &rewriter) const override {
Value lhs = op.getLhs();
Value rhs = op.getRhs();
auto lhsTy = dyn_cast<RankedTensorType>(lhs.getType());
auto rhsTy = dyn_cast<RankedTensorType>(rhs.getType());
auto resultTy = dyn_cast<RankedTensorType>(op.getType());
if (!lhsTy || !rhsTy) {
return rewriter.notifyMatchFailure(op, "lhs and rhs must be ranked");
}
if (lhsTy.getRank() != 2 || rhsTy.getRank() != 2) {
return rewriter.notifyMatchFailure(op, "lhs and rhs must be rank-2");
}
if (lhsTy.getDimSize(1) != 1)
return failure();
// Dynamically compute the shape of the result of the DotOp by querying
// the 0-th dimensions, of the left, and the 1st dimension of the right.
// Concatenating them togething to make the final shape.
Value batchSize = rewriter.create<mlir::stablehlo::GetDimensionSizeOp>(
op.getLoc(), lhs, rewriter.getI64IntegerAttr(0));
Value batchSize1 = rewriter.create<mlir::stablehlo::ReshapeOp>(
op.getLoc(), RankedTensorType::get({1}, rewriter.getI32Type()),
batchSize);
Value featureSize = rewriter.create<mlir::stablehlo::GetDimensionSizeOp>(
op.getLoc(), rhs, rewriter.getI64IntegerAttr(1));
Value featureSize1 = rewriter.create<mlir::stablehlo::ReshapeOp>(
op.getLoc(), RankedTensorType::get({1}, rewriter.getI32Type()),
featureSize);
Value outSize = rewriter.create<mlir::stablehlo::ConcatenateOp>(
op.getLoc(), RankedTensorType::get({2}, rewriter.getI32Type()),
ValueRange{batchSize1, featureSize1}, rewriter.getI64IntegerAttr(0));
lhs = rewriter.create<mlir::stablehlo::DynamicBroadcastInDimOp>(
op.getLoc(), resultTy.clone(lhsTy.getElementType()), lhs, outSize,
rewriter.getDenseI64ArrayAttr({0, 1}));
rhs = rewriter.create<mlir::stablehlo::DynamicBroadcastInDimOp>(
op.getLoc(), resultTy.clone(rhsTy.getElementType()), rhs, outSize,
rewriter.getDenseI64ArrayAttr({0, 1}));
auto computeETy = lhsTy.getElementType();
if (computeETy.getIntOrFloatBitWidth() < rhsTy.getElementTypeBitWidth())
computeETy = rhsTy.getElementType();
if (computeETy.getIntOrFloatBitWidth() < resultTy.getElementTypeBitWidth())
computeETy = resultTy.getElementType();
auto computeTy = resultTy.clone(computeETy);
rhs = rewriter.create<mlir::stablehlo::ConvertOp>(op.getLoc(), computeTy,
rhs);
lhs = rewriter.create<mlir::stablehlo::ConvertOp>(op.getLoc(), computeTy,
lhs);
auto result = rewriter.create<mlir::stablehlo::MulOp>(
op.getLoc(), resultTy.clone(computeETy), lhs, rhs);
rewriter.replaceOpWithNewOp<mlir::stablehlo::ConvertOp>(op, resultTy,
result);
return success();
}
};
// Rewrite RngBitGenerator with f32 return type to instead generate the same
// number of i32 outputs, then BitcastConvert to return f32.
struct RngBitcastFloat final
: OpRewritePattern<mlir::stablehlo::RngBitGeneratorOp> {
using OpRewritePattern::OpRewritePattern;
LogicalResult matchAndRewrite(mlir::stablehlo::RngBitGeneratorOp op,
PatternRewriter &rewriter) const override {
ImplicitLocOpBuilder builder(op.getLoc(), rewriter);
auto resultTy = dyn_cast<RankedTensorType>(op.getType(1));
auto stateTy = dyn_cast<RankedTensorType>(op.getType(0));
if (!isa<FloatType>(resultTy.getElementType())) {
return failure();
}
llvm::SmallVector<Type> castedTypes;
castedTypes.push_back(stateTy);
castedTypes.push_back(resultTy.clone(rewriter.getI32Type()));
TypeRange castedTypeRange = TypeRange{castedTypes};
auto resultOp = rewriter.create<mlir::stablehlo::RngBitGeneratorOp>(
op.getLoc(), castedTypeRange, op.getRngAlgorithm(), op.getOperand());
auto casted = rewriter.create<mlir::stablehlo::BitcastConvertOp>(
resultOp.getLoc(), resultTy, resultOp.getResult(1));
llvm::SmallVector<Value> results;
results.push_back(resultOp.getResult(0));
results.push_back(casted);
rewriter.replaceOp(op, results);
return success();
}
};
struct ZeroConcat final : OpRewritePattern<mlir::stablehlo::ConcatenateOp> {
using OpRewritePattern::OpRewritePattern;
LogicalResult matchAndRewrite(mlir::stablehlo::ConcatenateOp op,
PatternRewriter &rewriter) const override {
auto type = dyn_cast<RankedTensorType>(op.getType());
if (!type || !type.hasStaticShape())
return failure();
uint64_t axis = op.getDimension();
OperandRange origInputs = op.getInputs();
SmallVector<Value> nonzeroInputs;
for (auto input : origInputs) {
auto type = dyn_cast<RankedTensorType>(input.getType());
ArrayRef<int64_t> shape = type.getShape();
if (axis > shape.size())
return failure();
if (shape[axis] != 0)
nonzeroInputs.push_back(input);
}
if (nonzeroInputs.size() == origInputs.size())
return failure();
rewriter.replaceOpWithNewOp<mlir::stablehlo::ConcatenateOp>(
op, nonzeroInputs, /*dimension=*/axis);
return success();
}
};
// Similar to DotIsMul, this finds the case where a dot general
// can be represented using a mul operation. This includes possibly making
// an implicit cast explicit prior the mul.
struct DotGeneralIsMul final : OpRewritePattern<mlir::stablehlo::DotGeneralOp> {
using OpRewritePattern::OpRewritePattern;
LogicalResult matchAndRewrite(mlir::stablehlo::DotGeneralOp op,
PatternRewriter &rewriter) const override {
auto lhs = cast<Value>(op.getLhs());
auto rhs = cast<Value>(op.getRhs());
auto lhsTy = dyn_cast<RankedTensorType>(lhs.getType());
auto rhsTy = dyn_cast<RankedTensorType>(rhs.getType());
auto resultTy = dyn_cast<RankedTensorType>(op.getType());
ImplicitLocOpBuilder builder(op.getLoc(), rewriter);
if (!lhsTy || !rhsTy || !resultTy)
return failure();
auto dNums = op.getDotDimensionNumbers();
auto batchDimsL = dNums.getLhsBatchingDimensions();
auto batchDimsR = dNums.getRhsBatchingDimensions();
auto contractDimsL = dNums.getLhsContractingDimensions();
auto contractDimsR = dNums.getRhsContractingDimensions();
llvm::SmallVector<bool> isLhsParallelDim(lhsTy.getRank(), true);
llvm::SmallVector<bool> isRhsParallelDim(rhsTy.getRank(), true);
for (auto dim : batchDimsL)
isLhsParallelDim[dim] = false;
for (auto dim : batchDimsR)
isRhsParallelDim[dim] = false;
for (auto dim : contractDimsL)
isLhsParallelDim[dim] = false;
for (auto dim : contractDimsR)
isRhsParallelDim[dim] = false;
for (auto dim : contractDimsL) {
if (lhsTy.getDimSize(dim) != 1) {
return rewriter.notifyMatchFailure(op, "Non unit contract dimensions");
}
}
// Generate the permutation matrix to order BatchDims, ParallelDims,
// ContractDims.
llvm::SmallVector<int64_t> permLhs;
llvm::SmallVector<int64_t> permRhs;
permLhs.append(batchDimsL.begin(), batchDimsL.end());
permRhs.append(batchDimsR.begin(), batchDimsR.end());
for (auto [idx, value] : llvm::enumerate(isLhsParallelDim)) {
if (value)
permLhs.push_back(idx);
}
for (auto [idx, value] : llvm::enumerate(isRhsParallelDim)) {
if (value)
permRhs.push_back(idx);
}
llvm::append_range(permLhs, contractDimsL);
llvm::append_range(permRhs, contractDimsR);
// Determine the transpose shape based on the generate permutations.
llvm::SmallVector<int64_t> lhsTransposeShape;
llvm::SmallVector<int64_t> rhsTransposeShape;
for (auto dim : permLhs)
lhsTransposeShape.push_back(lhsTy.getDimSize(dim));
for (auto dim : permRhs)
rhsTransposeShape.push_back(rhsTy.getDimSize(dim));
// Transpose the left hand side and the right hand side.
lhs = builder.create<mlir::stablehlo::TransposeOp>(
RankedTensorType::get(lhsTransposeShape, lhsTy.getElementType()), lhs,
builder.getDenseI64ArrayAttr(permLhs));
lhsTy = llvm::cast<RankedTensorType>(lhs.getType());
rhs = builder.create<mlir::stablehlo::TransposeOp>(
RankedTensorType::get(rhsTransposeShape, rhsTy.getElementType()), rhs,
builder.getDenseI64ArrayAttr(permRhs));
rhsTy = llvm::cast<RankedTensorType>(rhs.getType());
auto dimI32Ty = RankedTensorType::get({1}, builder.getI32Type());
// Drop all of the non-concat dimensions from the lhs.
llvm::SmallVector<Value> lhsReshapeDims;
for (int i = 0, s = lhsTy.getRank() - contractDimsL.size(); i < s; i++) {
Value dim = builder.create<mlir::stablehlo::GetDimensionSizeOp>(lhs, i);
lhsReshapeDims.push_back(
builder.create<mlir::stablehlo::ReshapeOp>(dimI32Ty, dim));
}
Value lhsDynShape = builder.create<mlir::stablehlo::ConcatenateOp>(
RankedTensorType::get({static_cast<int64_t>(lhsReshapeDims.size())},
builder.getI32Type()),
lhsReshapeDims, 0);
lhsTy =
RankedTensorType::get(lhsTy.getShape().drop_back(contractDimsL.size()),
lhsTy.getElementType());
lhs = builder.create<mlir::stablehlo::DynamicReshapeOp>(lhsTy, lhs,
lhsDynShape);
// Drop all of the non concat dimensions from the rhs.
llvm::SmallVector<Value> rhsReshapeDims;
for (int i = 0, s = rhsTy.getRank() - contractDimsR.size(); i < s; i++) {
Value dim = builder.create<mlir::stablehlo::GetDimensionSizeOp>(rhs, i);
rhsReshapeDims.push_back(
builder.create<mlir::stablehlo::ReshapeOp>(dimI32Ty, dim));
}
Value rhsDynShape = builder.create<mlir::stablehlo::ConcatenateOp>(
RankedTensorType::get({static_cast<int64_t>(rhsReshapeDims.size())},
builder.getI32Type()),
rhsReshapeDims, 0);
rhsTy =
RankedTensorType::get(rhsTy.getShape().drop_back(contractDimsR.size()),
rhsTy.getElementType());
rhs = builder.create<mlir::stablehlo::DynamicReshapeOp>(rhsTy, rhs,
rhsDynShape);
// Compute the size of the output shape with dynamic shape support using the
// lhs and rhs dimensions.
llvm::SmallVector<Value> outputDims;
outputDims.append(lhsReshapeDims);
outputDims.append(rhsReshapeDims.begin() + batchDimsR.size(),
rhsReshapeDims.end());
Value outputShape = builder.create<mlir::stablehlo::ConcatenateOp>(
RankedTensorType::get({resultTy.getRank()}, builder.getI32Type()),
outputDims, 0);
// Broadcast the left hand side to match the expect output shape.
llvm::SmallVector<int64_t> lhsDimMapping(lhsTy.getRank());
std::iota(lhsDimMapping.begin(), lhsDimMapping.end(), 0);
auto lhsBroadcastTy =
RankedTensorType::get(resultTy.getShape(), lhsTy.getElementType());
lhs = builder.createOrFold<mlir::stablehlo::DynamicBroadcastInDimOp>(
lhsBroadcastTy, lhs, outputShape,
rewriter.getDenseI64ArrayAttr(lhsDimMapping));
// Broadcast the right hand side to match the expected output shape.
llvm::SmallVector<int64_t> rhsDimMapping(rhsTy.getRank());
std::iota(rhsDimMapping.begin(), rhsDimMapping.begin() + batchDimsR.size(),
0);
std::iota(rhsDimMapping.begin() + batchDimsR.size(), rhsDimMapping.end(),
lhsTy.getRank());
auto rhsBroadcastTy =
RankedTensorType::get(resultTy.getShape(), rhsTy.getElementType());
rhs = builder.createOrFold<mlir::stablehlo::DynamicBroadcastInDimOp>(
rhsBroadcastTy, rhs, outputShape,
rewriter.getDenseI64ArrayAttr(rhsDimMapping));
lhs = builder.createOrFold<mlir::stablehlo::ConvertOp>(resultTy, lhs);
rhs = builder.createOrFold<mlir::stablehlo::ConvertOp>(resultTy, rhs);
rewriter.replaceOpWithNewOp<mlir::stablehlo::MulOp>(op, resultTy, lhs, rhs);
return success();
}
};
// TODO(suderman): Move this to solve code.
struct CustomCallIsTopK final
: OpRewritePattern<mlir::stablehlo::CustomCallOp> {
using OpRewritePattern::OpRewritePattern;
LogicalResult matchAndRewrite(mlir::stablehlo::CustomCallOp op,
PatternRewriter &rewriter) const override {
if (op.getCallTargetName() != "TopK") {
return rewriter.notifyMatchFailure(op, "not a TopK custom call");
}
if (op.getNumOperands() != 1 ||
!(op.getNumResults() == 1 || op.getNumResults() == 2)) {
return rewriter.notifyMatchFailure(
op, "incorrect number of operands / results");
}
ArrayAttr computations = op.getCalledComputations();
if (computations.size() != 1) {
return rewriter.notifyMatchFailure(op,
"incorrect number of computations");
}
SymbolRefAttr computation = dyn_cast<SymbolRefAttr>(computations[0]);
if (!computation) {
return rewriter.notifyMatchFailure(op, "not a ref attr");
}
auto operand = op.getOperand(0);
auto operandTy = cast<ShapedType>(operand.getType());
if (!operandTy.hasRank() || operandTy.getRank() != 2) {
return rewriter.notifyMatchFailure(op, "rank-2 input not found");
}
ShapedType topVTy;
ShapedType topITy;
if (op.getNumResults() == 1) {
if (auto tupleTy = dyn_cast<TupleType>(op.getType(0))) {
if (tupleTy.size() != 2) {
return rewriter.notifyMatchFailure(
op, "tuple return does not tuple two values");
}
topVTy = dyn_cast<ShapedType>(tupleTy.getType(0));
topITy = dyn_cast<ShapedType>(tupleTy.getType(1));
}
}
if (op.getNumResults() == 2) {
topVTy = dyn_cast<ShapedType>(op.getType(0));
topITy = dyn_cast<ShapedType>(op.getType(1));
}
if (!topVTy || !topITy) {
return rewriter.notifyMatchFailure(op, "unknown return type behavior");
}
int64_t k = topVTy.getDimSize(1);
if (ShapedType::isDynamic(k)) {
return rewriter.notifyMatchFailure(op, "dynamic top-k k value");
}
auto moduleOp = op->getParentOfType<ModuleOp>();
auto funcOp = dyn_cast<func::FuncOp>(moduleOp.lookupSymbol(computation));
Block &block = funcOp.getRegion().front();
auto stablehloCompareOp =
dyn_cast<mlir::stablehlo::CompareOp>(block.front());
if (!stablehloCompareOp) {
return rewriter.notifyMatchFailure(op, "not stablehlo compare op");
}
auto returnOp = dyn_cast<func::ReturnOp>(block.getTerminator());
if (!returnOp) {
return rewriter.notifyMatchFailure(op, "could not find ReturnOp");
}
if (returnOp.getNumOperands() != 1 ||
returnOp.getOperand(0) != stablehloCompareOp.getResult()) {
return rewriter.notifyMatchFailure(op, "ReturnOp operand not compare op");
}
auto direction = stablehloCompareOp.getComparisonDirection();
bool getTop = direction == mlir::stablehlo::ComparisonDirection::GT ||
direction == mlir::stablehlo::ComparisonDirection::GE;
if (!getTop) {
return rewriter.notifyMatchFailure(op,
"Unsupported comparison direction");
}
auto newTopK = rewriter.create<chlo::TopKOp>(
op.getLoc(), TypeRange{topVTy, topITy}, operand, k);
if (op.getNumResults() == 2) {
rewriter.replaceOp(op, newTopK.getResults());
return success();
}
if (auto tupleTy = dyn_cast<TupleType>(op.getType(0))) {
rewriter.replaceOpWithNewOp<mlir::stablehlo::TupleOp>(
op, op.getType(0), newTopK.getResults());
return success();
}
return failure();
}
};
// Recursive helper function that identifies an Iota followed by a set of
// broadcasts where the last dimension of the iota is preserved throughout.
bool isIotaOrIotaBroadcast(PatternRewriter &rewriter, Value input) {
if (auto iotaOp =
dyn_cast_or_null<mlir::stablehlo::IotaOp>(input.getDefiningOp())) {
int64_t iotaDim = iotaOp.getIotaDimension();
auto iotaLastDim = cast<ShapedType>(iotaOp.getType()).getRank() - 1;
if (iotaDim == iotaLastDim) {
return true;
}
(void)rewriter.notifyMatchFailure(iotaOp, "iota must be on last dimension");
return false;
}
if (auto broadcastOp = dyn_cast_or_null<mlir::stablehlo::BroadcastInDimOp>(
input.getDefiningOp())) {
auto broadcastLastDim =
cast<ShapedType>(broadcastOp.getType()).getRank() - 1;
if (broadcastOp.getBroadcastDimensions().back() != broadcastLastDim) {
(void)rewriter.notifyMatchFailure(
broadcastOp, "last dimension must be maintained in broadcast");
return false;
}
return isIotaOrIotaBroadcast(rewriter, broadcastOp.getOperand());
}
return false;
}
struct IotaSortSliceIsTopK final : OpRewritePattern<mlir::stablehlo::SortOp> {
using OpRewritePattern::OpRewritePattern;
LogicalResult matchAndRewrite(mlir::stablehlo::SortOp op,
PatternRewriter &rewriter) const override {
auto opOperands = op.getOperands();
auto opResults = op.getResults();
Value topKInput;
if (opOperands.size() != 2 || opResults.size() != 2) {
return rewriter.notifyMatchFailure(
op, "slice that maps to TopK must have exactly two inputs/outputs");
}
Value inputIota;
// Check that one of the inputs is iota, assume that the other one is the
// input.
for (Value operand : opOperands) {
if (isIotaOrIotaBroadcast(rewriter, operand)) {
inputIota = operand;
} else {
topKInput = operand;
}
}
if (!inputIota) {
return rewriter.notifyMatchFailure(op, "sort isn't called from Iota");
}
Block &block = op.getRegion().front();
auto stablehloCompareOp =
dyn_cast<mlir::stablehlo::CompareOp>(block.front());
if (!stablehloCompareOp) {
return rewriter.notifyMatchFailure(op, "not stablehlo compare op");
}
auto direction = stablehloCompareOp.getComparisonDirection();
bool getTop = direction == mlir::stablehlo::ComparisonDirection::GT ||
direction == mlir::stablehlo::ComparisonDirection::GE;
if (!getTop) {
return rewriter.notifyMatchFailure(op,
"unsupported comparison direction");
}
Value topV, topI;
int64_t k;
// Check that the output of the sort op gets fed into a slice.
for (auto [idx, result] : llvm::enumerate(opResults)) {
if (result.getUsers().empty())
return rewriter.notifyMatchFailure(
op, "sort isn't calling into a slice op");
auto sliceOp =
dyn_cast<mlir::stablehlo::SliceOp>(*result.getUsers().begin());
if (!sliceOp) {
return rewriter.notifyMatchFailure(
op, "sort isn't calling into a slice op");
}
for (auto stride : sliceOp.getStrides()) {
if (stride != 1) {
return rewriter.notifyMatchFailure(
op, "all slice strides must be 1 in order to match to TopK");
}
}
// Treat the first slice as inputs, the second as indices.
if (idx == 0) {
topV = sliceOp.getResult();
k = sliceOp.getLimitIndices().back();
} else {
topI = sliceOp.getResult();
}
}
auto topK = rewriter.create<chlo::TopKOp>(
op.getLoc(), TypeRange{topV.getType(), topI.getType()}, topKInput, k);
topV.replaceAllUsesWith(topK.getResults()[0]);
topI.replaceAllUsesWith(topK.getResults()[1]);
return success();
}
};
// TODO(suderman): Move this to solve code.
struct ApproxTopK final : OpRewritePattern<mlir::stablehlo::CustomCallOp> {
using OpRewritePattern::OpRewritePattern;
LogicalResult matchAndRewrite(mlir::stablehlo::CustomCallOp op,
PatternRewriter &rewriter) const override {
if (op.getCallTargetName() != "ApproxTopK")
return rewriter.notifyMatchFailure(op, "not ApproxTopK operation.");
auto computationName =
dyn_cast<SymbolRefAttr>(op.getCalledComputationsAttr()[0]);
Operation *funcOp;
for (auto parent = op->getParentOp(); parent;
parent = parent->getParentOp()) {
funcOp = SymbolTable::lookupNearestSymbolFrom<func::FuncOp>(
parent, computationName);
if (funcOp)
break;
}
if (!funcOp)
return rewriter.notifyMatchFailure(op, "computation function not found.");
int64_t k = cast<ShapedType>(op.getType(0)).getShape().back();
auto input = op.getOperand(0);
auto iota = op.getOperand(1);
if (auto iotaOp =
dyn_cast_or_null<mlir::stablehlo::IotaOp>(iota.getDefiningOp())) {
int64_t iotaDim = iotaOp.getIotaDimension();
auto iotaLastDim = cast<ShapedType>(iotaOp.getType()).getRank() - 1;
if (iotaDim != iotaLastDim) {
return rewriter.notifyMatchFailure(op, "Iota of last dim not found.");
}
}
Block &block = funcOp->getRegion(0).front();
auto stablehloCompareOp =
dyn_cast<mlir::stablehlo::CompareOp>(block.front());
if (!stablehloCompareOp) {
return rewriter.notifyMatchFailure(op, "not stablehlo compare op");
}
auto returnOp = block.getTerminator();
auto freturnOp = dyn_cast<func::ReturnOp>(returnOp);
auto sreturnOp = dyn_cast<mlir::stablehlo::ReturnOp>(returnOp);
if (!freturnOp && !sreturnOp) {
return rewriter.notifyMatchFailure(op, "could not find ReturnOp");
}
if (returnOp->getNumOperands() != 1 ||
returnOp->getOperand(0) != stablehloCompareOp.getResult()) {
return rewriter.notifyMatchFailure(op, "ReturnOp operand not compare op");
}
auto direction = stablehloCompareOp.getComparisonDirection();
bool getTop = direction == mlir::stablehlo::ComparisonDirection::GT ||
direction == mlir::stablehlo::ComparisonDirection::GE;
if (getTop) {
auto topK =
rewriter.create<chlo::TopKOp>(op.getLoc(), op.getResultTypes(), input,
rewriter.getI64IntegerAttr(k));
rewriter.replaceOp(op, topK);
return success();
}
bool getBottom = direction == mlir::stablehlo::ComparisonDirection::LT ||
direction == mlir::stablehlo::ComparisonDirection::LE;
if (getBottom) {
input = rewriter.create<mlir::stablehlo::NegOp>(op.getLoc(), input);
auto topK =
rewriter.create<chlo::TopKOp>(op.getLoc(), op.getResultTypes(), input,
rewriter.getI64IntegerAttr(k));
rewriter.replaceOp(op, topK);
return success();
}
return failure();
}
};
struct StableHLOToStableHLOPreprocessing final
: impl::StableHLOToStableHLOPreprocessingBase<
StableHLOToStableHLOPreprocessing> {
using StableHLOToStableHLOPreprocessingBase::
StableHLOToStableHLOPreprocessingBase;
void getDependentDialects(DialectRegistry &registry) const override {
registry.insert<shape::ShapeDialect, mlir::stablehlo::StablehloDialect,
chlo::ChloDialect, tensor::TensorDialect>();
}
void runOnOperation() override {
MLIRContext *context = &getContext();
ConversionTarget conversionTarget(*context);
RewritePatternSet conversionPatterns(context);
conversionTarget
.addLegalDialect<shape::ShapeDialect, chlo::ChloDialect,
mlir::stablehlo::StablehloDialect, math::MathDialect,
mlir::func::FuncDialect, mlir::arith::ArithDialect,
mlir::tensor::TensorDialect>();
if (failed(applyPartialConversion(getOperation(), conversionTarget,
std::move(conversionPatterns)))) {
return signalPassFailure();
}
RewritePatternSet patterns(context);
// General StableHLO canonicalization patterns. Run these with a high
// benefit to enable more rewrites and avoid needless expansions that could
// be more difficult to fold away. Note that we need to manually add these
// because StableHLO does not provide constant folders or canonicalization
// patterns for its ops.
populateCanonicalizationPatterns(context, &patterns, PatternBenefit{1024});
// TODO: Remove once we have a general contraction to matmul pass.
populatePreprocessingEinsumToDotGeneralPatterns(context, &patterns);
populatePreprocessingUnfuseBatchNormPatterns(context, &patterns);
populatePreprocessingComplexPatterns(context, &patterns);
populatePreprocessingGatherToTorchIndexSelectPatterns(context, &patterns);
patterns.insert<ExpandRngNormal, MulCastOfBool>(context);
// rng float conversion pattern
patterns.insert<RngBitcastFloat>(context);
// scatter canonicalization patterns
patterns.insert<ScatterInt64Indices, ScatterImplicitIndex,
ScatterImplicitBatch, ScatterMaterializeInsertedDim,
ScatterCollapseBatch, ScatterBatchFirst>(context);
// dot_general canonicalization patterns.
populatePreprocessingDotGeneralToDotPatterns(context, &patterns);
// TODO(jpienaar): This may be redundant with lower_general_dot. Remove if
// so.
patterns.insert<TransposeReshapeGenericDotGeneral>(context,
/*benefit=*/200);
patterns.insert<DotGeneralIsMul>(context, /*benefit=*/300);
// Fusion operations.
// TODO: Reconsider performing this optimization in preprocessing.
patterns.insert<FuseWidenOperands<mlir::stablehlo::DotOp>,
FuseWidenOperands<mlir::stablehlo::DotGeneralOp>,
FuseWidenOperands<mlir::stablehlo::ConvolutionOp>>(
context,
/*benefit=*/400);
// Identify known custom calls and convert them to equivalent StableHLO.
patterns.insert<CustomCallIsTopK>(context);
// Identify an iota->sort->slice pattern that maps to TopK.
patterns.insert<IotaSortSliceIsTopK, ApproxTopK>(context);
// Additional canonicalizers that simplify to computationally
// less-complex operations.
patterns.insert<DotToMul, ZeroConcat>(context);
// Unary elementwise op.
patterns.insert<
ReorderBroadcastInDimOpAndElementwiseOp<mlir::stablehlo::AbsOp>,
ReorderBroadcastInDimOpAndElementwiseOp<mlir::stablehlo::CeilOp>,
ReorderBroadcastInDimOpAndElementwiseOp<mlir::stablehlo::ConvertOp>,
ReorderBroadcastInDimOpAndElementwiseOp<mlir::stablehlo::ClzOp>,
ReorderBroadcastInDimOpAndElementwiseOp<mlir::stablehlo::CosineOp>,
ReorderBroadcastInDimOpAndElementwiseOp<mlir::stablehlo::ExpOp>,
ReorderBroadcastInDimOpAndElementwiseOp<mlir::stablehlo::Expm1Op>,
ReorderBroadcastInDimOpAndElementwiseOp<mlir::stablehlo::FloorOp>,
ReorderBroadcastInDimOpAndElementwiseOp<mlir::stablehlo::ImagOp>,
ReorderBroadcastInDimOpAndElementwiseOp<mlir::stablehlo::IsFiniteOp>,
ReorderBroadcastInDimOpAndElementwiseOp<mlir::stablehlo::LogOp>,
ReorderBroadcastInDimOpAndElementwiseOp<mlir::stablehlo::Log1pOp>,
ReorderBroadcastInDimOpAndElementwiseOp<mlir::stablehlo::LogisticOp>,
ReorderBroadcastInDimOpAndElementwiseOp<mlir::stablehlo::NotOp>,
ReorderBroadcastInDimOpAndElementwiseOp<mlir::stablehlo::NegOp>,
ReorderBroadcastInDimOpAndElementwiseOp<
mlir::stablehlo::PopulationCountOp>,
ReorderBroadcastInDimOpAndElementwiseOp<mlir::stablehlo::RealOp>,
ReorderBroadcastInDimOpAndElementwiseOp<mlir::stablehlo::RoundOp>,
ReorderBroadcastInDimOpAndElementwiseOp<mlir::stablehlo::RsqrtOp>,
ReorderBroadcastInDimOpAndElementwiseOp<mlir::stablehlo::SignOp>,
ReorderBroadcastInDimOpAndElementwiseOp<mlir::stablehlo::SineOp>,
ReorderBroadcastInDimOpAndElementwiseOp<mlir::stablehlo::SqrtOp>,
ReorderBroadcastInDimOpAndElementwiseOp<mlir::stablehlo::TanhOp>>(
context);
// Binary elementwise op.
patterns.insert<
ReorderBroadcastInDimOpAndElementwiseOp<mlir::stablehlo::AddOp>,
ReorderBroadcastInDimOpAndElementwiseOp<mlir::stablehlo::Atan2Op>,
ReorderBroadcastInDimOpAndElementwiseOp<mlir::stablehlo::ComplexOp>,
ReorderBroadcastInDimOpAndElementwiseOp<mlir::stablehlo::DivOp>,
ReorderBroadcastInDimOpAndElementwiseOp<mlir::stablehlo::MaxOp>,
ReorderBroadcastInDimOpAndElementwiseOp<mlir::stablehlo::MinOp>,
ReorderBroadcastInDimOpAndElementwiseOp<mlir::stablehlo::MulOp>,
ReorderBroadcastInDimOpAndElementwiseOp<mlir::stablehlo::PowOp>,
ReorderBroadcastInDimOpAndElementwiseOp<mlir::stablehlo::RemOp>,
ReorderBroadcastInDimOpAndElementwiseOp<mlir::stablehlo::ShiftLeftOp>,
ReorderBroadcastInDimOpAndElementwiseOp<
mlir::stablehlo::ShiftRightArithmeticOp>,
ReorderBroadcastInDimOpAndElementwiseOp<
mlir::stablehlo::ShiftRightLogicalOp>,
ReorderBroadcastInDimOpAndElementwiseOp<mlir::stablehlo::SubtractOp>,
ReorderBroadcastInDimOpAndElementwiseOp<mlir::stablehlo::AndOp>,
ReorderBroadcastInDimOpAndElementwiseOp<mlir::stablehlo::OrOp>,
ReorderBroadcastInDimOpAndElementwiseOp<mlir::stablehlo::XorOp>>(
context);
if (orderConvFeatures) {
patterns.insert<ReorderConvOpInputDimensions>(context);
patterns.insert<ReorderConvOpKernelDimensions>(context);
patterns.insert<ReorderConvOpOutputDimensions>(context);
}
if (failed(applyPatternsGreedily(getOperation(), std::move(patterns)))) {
return signalPassFailure();
}
}
};
} // namespace
} // namespace mlir::iree_compiler::stablehlo