Explicitly ordering transpose and reshape patterns when legalizing mhlo.dot_general
1. RankReducedDotGeneral Assume batchingDims at the starting position and the parallelDims is consecutive
2. after parallelDims reduce Also need to reduce rank
[Solution]:
move TransposeGenericDotGeneral to the front of RankReducedDotGeneral
reduce rank after reduce parallelDims
Fixes https://github.com/google/iree/issues/7272
diff --git a/iree/compiler/InputConversion/MHLO/MHLOToMHLOPreprocessing.cpp b/iree/compiler/InputConversion/MHLO/MHLOToMHLOPreprocessing.cpp
index ca32815..abf39f2 100644
--- a/iree/compiler/InputConversion/MHLO/MHLOToMHLOPreprocessing.cpp
+++ b/iree/compiler/InputConversion/MHLO/MHLOToMHLOPreprocessing.cpp
@@ -493,11 +493,16 @@
// inserts transposes so the dot_general always has the form:
// {batch_dims, parallel_dims, contraction_dims}.
// {batch_dims, contraction_dims, parallel_dims}
-class TransposeGenericDotGeneral : public OpRewritePattern<mhlo::DotGeneralOp> {
+// 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.
+class TransposeReshapeGenericDotGeneral
+ : public OpRewritePattern<mhlo::DotGeneralOp> {
public:
using OpRewritePattern<mhlo::DotGeneralOp>::OpRewritePattern;
- Value TransposeIfNonConsecutive(OpBuilder b, Location loc, Value src,
+ Value TransposeIfNonConsecutive(OpBuilder &b, Location loc, Value src,
ArrayRef<int64_t> targetOrder) const {
if (isConsecutive(targetOrder)) return src;
auto type = src.getType().cast<RankedTensorType>();
@@ -510,6 +515,23 @@
b.getI64TensorAttr(targetOrder));
}
+ Value ReshapeIfMorethan3D(OpBuilder &b, Location loc, Value src,
+ size_t dimsBorder0, size_t dimsBorder1) const {
+ auto type = src.getType().cast<RankedTensorType>();
+ if (type.getRank() <= 3) return src;
+ auto shape = type.getShape();
+ SmallVector<int64_t, 4> 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<mhlo::ReshapeOp>(
+ loc, RankedTensorType::get(result_shape, type.getElementType()), src);
+ }
+
LogicalResult matchAndRewrite(mhlo::DotGeneralOp op,
PatternRewriter &rewriter) const override {
auto lhsShapeType = op.lhs().getType().dyn_cast<RankedTensorType>();
@@ -559,167 +581,53 @@
lhsTargetOrder);
Value rhs = TransposeIfNonConsecutive(rewriter, op.getLoc(), op.rhs(),
rhsTargetOrder);
- if (lhs == op.lhs() && rhs == op.rhs()) return failure();
+ // 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();
- auto lhsBatchingDimsAttr =
- llvm::to_vector<4>(llvm::seq<int64_t>(0, lhsBatchingDims.size()));
- auto rhsBatchingDimsAttr =
- llvm::to_vector<4>(llvm::seq<int64_t>(0, rhsBatchingDims.size()));
- auto lhsContractingDimsAttr = llvm::to_vector<4>(
- llvm::seq<int64_t>(lhsContractionBase, lhsShapeType.getRank()));
- auto rhsContractingDimsAttr = llvm::to_vector<4>(
- llvm::seq<int64_t>(rhsContractionBase, numRhsContractionDims));
+
+ lhs = ReshapeIfMorethan3D(rewriter, op.getLoc(), lhs,
+ rhsBatchingDims.size(), lhsContractionBase);
+ rhs = ReshapeIfMorethan3D(rewriter, op.getLoc(), rhs,
+ rhsBatchingDims.size(), numRhsContractionDims);
+
+ if (lhs == op.lhs() && rhs == op.rhs()) return failure();
+
auto dimensionNumbers = mhlo::DotDimensionNumbersAttr::get(
- rewriter.getContext(), lhsBatchingDimsAttr, rhsBatchingDimsAttr,
- lhsContractingDimsAttr, rhsContractingDimsAttr);
+ rewriter.getContext(), /*lhsBatchingDimensions=*/0,
+ /*rhsBatchingDimensions=*/0,
+ /*lhsContractingDimensions=*/2, /*rhsContractingDimensions=*/1);
+ auto lhsNewType = lhs.getType().cast<RankedTensorType>();
+ auto rhsNewType = rhs.getType().cast<RankedTensorType>();
+
+ // if lhs's shape or rhs's shape has collapsed, we need reshape the result
+ bool needReshapeResult = lhsNewType.getRank() < lhsShapeType.getRank() ||
+ rhsNewType.getRank() < rhsShapeType.getRank();
+ // batching、lhs parallel、rhs parallel this order is a convension
+ SmallVector<int64_t, 4> newShape = {lhsNewType.getShape()[0],
+ lhsNewType.getShape()[1],
+ rhsNewType.getShape()[2]};
+ auto newResultType =
+ needReshapeResult
+ ? RankedTensorType::get(newShape, resultType.getElementType())
+ : op.getType();
Value result = rewriter.create<mhlo::DotGeneralOp>(
- op.getLoc(), op.getType(), lhs, rhs, dimensionNumbers,
- op.precision_configAttr());
- rewriter.replaceOp(op, result);
- return success();
- }
-};
-
-// Rewrite mhlo.dot_general to operate on rank-3 tensors when reduction 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.
-class RankReducedDotGeneral : public OpRewritePattern<mhlo::DotGeneralOp> {
- public:
- using OpRewritePattern<mhlo::DotGeneralOp>::OpRewritePattern;
-
- LogicalResult matchAndRewrite(mhlo::DotGeneralOp op,
- PatternRewriter &rewriter) const override {
- auto lhsShapeType = op.lhs().getType().dyn_cast<ShapedType>();
- auto rhsShapeType = op.rhs().getType().dyn_cast<ShapedType>();
- auto resultType = op.getResult().getType().dyn_cast<ShapedType>();
-
- if (!lhsShapeType || !rhsShapeType || !resultType) return failure();
- if (!lhsShapeType.hasStaticShape() || !rhsShapeType.hasStaticShape())
- return failure();
- if (resultType.getRank() <= 3) return failure();
-
- mhlo::DotDimensionNumbersAttr dimNumbers = op.dot_dimension_numbers();
- auto lhsBatchingDims =
- llvm::to_vector<4>(dimNumbers.getLhsBatchingDimensions());
- auto rhsBatchingDims =
- llvm::to_vector<4>(dimNumbers.getRhsBatchingDimensions());
- auto lhsContractingDims =
- llvm::to_vector<4>(dimNumbers.getLhsContractingDimensions());
- auto rhsContractingDims =
- llvm::to_vector<4>(dimNumbers.getRhsContractingDimensions());
-
- if (lhsBatchingDims.empty() || rhsBatchingDims.empty()) return failure();
-
- llvm::sort(lhsBatchingDims);
- llvm::sort(lhsContractingDims);
- llvm::sort(rhsBatchingDims);
- llvm::sort(rhsContractingDims);
-
- auto isDomainSplit = [](ArrayRef<int64_t> shape,
- ArrayRef<int64_t> batchingDims,
- ArrayRef<int64_t> contractingDims) {
- // Batching and contracting are contiguous.
- if ((contractingDims.front() - batchingDims.back()) == 1) return false;
- // Contracting dims are inner most.
- if (contractingDims.back() == (shape.size() - 1)) return false;
- return true;
- };
-
- if (!isConsecutive(lhsBatchingDims) || !isConsecutive(lhsContractingDims) ||
- !isConsecutive(rhsBatchingDims) || !isConsecutive(rhsContractingDims))
- return failure();
-
- if (isDomainSplit(lhsShapeType.getShape(), lhsBatchingDims,
- lhsContractingDims) ||
- isDomainSplit(rhsShapeType.getShape(), rhsBatchingDims,
- rhsContractingDims))
- return failure();
-
- // Collapsing shape into a rank-3 tensor, returns newCollabsedShape
- // contraction and parallel dim indices.
- auto computeCollapsedShape = [](ArrayRef<int64_t> shape,
- ArrayRef<int64_t> batchingDims,
- ArrayRef<int64_t> contractingDims) {
- auto newRank =
- shape.size() - batchingDims.size() - contractingDims.size() + 2;
- auto batchingSize = std::accumulate(
- batchingDims.begin(), batchingDims.end(), 1,
- [shape](const int64_t accum, const int64_t index) -> int64_t {
- return accum * shape[index];
- });
- auto contractingSize = std::accumulate(
- contractingDims.begin(), contractingDims.end(), 1,
- [shape](const int64_t accum, const int64_t index) -> int64_t {
- return accum * shape[index];
- });
-
- int parallelDimIndex, contractingDimIndex, parallelDimSize = 1;
- if (contractingDims.front() - batchingDims.back() > 1) {
- parallelDimIndex = 1;
- contractingDimIndex = 2;
- for (int i = batchingDims.back() + 1; i < contractingDims.front();
- ++i) {
- parallelDimSize *= shape[i];
- }
- } else {
- contractingDimIndex = 1;
- parallelDimIndex = 2;
- for (int i = contractingDims.back() + 1; i < shape.size(); ++i) {
- parallelDimSize *= shape[i];
- }
- }
- llvm::SmallVector<int64_t, 4> newShape(newRank);
- newShape[0] = batchingSize;
- newShape[contractingDimIndex] = contractingSize;
- newShape[parallelDimIndex] = parallelDimSize;
- return std::make_tuple(newShape, contractingDimIndex, parallelDimIndex);
- };
-
- int lhsContractingDimIndex, rhsContractingDimIndex, lhsParallelDimIndex,
- rhsParallelDimIndex;
- SmallVector<int64_t, 4> lhsNewShape, rhsNewShape;
- std::tie(lhsNewShape, lhsContractingDimIndex, lhsParallelDimIndex) =
- computeCollapsedShape(lhsShapeType.getShape(), lhsBatchingDims,
- lhsContractingDims);
-
- std::tie(rhsNewShape, rhsContractingDimIndex, rhsParallelDimIndex) =
- computeCollapsedShape(rhsShapeType.getShape(), rhsBatchingDims,
- rhsContractingDims);
- SmallVector<int64_t, 4> resultNewShape = {lhsNewShape[0],
- lhsNewShape[lhsParallelDimIndex],
- rhsNewShape[rhsParallelDimIndex]};
- Type dotGeneralResultType =
- RankedTensorType::get(resultNewShape, resultType.getElementType());
-
- auto loc = op.getLoc();
- Value reshapedLhs = rewriter.create<mhlo::ReshapeOp>(
- loc, RankedTensorType::get(lhsNewShape, lhsShapeType.getElementType()),
- op.lhs());
- Value reshapedRhs = rewriter.create<mhlo::ReshapeOp>(
- loc, RankedTensorType::get(rhsNewShape, rhsShapeType.getElementType()),
- op.rhs());
- auto dimensionNumbers = mhlo::DotDimensionNumbersAttr::get(
- rewriter.getContext(),
- /*lhs_batching_dimensions=*/{0},
- /*rhs_batching_dimensions=*/{0},
- /*lhs_contracting_dimensions=*/{lhsContractingDimIndex},
- /*rhs_contracting_dimensions=*/
- {rhsContractingDimIndex});
- Value dotGeneralResult = rewriter.create<mhlo::DotGeneralOp>(
- loc, dotGeneralResultType, reshapedLhs, reshapedRhs, dimensionNumbers,
+ op.getLoc(), newResultType, lhs, rhs, dimensionNumbers,
op.precision_configAttr());
- Value result =
- rewriter.create<mhlo::ReshapeOp>(loc, resultType, dotGeneralResult);
+ if (needReshapeResult) {
+ result =
+ rewriter.create<mhlo::ReshapeOp>(op.getLoc(), resultType, result);
+ }
rewriter.replaceOp(op, result);
-
return success();
}
};
@@ -915,7 +823,7 @@
// dot_general canoncalization patterns.
mhlo::PopulateGeneralDotOpLoweringPatterns(&patterns, context);
- patterns.insert<RankReducedDotGeneral, TransposeGenericDotGeneral>(context);
+ patterns.insert<TransposeReshapeGenericDotGeneral>(context);
// Unary elementwise op.
patterns.insert<
diff --git a/iree/compiler/InputConversion/MHLO/test/mhlo_to_mhlo_preprocessing_canoncalize_dot_general.mlir b/iree/compiler/InputConversion/MHLO/test/mhlo_to_mhlo_preprocessing_canoncalize_dot_general.mlir
index fbef414..ee8ddf6 100644
--- a/iree/compiler/InputConversion/MHLO/test/mhlo_to_mhlo_preprocessing_canoncalize_dot_general.mlir
+++ b/iree/compiler/InputConversion/MHLO/test/mhlo_to_mhlo_preprocessing_canoncalize_dot_general.mlir
@@ -56,10 +56,10 @@
return %0 : tensor<1x8x32x32xf32>
}
// CHECK: dot_general_to_dot_general_rank_reduced_a_transposed(%[[ARG0:.+]]: tensor<1x8x64x32xf32>, %[[ARG1:.+]]: tensor<1x8x64x32xf32>) -> tensor<1x8x32x32xf32>
-// CHECK: %[[ARG0_RESHAPED:.+]] = "mhlo.reshape"(%[[ARG0]]) : (tensor<1x8x64x32xf32>) -> tensor<8x64x32xf32>
+// CHECK: %[[ARG0_RESHAPED_TR:.+]] = "mhlo.transpose"(%[[ARG0]]) {permutation = dense<[0, 1, 3, 2]> : tensor<4xi64>} : (tensor<1x8x64x32xf32>) -> tensor<1x8x32x64xf32>
+// CHECK: %[[ARG0_RESHAPED:.+]] = "mhlo.reshape"(%[[ARG0_RESHAPED_TR]]) : (tensor<1x8x32x64xf32>) -> tensor<8x32x64xf32>
// CHECK: %[[ARG1_RSSHAPED:.+]] = "mhlo.reshape"(%[[ARG1]]) : (tensor<1x8x64x32xf32>) -> tensor<8x64x32xf32>
-// CHECK: %[[ARG0_RESHAPED_TR:.+]] = "mhlo.transpose"(%[[ARG0_RESHAPED]]) {permutation = dense<[0, 2, 1]> : tensor<3xi64>} : (tensor<8x64x32xf32>) -> tensor<8x32x64xf32>
-// CHECK: %[[DOT_RESULT:.+]] = "mhlo.dot_general"(%[[ARG0_RESHAPED_TR]], %[[ARG1_RSSHAPED]])
+// CHECK: %[[DOT_RESULT:.+]] = "mhlo.dot_general"(%[[ARG0_RESHAPED]], %[[ARG1_RSSHAPED]])
// CHECK: %[[RESULT:.+]] = "mhlo.reshape"(%[[DOT_RESULT]]) : (tensor<8x32x32xf32>) -> tensor<1x8x32x32xf32>
// -----
@@ -69,7 +69,7 @@
dot_dimension_numbers = #mhlo.dot<
lhs_batching_dimensions = [0, 1],
lhs_contracting_dimensions = [3],
- rhs_batching_dimensions = [0,1 ],
+ rhs_batching_dimensions = [0, 1],
rhs_contracting_dimensions = [3],
>,
precision_config = ["DEFAULT", "DEFAULT"]
@@ -77,10 +77,10 @@
return %0 : tensor<1x8x32x32xf32>
}
// CHECK: dot_general_to_dot_general_rank_reduced_b_transposed(%[[ARG0:.+]]: tensor<1x8x32x64xf32>, %[[ARG1:.+]]: tensor<1x8x32x64xf32>) -> tensor<1x8x32x32xf32>
+// CHECK: %[[ARG1_RESHAPED_TR:.+]] = "mhlo.transpose"(%[[ARG1]]) {permutation = dense<[0, 1, 3, 2]> : tensor<4xi64>} : (tensor<1x8x32x64xf32>) -> tensor<1x8x64x32xf32>
// CHECK: %[[ARG0_RESHAPED:.+]] = "mhlo.reshape"(%[[ARG0]]) : (tensor<1x8x32x64xf32>) -> tensor<8x32x64xf32>
-// CHECK: %[[ARG1_RESHAPED:.+]] = "mhlo.reshape"(%[[ARG1]]) : (tensor<1x8x32x64xf32>) -> tensor<8x32x64xf32>
-// CHECK: %[[ARG1_RESHAPED_TR:.+]] = "mhlo.transpose"(%[[ARG1_RESHAPED]]) {permutation = dense<[0, 2, 1]> : tensor<3xi64>} : (tensor<8x32x64xf32>) -> tensor<8x64x32xf32>
-// CHECK: %[[DOT_RESULT:.+]] = "mhlo.dot_general"(%[[ARG0_RESHAPED]], %[[ARG1_RESHAPED_TR]])
+// CHECK: %[[ARG1_RESHAPED:.+]] = "mhlo.reshape"(%[[ARG1_RESHAPED_TR]]) : (tensor<1x8x64x32xf32>) -> tensor<8x64x32xf32>
+// CHECK: %[[DOT_RESULT:.+]] = "mhlo.dot_general"(%[[ARG0_RESHAPED]], %[[ARG1_RESHAPED]])
// CHECK: %[[RESULT:.+]] = "mhlo.reshape"(%[[DOT_RESULT]]) : (tensor<8x32x32xf32>) -> tensor<1x8x32x32xf32>
// CHECK: return %[[RESULT]] : tensor<1x8x32x32xf32>
@@ -100,11 +100,11 @@
return %0 : tensor<1x8x32x32xf32>
}
// CHECK: dot_general_to_dot_general_rank_reduced_ab_transposed(%[[ARG0:.+]]: tensor<1x8x64x32xf32>, %[[ARG1:.+]]: tensor<1x8x32x64xf32>) -> tensor<1x8x32x32xf32>
-// CHECK: %[[ARG0_RESHAPED:.+]] = "mhlo.reshape"(%[[ARG0]]) : (tensor<1x8x64x32xf32>) -> tensor<8x64x32xf32>
-// CHECK: %[[ARG1_RESHAPED:.+]] = "mhlo.reshape"(%[[ARG1]]) : (tensor<1x8x32x64xf32>) -> tensor<8x32x64xf32>
-// CHECK: %[[ARG0_RESHAPED_TR:.+]] = "mhlo.transpose"(%[[ARG0_RESHAPED]]) {permutation = dense<[0, 2, 1]> : tensor<3xi64>} : (tensor<8x64x32xf32>) -> tensor<8x32x64xf32>
-// CHECK: %[[ARG1_RESHAPED_TR:.+]] = "mhlo.transpose"(%[[ARG1_RESHAPED]]) {permutation = dense<[0, 2, 1]> : tensor<3xi64>} : (tensor<8x32x64xf32>) -> tensor<8x64x32xf32>
-// CHECK: %[[DOT_RESULT:.+]] = "mhlo.dot_general"(%[[ARG0_RESHAPED_TR]], %[[ARG1_RESHAPED_TR]])
+// CHECK: %[[ARG0_RESHAPED_TR:.+]] = "mhlo.transpose"(%[[ARG0]]) {permutation = dense<[0, 1, 3, 2]> : tensor<4xi64>} : (tensor<1x8x64x32xf32>) -> tensor<1x8x32x64xf32>
+// CHECK: %[[ARG1_RESHAPED_TR:.+]] = "mhlo.transpose"(%[[ARG1]]) {permutation = dense<[0, 1, 3, 2]> : tensor<4xi64>} : (tensor<1x8x32x64xf32>) -> tensor<1x8x64x32xf32>
+// CHECK: %[[ARG0_RESHAPED:.+]] = "mhlo.reshape"(%[[ARG0_RESHAPED_TR]]) : (tensor<1x8x32x64xf32>) -> tensor<8x32x64xf32>
+// CHECK: %[[ARG1_RESHAPED:.+]] = "mhlo.reshape"(%[[ARG1_RESHAPED_TR]]) : (tensor<1x8x64x32xf32>) -> tensor<8x64x32xf32>
+// CHECK: %[[DOT_RESULT:.+]] = "mhlo.dot_general"(%[[ARG0_RESHAPED]], %[[ARG1_RESHAPED]])
// CHECK: %[[RESULT:.+]] = "mhlo.reshape"(%[[DOT_RESULT]]) : (tensor<8x32x32xf32>) -> tensor<1x8x32x32xf32>
// CHECK: return %[[RESULT]] : tensor<1x8x32x32xf32>
@@ -135,3 +135,31 @@
// CHECK: %[[DOT_RESULT:.+]] = "mhlo.dot_general"(%[[ARG0_RESHAPED]], %[[ARG1_RESHAPED]])
// CHECK: %[[RESULT:.+]] = "mhlo.reshape"(%[[DOT_RESULT]]) : (tensor<8x1x512xf32>) -> tensor<1x8x1x512xf32>
// CHECK: return %[[RESULT]] : tensor<1x8x1x512xf32>
+
+// -----
+
+func @dot_general_1d_batching_1d_contracting(%arg0: tensor<64x155x4x36xf32>, %arg1: tensor<309x4x36xf32>) -> tensor<4x64x155x309xf32> {
+ %0 = "mhlo.dot_general"(%arg0, %arg1) {
+ dot_dimension_numbers = #mhlo.dot<
+ lhs_batching_dimensions = [2],
+ rhs_batching_dimensions = [1],
+ lhs_contracting_dimensions = [3],
+ rhs_contracting_dimensions = [2]
+ >} : (tensor<64x155x4x36xf32>, tensor<309x4x36xf32>) -> tensor<4x64x155x309xf32>
+ return %0 : tensor<4x64x155x309xf32>
+}
+
+// CHECK-LABEL: func @dot_general_1d_batching_1d_contracting
+// CHECK-SAME: %[[ARG0:[a-zA-Z0-9_]+]]
+// CHECK-SAME: %[[ARG1:[a-zA-Z0-9_]+]]
+// CHECK: %[[ARG0_RESHAPED_TR:.+]] = "mhlo.transpose"(%[[ARG0]])
+// CHECK-SAME: {permutation = dense<[2, 0, 1, 3]> : tensor<4xi64>}
+// CHECK-SAME: (tensor<64x155x4x36xf32>) -> tensor<4x64x155x36xf32>
+// CHECK: %[[ARG1_RESHAPED_TR:.+]] = "mhlo.transpose"(%[[ARG1]])
+// CHECK-SAME: {permutation = dense<[1, 2, 0]> : tensor<3xi64>}
+// CHECK-SAME: (tensor<309x4x36xf32>) -> tensor<4x36x309xf32>
+// CHECK: %[[ARG0_RESHAPED:.+]] = "mhlo.reshape"(%[[ARG0_RESHAPED_TR]])
+// CHECK-SAME: (tensor<4x64x155x36xf32>) -> tensor<4x9920x36xf32>
+// CHECK: %[[DOT_RESULT:.+]] = "mhlo.dot_general"(%[[ARG0_RESHAPED]], %[[ARG1_RESHAPED_TR]])
+// CHECK-SAME: (tensor<4x9920x36xf32>, tensor<4x36x309xf32>) -> tensor<4x9920x309xf32>
+// CHECK: "mhlo.reshape"(%[[DOT_RESULT]]) : (tensor<4x9920x309xf32>) -> tensor<4x64x155x309xf32>
diff --git a/iree/test/e2e/xla_ops/dot_general.mlir b/iree/test/e2e/xla_ops/dot_general.mlir
index 48e106f..d56b98a 100644
--- a/iree/test/e2e/xla_ops/dot_general.mlir
+++ b/iree/test/e2e/xla_ops/dot_general.mlir
@@ -141,3 +141,33 @@
check.expect_almost_eq_const(%res, dense<409.596> : tensor<4x32x64xf32>) : tensor<4x32x64xf32>
return
}
+
+func @dot_general_nontrivial_batching_mutliple_parallel_dimension() {
+ %lhs = util.unfoldable_constant dense<[
+ [[[0.0], [1.0]], [[2.0], [3.0]], [[ 4.0], [ 5.0]]],
+ [[[6.0], [7.0]], [[8.0], [9.0]], [[10.0], [11.0]]]
+ ]> : tensor<2x3x2x1xf32>
+ %rhs = util.unfoldable_constant dense<[
+ [[0.0], [1.0]], [[2.0], [3.0]]
+ ]> : tensor<2x2x1xf32>
+ %res = "mhlo.dot_general"(%lhs, %rhs) {
+ dot_dimension_numbers = #mhlo.dot<
+ lhs_batching_dimensions = [2],
+ rhs_batching_dimensions = [1],
+ lhs_contracting_dimensions = [3],
+ rhs_contracting_dimensions = [2]
+ >,
+ precision_config = ["DEFAULT", "DEFAULT"]
+ } : (tensor<2x3x2x1xf32>, tensor<2x2x1xf32>) -> tensor<2x2x3x2xf32>
+ check.expect_almost_eq_const(%res, dense<[
+ [
+ [[0.0, 0.0], [0.0, 4.0], [0.0, 8.0]],
+ [[0.0, 12.0], [0.0, 16.0], [0.0, 20.0]]
+ ],
+ [
+ [[1.0, 3.0], [3.0, 9.0], [ 5.0, 15.0]],
+ [[7.0, 21.0], [9.0, 27.0], [11.0, 33.0]]
+ ]
+ ]> : tensor<2x2x3x2xf32>) : tensor<2x2x3x2xf32>
+ return
+}
