Advance llvm-project and stablehlo. (#2619)

llvm-project: bbd2b08b95fe76bea138c1b03c1cd42ed3ee04df
stablehlo: ab709fe48de88c67717abfbd7ef17425eb95ddaf

These commits were chosen in order to account for an MLIR API break from
3dbac2c007
which required a patch to stablehlo. We integrate a bit beyond that
commit to deal with some revert/reapply cycles in the intervening range
which were discovered in another downstream.

Further, it requires adaptation to the stablehlo API breaks introduced
from https://github.com/openxla/stablehlo/pull/1872 which are along for
the ride.

Since some stablehlo builders were changed to directly take int64_t
array refs, also traced that up some call stacks to eliminate some
signed/unsigned mismatches that result.

Also adds a few TOSA tests to the passing set that seem to work now.
pull/2623/head snapshot-20231208.1046
Stella Laurenzo 2023-12-07 23:13:42 -08:00 committed by GitHub
parent 63505ad6b2
commit 8252656b6d
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
7 changed files with 43 additions and 51 deletions

@ -1 +1 @@
Subproject commit 5e5a22caf88ac1ccfa8dc5720295fdeba0ad9372
Subproject commit bbd2b08b95fe76bea138c1b03c1cd42ed3ee04df

2
externals/stablehlo vendored

@ -1 +1 @@
Subproject commit 83f095e7217c897f1eccac5652600ceb944cb0e0
Subproject commit ab709fe48de88c67717abfbd7ef17425eb95ddaf

View File

@ -51,7 +51,7 @@ Value promoteType(PatternRewriter &rewriter, Location loc, Value input,
Value promoteAndBroadcast(ConversionPatternRewriter &rewriter, Value input,
TensorType outType);
SmallVector<size_t> toPositiveDims(ArrayRef<int64_t> dims, int64_t rank);
SmallVector<int64_t> toPositiveDims(ArrayRef<int64_t> dims, int64_t rank);
// Get the dimension sizes of the input tensor, given the dimension axes
FailureOr<SmallVector<Value, 4>> getDimSizesOfTensor(PatternRewriter &rewriter,

View File

@ -615,12 +615,8 @@ public:
SmallVector<int64_t> permValues(inputRank);
std::iota(std::begin(permValues), std::end(permValues), 0);
std::swap(permValues[dim0], permValues[dim1]);
DenseIntElementsAttr permutation = DenseIntElementsAttr::get(
RankedTensorType::get({static_cast<long int>(permValues.size())},
rewriter.getI64Type()),
permValues);
rewriter.replaceOpWithNewOp<stablehlo::TransposeOp>(op, outType, self,
permutation);
permValues);
return success();
}
};
@ -793,12 +789,8 @@ LogicalResult ConvertAtenOp<AtenPermuteOp>::matchAndRewrite(
return op.emitError("not all dims are valid");
}
DenseIntElementsAttr permutation = DenseIntElementsAttr::get(
RankedTensorType::get({static_cast<long int>(permValues.size())},
rewriter.getI64Type()),
permValues);
rewriter.replaceOpWithNewOp<stablehlo::TransposeOp>(op, outType, self,
permutation);
permValues);
return success();
}
@ -1755,8 +1747,7 @@ LogicalResult ConvertAtenOp<AtenFlipOp>::matchAndRewrite(
}
}
rewriter.replaceOpWithNewOp<stablehlo::ReverseOp>(
op, outType, self, rewriter.getI64TensorAttr(dims));
rewriter.replaceOpWithNewOp<stablehlo::ReverseOp>(op, outType, self, dims);
return success();
}

View File

