[Codegen] Introduce VectorizableOpInterface and migrate LinalgExt::GatherOp by hanhanW · Pull Request #23653 · iree-org/iree

hanhanW · 2026-03-05T00:20:31Z

Introduce a VectorizableOpInterface that provides an unified extension point for vectorization. Operations implement this interface to declare, whether it is vectorizable, how they should be vectorized, and the driver dispatches through it.

Implement external model for IREE::LinalgExt::GatherOp, inlining the vectorization logic from vectorizeLinalgExtGatherToTransferGather.
Update GenericVectorization.cpp to dispatch through the interface for ops that implement it, with TypeSwitch fallback for unmigrated ops.

No new tests because it is an NFC in terms of functionality. It just follows different mechanism for vectorization.

It is a step towards https://lists.lfaidata.foundation/g/iree-technical-discussion/message/15

Assisted-by: Claude

…therOp Introduce a VectorizableOpInterface that provides an unified extension point for vectorization. Operations implement this interface to declare, whether it is vectorizable, how they should be vectorized, and the driver dispatches through it. - Implement external model for IREE::LinalgExt::GatherOp, inlining the vectorization logic from vectorizeLinalgExtGatherToTransferGather. - Update GenericVectorization.cpp to dispatch through the interface for ops that implement it, with TypeSwitch fallback for unmigrated ops. It is a step towards https://lists.lfaidata.foundation/g/iree-technical-discussion/message/15 Assisted-by: Claude Signed-off-by: hanhanW <[email protected]>

Signed-off-by: hanhanW <[email protected]>

sommerlukas

Just a small nit, but please also wait for @Groverkss to review.

sommerlukas · 2026-03-10T08:55:18Z

compiler/src/iree/compiler/Codegen/Interfaces/VectorizableOpInterface.cpp

