[mlir] `-test-vector-unrolling-patterns` pass failed with assertion error "cast<Ty>() argument of incompatible type!"

[Colloportus0]

MLIR built at commit 585cbe3
Reproduced with:
mlir-opt -test-vector-unrolling-patterns=unroll-based-on-type temp.mlir

temp.mlir:

module {
  func.func nested @func1(%arg0: i64, %arg1: memref<?x?xf16>, %arg2: i16) -> f32 {
    %cst_1 = arith.constant 1.0 : f32
    %23 = vector.broadcast %cst_1 : f32 to vector<23x8xf32>
    %24 = vector.extract_strided_slice %23 {offsets = [14], sizes = [6], strides = [1]} : vector<23x8xf32> to vector<6x8xf32>
    %115 = vector.multi_reduction <minf>, %24, %cst_1 [0, 1] : vector<6x8xf32> to f32
    return %115 : f32
  }
}

trace:

Assertion failed: (isa<To>(Val) && "cast<Ty>() argument of incompatible type!"), function cast, file Casting.h, line 566.
PLEASE submit a bug report to https://github.com/llvm/llvm-project/issues/ and include the crash backtrace.
Stack dump:
0.	Program arguments: mlir-opt -test-vector-unrolling-patterns=unroll-based-on-type temp.mlir
 #0 0x000000010310cf18 llvm::sys::PrintStackTrace(llvm::raw_ostream&, int) (/workspace/build/bin/mlir-opt+0x1002f8f18)
 #1 0x000000010310b1b4 llvm::sys::RunSignalHandlers() (/workspace/build/bin/mlir-opt+0x1002f71b4)
 #2 0x000000010310d5c4 SignalHandler(int) (/workspace/build/bin/mlir-opt+0x1002f95c4)
 #3 0x00000001be7894c4 (/usr/lib/system/libsystem_platform.dylib+0x1803414c4)
 #4 0x00000001be771ee0 (/usr/lib/system/libsystem_pthread.dylib+0x180329ee0)
 #5 0x00000001be6ac340 (/usr/lib/system/libsystem_c.dylib+0x180264340)
 #6 0x00000001be6ab754 (/usr/lib/system/libsystem_c.dylib+0x180263754)
 #7 0x00000001074e7bf0 inferStridedSliceOpResultType(mlir::VectorType, mlir::ArrayAttr, mlir::ArrayAttr, mlir::ArrayAttr) (.cold.1) (/workspace/build/bin/mlir-opt+0x1046d3bf0)
 #8 0x0000000104db4dd4 mlir::vector::ExtractStridedSliceOp::build(mlir::OpBuilder&, mlir::OperationState&, mlir::Value, llvm::ArrayRef<long long>, llvm::ArrayRef<long long>, llvm::ArrayRef<long long>) (/workspace/build/bin/mlir-opt+0x101fa0dd4)
 #9 0x0000000104ec294c mlir::vector::ExtractStridedSliceOp mlir::OpBuilder::create<mlir::vector::ExtractStridedSliceOp, mlir::Value, llvm::SmallVector<long long, 6u>&, llvm::SmallVector<long long, 6u>&, llvm::SmallVector<long long, 6u>&>(mlir::Location, mlir::Value&&, llvm::SmallVector<long long, 6u>&, llvm::SmallVector<long long, 6u>&, llvm::SmallVector<long long, 6u>&) (/workspace/build/bin/mlir-opt+0x1020ae94c)