@ -62,13 +62,9 @@ Value getPermutedTensor(PatternRewriter &rewriter, Operation *op, Value input,
newShape.push_back(inpShape[d]);
}
auto attrTy = RankedTensorType::get({static_cast<int64_t>(transDims.size())},
rewriter.getIntegerType(64));
auto permuteAttr = DenseIntElementsAttr::get(attrTy, transDims);
auto outTy = RankedTensorType::get(newShape, inputTy.getElementType());
auto result = rewriter.create<stablehlo::TransposeOp>(op->getLoc(), outTy,
input, permuteAttr);
input, transDims);
return result.getResult();
}
@ -500,8 +496,8 @@ public:
for (int64_t i = 0; i <= rank; i++)
transposeDims[i] = i;
std::swap(transposeDims[rank - 1], transposeDims[rank - 2]);
weight = rewriter.create<stablehlo::TransposeOp>(
op->getLoc(), weight, rewriter.getI64TensorAttr(transposeDims));
weight = rewriter.create<stablehlo::TransposeOp>(op->getLoc(), weight,
transposeDims);
// 3. [H, W, ..., G, OC, IC//G] => [H, W, ..., G*OC, IC//G]
weightShapeInt.erase(weightShapeInt.end() - 2);
@ -546,12 +542,10 @@ public:
}
auto transposeTy =
RankedTensorType::get(transposeShape, weightTy.getElementType());
DenseIntElementsAttr permAttr = DenseIntElementsAttr::get(
RankedTensorType::get({nDims}, rewriter.getI64Type()), perm);
auto transposeOp = rewriter.create<stablehlo::TransposeOp>(
op->getLoc(), transposeTy, weight, permAttr);
op->getLoc(), transposeTy, weight, perm);
auto reverseOp = rewriter.create<stablehlo::ReverseOp>(
op->getLoc(), transposeOp, rewriter.getI64TensorAttr({0, 1}));
op->getLoc(), transposeOp, ArrayRef<int64_t>{0, 1});
// Prepare for transposed convolution
SmallVector<int64_t> stablehloStrideVec(nSpatialDims, 1);

View File