+                                          DictionaryAttr options) const {
+    auto gatherOp = cast<IREE::LinalgExt::GatherOp>(op);
+    int64_t batchRank = gatherOp.getBatchRank();
+    Location loc = gatherOp.getLoc();


Might be worth creating an ImplicitLocOpBuilder here and use it in the rest of this function.

I tried ImplicitLocOpBuilder when it was created, but I did not like the idea much. It was useful if you create operations in a single line, though I still don't like the idea.

The other point is that it extends OpBuilder, not RewriterBase, so methods like modifyOpInPlace don't work.

I'm not sure what you mean by

operations in a single line

If we create an ImplicitLocOpBuilder once in the beginning of vectorize, there is nine places where we pass rewriter, loc and could be passing that builder instead.

For things like modifyOpInPlace, you can keep using the rewriter, I think.

Groverkss · 2026-03-10T16:18:25Z

compiler/src/iree/compiler/Codegen/Interfaces/VectorizableOpInterface.cpp

+    auto indicesVecRead = vector::TransferReadOp::create(
+        rewriter, loc, indicesVecTy.clone(indicesTy.getElementType()),
+        gatherOp.getIndices(), SmallVector<Value>(indicesTy.getRank(), zero),
+        std::nullopt);


Note that eventually I would like to move this to an interface. Using transfer_read/transfer_write to vectorize something is not a very nice way of vectorization. We should be doing dialect conversion imo. We can do that after the VectorTile size analysis lands.

Is there a plan for VectorTile size analysis yet? Is it something similar to what we implemented:

iree/compiler/src/iree/compiler/Codegen/Utils/Utils.h

Lines 363 to 372 in 4e3da92

struct VectorizationTileSizes {

SmallVector<int64_t> destShape;

SmallVector<int64_t> vectorSizes;

SmallVector<bool> vectorScalableFlags;

};

/// Returns a VectorizationTileSizes which contains the inferred bounded result

/// shape and vector input sizes. This is useful to infer the sizes from a

/// chain.

std::optional<VectorizationTileSizes> inferSizesFromIR(Value val);

There is a PR with a first attempt at the analysis here: #23668

The idea is that we take the tile sizes from to_layout operations (that come from lowering_config) and propagate that to the other compute ops. It's then annotated as a discardable attribute, so we can use it in GenericVectorization. There, we use three sources of information in this order:

lowering_config

Result of the VectorTileSizeAnalysis

inferSizesFromIR

The quick comment is about getVectorScalableFlags method. The lowering_config can implement the method to provide vector tile sizes. It's a little fragile because lowering config can be messed up during the transformation, so it is default off on x86 backends.

iree/compiler/src/iree/compiler/Codegen/Common/GenericVectorization.cpp

Lines 41 to 43 in 4e3da92

std::optional<SmallVector<int64_t>> vectorSizes =

loweringConfig.getVectorSizes();

SmallVector<bool> scalableFlags = loweringConfig.getVectorScalableFlags();

So, you mean we need to propagate the scalable flags in the analysis?

I meant the design idea was using the interface method, if the lowering config has the information. Scalable flags is just something similar to vector size and only some CPU backends use it.

hanhanW force-pushed the users/hanhanW/vec-iface-c1 branch from bc0a2cb to 16bcb57 Compare March 5, 2026 00:48

hanhanW force-pushed the users/hanhanW/vec-iface-c1 branch from 16bcb57 to c35263b Compare March 5, 2026 00:48

hanhanW marked this pull request as ready for review March 5, 2026 00:57

hanhanW requested review from Groverkss, MaheshRavishankar, Max191 and qedawkins as code owners March 5, 2026 00:57

Rename iface -> vectorizableOp for VectorizableOpInterface variables.

1df2af0

Signed-off-by: hanhanW <[email protected]>

This was referenced Mar 6, 2026

[Codegen] Migrate ArgCompareOp to VectorizableOpInterface #23656

Merged

[Codegen] Migrate ToLayoutOp to VectorizableOpInterface #23658

Merged

[Codegen] Migrate MapStoreOp to VectorizableOpInterface #23662

Merged

Improve documentation

c892ec4

Signed-off-by: hanhanW <[email protected]>

This was referenced Mar 9, 2026

[Codegen] Migrate PadOp/PackOp/UnPackOp to VectorizableOpInterface #23712

Merged

[Codegen] Migrate all ops to VectorizableOpInterface and remove TypeSwitch #23713

Merged

sommerlukas approved these changes Mar 10, 2026

View reviewed changes

Groverkss approved these changes Mar 10, 2026

View reviewed changes

Groverkss reviewed Mar 10, 2026

View reviewed changes

hanhanW merged commit 3b4bb3d into main Mar 10, 2026
59 of 61 checks passed

hanhanW deleted the users/hanhanW/vec-iface-c1 branch March 10, 2026 16:39

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[Codegen] Introduce VectorizableOpInterface and migrate LinalgExt::GatherOp#23653

[Codegen] Introduce VectorizableOpInterface and migrate LinalgExt::GatherOp#23653
hanhanW merged 3 commits intomainfrom
users/hanhanW/vec-iface-c1

hanhanW commented Mar 5, 2026 •

edited

Loading

Uh oh!

sommerlukas left a comment

Uh oh!

sommerlukas Mar 10, 2026

Uh oh!

hanhanW Mar 10, 2026

Uh oh!

sommerlukas Mar 10, 2026

Uh oh!

Groverkss Mar 10, 2026

Uh oh!

hanhanW Mar 10, 2026

Uh oh!

sommerlukas Mar 10, 2026

Uh oh!

hanhanW Mar 10, 2026

Uh oh!

sommerlukas Mar 10, 2026

Uh oh!

hanhanW Mar 10, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

	struct VectorizationTileSizes {
	SmallVector<int64_t> destShape;
	SmallVector<int64_t> vectorSizes;
	SmallVector<bool> vectorScalableFlags;
	};

	/// Returns a VectorizationTileSizes which contains the inferred bounded result
	/// shape and vector input sizes. This is useful to infer the sizes from a
	/// chain.
	std::optional<VectorizationTileSizes> inferSizesFromIR(Value val);

	std::optional<SmallVector<int64_t>> vectorSizes =
	loweringConfig.getVectorSizes();
	SmallVector<bool> scalableFlags = loweringConfig.getVectorScalableFlags();

Conversation

hanhanW commented Mar 5, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

sommerlukas left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

hanhanW commented Mar 5, 2026 •

edited

Loading