#10 0x0000000104ec3880 (anonymous namespace)::UnrollMultiReductionPattern::matchAndRewrite(mlir::vector::MultiDimReductionOp, mlir::PatternRewriter&) const (/workspace/build/bin/mlir-opt+0x1020af880)
#11 0x0000000106d8f0f0 mlir::PatternApplicator::matchAndRewrite(mlir::Operation*, mlir::PatternRewriter&, llvm::function_ref<bool (mlir::Pattern const&)>, llvm::function_ref<void (mlir::Pattern const&)>, llvm::function_ref<mlir::LogicalResult (mlir::Pattern const&)>)::$_2::operator()() const (/workspace/build/bin/mlir-opt+0x103f7b0f0)
#12 0x0000000106d8c3ec mlir::PatternApplicator::matchAndRewrite(mlir::Operation*, mlir::PatternRewriter&, llvm::function_ref<bool (mlir::Pattern const&)>, llvm::function_ref<void (mlir::Pattern const&)>, llvm::function_ref<mlir::LogicalResult (mlir::Pattern const&)>) (/workspace/build/bin/mlir-opt+0x103f783ec)
#13 0x00000001057bb420 (anonymous namespace)::GreedyPatternRewriteDriver::processWorklist() (/workspace/build/bin/mlir-opt+0x1029a7420)
#14 0x00000001057b8c50 mlir::applyPatternsAndFoldGreedily(mlir::Region&, mlir::FrozenRewritePatternSet const&, mlir::GreedyRewriteConfig, bool*) (/workspace/build/bin/mlir-opt+0x1029a4c50)
#15 0x0000000105714640 (anonymous namespace)::TestVectorUnrollingPatterns::runOnOperation() (/workspace/build/bin/mlir-opt+0x102900640)
#16 0x000000010575c2a8 mlir::detail::OpToOpPassAdaptor::run(mlir::Pass*, mlir::Operation*, mlir::AnalysisManager, bool, unsigned int) (/workspace/build/bin/mlir-opt+0x1029482a8)
#17 0x000000010575c998 mlir::detail::OpToOpPassAdaptor::runPipeline(mlir::OpPassManager&, mlir::Operation*, mlir::AnalysisManager, bool, unsigned int, mlir::PassInstrumentor*, mlir::PassInstrumentation::PipelineParentInfo const*) (/workspace/build/bin/mlir-opt+0x102948998)
#18 0x000000010575da4c mlir::detail::OpToOpPassAdaptor::runOnOperationAsyncImpl(bool) (/workspace/build/bin/mlir-opt+0x102949a4c)
#19 0x000000010575c404 mlir::detail::OpToOpPassAdaptor::run(mlir::Pass*, mlir::Operation*, mlir::AnalysisManager, bool, unsigned int) (/workspace/build/bin/mlir-opt+0x102948404)
#20 0x000000010575c998 mlir::detail::OpToOpPassAdaptor::runPipeline(mlir::OpPassManager&, mlir::Operation*, mlir::AnalysisManager, bool, unsigned int, mlir::PassInstrumentor*, mlir::PassInstrumentation::PipelineParentInfo const*) (/workspace/build/bin/mlir-opt+0x102948998)
#21 0x000000010575e82c mlir::PassManager::run(mlir::Operation*) (/workspace/build/bin/mlir-opt+0x10294a82c)
#22 0x0000000105756d14 performActions(llvm::raw_ostream&, std::__1::shared_ptr<llvm::SourceMgr> const&, mlir::MLIRContext*, mlir::MlirOptMainConfig const&) (/workspace/build/bin/mlir-opt+0x102942d14)
#23 0x0000000105756368 mlir::LogicalResult llvm::function_ref<mlir::LogicalResult (std::__1::unique_ptr<llvm::MemoryBuffer, std::__1::default_delete<llvm::MemoryBuffer>>, llvm::raw_ostream&)>::callback_fn<mlir::MlirOptMain(llvm::raw_ostream&, std::__1::unique_ptr<llvm::MemoryBuffer, std::__1::default_delete<llvm::MemoryBuffer>>, mlir::DialectRegistry&, mlir::MlirOptMainConfig const&)::$_1>(long, std::__1::unique_ptr<llvm::MemoryBuffer, std::__1::default_delete<llvm::MemoryBuffer>>, llvm::raw_ostream&) (/workspace/build/bin/mlir-opt+0x102942368)
#24 0x00000001057d434c mlir::splitAndProcessBuffer(std::__1::unique_ptr<llvm::MemoryBuffer, std::__1::default_delete<llvm::MemoryBuffer>>, llvm::function_ref<mlir::LogicalResult (std::__1::unique_ptr<llvm::MemoryBuffer, std::__1::default_delete<llvm::MemoryBuffer>>, llvm::raw_ostream&)>, llvm::raw_ostream&, bool, bool) (/workspace/build/bin/mlir-opt+0x1029c034c)
#25 0x0000000105750d1c mlir::MlirOptMain(llvm::raw_ostream&, std::__1::unique_ptr<llvm::MemoryBuffer, std::__1::default_delete<llvm::MemoryBuffer>>, mlir::DialectRegistry&, mlir::MlirOptMainConfig const&) (/workspace/build/bin/mlir-opt+0x10293cd1c)
#26 0x000000010575118c mlir::MlirOptMain(int, char**, llvm::StringRef, mlir::DialectRegistry&) (/workspace/build/bin/mlir-opt+0x10293d18c)
#27 0x0000000102e17510 main (/workspace/build/bin/mlir-opt+0x100003510)
#28 0x000000010e5b1088 

[llvmbot]

@llvm/issue-subscribers-mlir

[rikhuijzer]

I have found the root cause, but I give up on fixing this for now because I cannot figure it what to change. Maybe later.

The root cause is that vector.multi_reduction here returns an f32. This results in a problem because loop unrolling for vector.multi_reduction means applying multiple vector.extract_strided_slices. Per convention, vector.extract_strided_slice will always return a vector, but the vector.multi_reduction requests a f32.

A minimum example that also crashes:

func.func @multi_reduction(%0: vector<4x2xf32>, %acc1: f32) -> f32 {
  %2 = vector.multi_reduction <add>, %0, %acc1 [0, 1] : vector<4x2xf32> to f32
  return %2 : f32
}

To see that these return types make sense, take the following Julia example:

julia> A = reshape(1:(3*3), 3, 3)
3×3 reshape(::UnitRange{Int64}, 3, 3) with eltype Int64:
 1  4  7
 2  5  8
 3  6  9

julia> A[2:2, 2:2] # extract strided slice with offset = [0, 0], sizes = [1, 1], stride = [1, 1] 
1×1 Matrix{Int64}:
 5