@ -250,12 +250,12 @@ Value promoteAndBroadcast(ConversionPatternRewriter &rewriter, Value input,
return bcast_op.getResult();
}
SmallVector<size_t> toPositiveDims(ArrayRef<int64_t> dims, int64_t rank) {
SmallVector<size_t> posDims;
SmallVector<int64_t> toPositiveDims(ArrayRef<int64_t> dims, int64_t rank) {
SmallVector<int64_t> posDims;
posDims.reserve(rank);
std::transform(
dims.begin(), dims.end(), std::back_inserter(posDims),
[rank](int64_t d) -> size_t { return toPositiveDim(d, rank); });
[rank](int64_t d) -> int64_t { return toPositiveDim(d, rank); });
return posDims;
}
@ -316,10 +316,10 @@ FailureOr<Value> unsqueezeTensor(PatternRewriter &rewriter, Operation *op,
op, "failed to get dimension sizes of the input");
auto dimSizes = *dimSizesInfo;
auto rank = dimSizes.size();
size_t newRank = rank + inputUnsqzDims.size();
int64_t rank = dimSizes.size();
int64_t newRank = rank + inputUnsqzDims.size();
auto unsqzDims = toPositiveDims(inputUnsqzDims, newRank);
for (size_t k = 0, sz = unsqzDims.size(); k < sz; ++k)
for (int64_t k = 0, sz = unsqzDims.size(); k < sz; ++k)
if (k > 1 && unsqzDims[k] <= unsqzDims[k - 1])
return rewriter.notifyMatchFailure(
op, "unsqueeze dimensions must be specified in order");
@ -335,8 +335,8 @@ FailureOr<Value> unsqueezeTensor(PatternRewriter &rewriter, Operation *op,
std::vector<int64_t> newShape;
newDimSizes.reserve(newRank);
newShape.reserve(newRank);
for (size_t k = 0, i = 0, j = 0; k < newRank; ++k) {
if (j < unsqzDims.size() && unsqzDims[j] == k) {
for (int64_t k = 0, i = 0, j = 0; k < newRank; ++k) {
if (j < static_cast<int64_t>(unsqzDims.size()) && unsqzDims[j] == k) {
newDimSizes.push_back(one);
newShape.push_back(1);
j++;

View File

@ -13,6 +13,8 @@
from torch_mlir_e2e_test.test_suite import COMMON_TORCH_MLIR_LOWERING_XFAILS
from torch_mlir._version import torch_version_for_comparison, version
print(f"TORCH_VERSION_FOR_COMPARISON =", torch_version_for_comparison())
LINALG_XFAIL_SET = COMMON_TORCH_MLIR_LOWERING_XFAILS | {
# Lowering Torch Backend IR -> Linalg-on-Tensors Backend IR failed
# 'linalg.depthwise_conv_2d_nchw_chw' op inferred input/output operand #1 has shape's dimension #0 to be 4, but found 8
@ -21,6 +23,14 @@ LINALG_XFAIL_SET = COMMON_TORCH_MLIR_LOWERING_XFAILS | {
"IscloseStaticModuleTrue_basic"
}
if torch_version_for_comparison() >= version.parse("2.2.0.dev20231204"):
LINALG_XFAIL_SET |= {
"Conv2dWithPaddingDilationStrideStaticModule_grouped",
"Conv2dWithPaddingDilationStrideStaticModule_grouped_multiplier",
"ConvolutionModule2DGroups_basic",
}
TORCHDYNAMO_XFAIL_SET = {
#### General TorchDynamo/PyTorch errors
@ -306,10 +316,11 @@ TORCHDYNAMO_XFAIL_SET = {
"ArangeStartOutViewModule_basic",
}
if torch_version_for_comparison() < version.parse("2.1.0.dev"):
TORCHDYNAMO_XFAIL_SET -= {
"ScaledDotProductAttentionSameModule_basic",
"ScaledDotProductAttentionDifferentModule_basic",
if torch_version_for_comparison() >= version.parse("2.2.0.dev20231204"):
TORCHDYNAMO_XFAIL_SET |= {
"Conv2dWithPaddingDilationStrideStaticModule_grouped",
"Conv2dWithPaddingDilationStrideStaticModule_grouped_multiplier",
"ConvolutionModule2DGroups_basic",
}
TORCHDYNAMO_CRASHING_SET = {
@ -1305,6 +1316,10 @@ TOSA_PASS_SET = {
"MeanModule_basic",
"ArangeStartOutModule_basic",
"ArangeStartOutViewModule_basic",
"Conv2dBiasNoPaddingModule_basic",
"Conv2dNoPaddingModule_basic",
"Conv2dWithPaddingDilationStrideModule_basic",
"Conv2dWithPaddingModule_basic",
}
MAKE_FX_TOSA_PASS_SET = (TOSA_PASS_SET | {
@ -1335,21 +1350,13 @@ MAKE_FX_TOSA_PASS_SET = (TOSA_PASS_SET | {
# failed to legalize operation 'torch.aten.to.dtype' that was explicitly marked illegal
"AtenEyeModuleInt2D_basic",
"AtenEyeMModuleInt2D_basic",
"Conv2dBiasNoPaddingModule_basic",
"Conv2dNoPaddingModule_basic",
"Conv2dWithPaddingDilationStrideModule_basic",
"Conv2dWithPaddingModule_basic",
}
if torch_version_for_comparison() < version.parse("2.1.0.dev"):
MAKE_FX_TOSA_PASS_SET -= {
# 'tensor.expand_shape' op expected rank expansion, but found source rank 1 >= result rank 1
"ReshapeCollapseModule_basic",
# failed to lower torch.aten.empty.memory_format
"BatchNorm1DModule_basic",
"BatchNorm1DWith2DInputModule_basic",
"BatchNorm2DModule_basic",
"BatchNorm3DModule_basic",
"BatchNorm1DStaticShapeModule_basic",
}
LTC_CRASHING_SET = {
# TODO: update test to move all inputs to the lazy device. Otherwise test fails with:
# Check failed: lazy_tensor Input tensor is not a lazy tensor: CPUBoolType.