So, indeed, it makes sense that a strided slice returns a n-dimensional vector even for a single scalar result.

Next, reduce doesn't necessarily have to return a vector:

julia> reduce(min, A[2:2, 2:2])
5

This problem with mismatching types most likely lies somewhere in the UnrollMultiReductionPattern matchAndRewrite:

llvm-project/mlir/lib/Dialect/Vector/Transforms/VectorUnroll.cpp

Lines 348 to 415 in c6f0f88

	struct UnrollMultiReductionPattern
	: public OpRewritePattern<vector::MultiDimReductionOp> {
	UnrollMultiReductionPattern(MLIRContext *context,
	const vector::UnrollVectorOptions &options,
	PatternBenefit benefit = 1)
	: OpRewritePattern<vector::MultiDimReductionOp>(context, benefit),
	options(options) {}
	LogicalResult matchAndRewrite(vector::MultiDimReductionOp reductionOp,
	PatternRewriter &rewriter) const override {
	std::optional<SmallVector<int64_t>> targetShape =
	getTargetShape(options, reductionOp);
	if (!targetShape)
	return failure();
	SmallVector<int64_t> originalSize = *reductionOp.getShapeForUnroll();
	llvm::MapVector<
	SmallVector<int64_t>, Value,
	llvm::DenseMap<SmallVector<int64_t>, unsigned, OffsetMapInfo>>
	accCache;
	Location loc = reductionOp.getLoc();
	// Stride of the ratios, this gives us the offsets of sliceCount in a basis
	// of multiples of the targetShape.
	for (SmallVector<int64_t> offsets :
	StaticTileOffsetRange(originalSize, *targetShape)) {
	SmallVector<Value> operands;
	SmallVector<int64_t> operandStrides(offsets.size(), 1);
	Value slicedOperand = rewriter.create<vector::ExtractStridedSliceOp>(
	loc, reductionOp.getSource(), offsets, *targetShape, operandStrides);
	operands.push_back(slicedOperand);
	SmallVector<int64_t> dstShape;
	SmallVector<int64_t> destOffset;
	for (size_t i : llvm::seq(size_t(0), targetShape->size())) {
	if (!reductionOp.isReducedDim(i)) {
	destOffset.push_back(offsets[i]);
	dstShape.push_back((*targetShape)[i]);
	}
	}
	Value acc;
	SmallVector<int64_t> accStrides(destOffset.size(), 1);
	// If a version of the accumulator has already been computed, use it
	// otherwise extract the first version from the original operand.
	auto accIt = accCache.find(destOffset);
	if (accIt != accCache.end())
	acc = accIt->second;
	else
	acc = rewriter.create<vector::ExtractStridedSliceOp>(
	loc, reductionOp.getAcc(), destOffset, dstShape, accStrides);
	operands.push_back(acc);
	auto targetType = VectorType::get(
	dstShape, reductionOp.getSourceVectorType().getElementType());
	Operation *newOp = cloneOpWithOperandsAndTypes(rewriter, loc, reductionOp,
	operands, targetType);
	Value result = newOp->getResult(0);
	accCache[destOffset] = result;
	}
	// Assemble back the accumulator into a single vector.
	Value result = rewriter.create<arith::ConstantOp>(
	loc, reductionOp.getDestType(),
	rewriter.getZeroAttr(reductionOp.getDestType()));
	for (const auto &it : accCache) {
	SmallVector<int64_t> dstStrides(it.first.size(), 1);
	result = rewriter.create<vector::InsertStridedSliceOp>(
	loc, it.second, result, it.first, dstStrides);
	}
	rewriter.replaceOp(reductionOp, result);
	return success();
	}

[Prakhar-Dixit]

I was looking forward to working on this issue. Can this issue be assigned to me?
@banach-space @Lewuathe @makslevental

[makslevental]

@Prakhar-Dixit yea sure - good luck and god speed 😄

[Prakhar-Dixit]

After investigating this issue, I discovered that the ExtractStridedSliceOp only handles vectors and does not support reducing vectors to scalar values. I reviewed the previous version of the code (see VectorUnrollDistribute.cpp in this commit) to check if anything was missed. The code was essentially the same and still crashes for this particular case.

Minimal Example of the Crashing IR:

func.func @multi_reduction(%0: vector<4x2xf32>, %acc1: f32) -> f32 {
  %2 = vector.multi_reduction <add>, %0, %acc1 [0, 1] : vector<4x2xf32> to f32
  return %2 : f32
}

The IR appears valid because it lowers correctly and shows no signs of being invalid. Therefore, I believe the proper fix is to bail out instead of handling scalars, as I assume it was never intended to do so when the operation is asked to reduce a vector to a scalar. For example:

auto resultType = reductionOp->getResult(0).getType();
if (mlir::isa<mlir::FloatType>(resultType) || mlir::isa<mlir::IntegerType>(resultType)) {
  return failure();
}

Would love to hear your thoughts on this,
@makslevental @dcaballe @nicolasvasilache.