[CPU] Simplify how tile sizes are updated #16435

banach-space · 2024-02-16T15:49:39Z

This is a follow-up for #16350 and is meant to simplify how tile sizes
are updated. In particular, a new wrapper for tile sizes is added,
SizesAndScalableFlagsTuple, that enables the following simplification:

BEFORE

vecTileSizes[idx] = innerVecTileSizes[idx];
vecScalableTileFlags[idx] = innerVecScalableTileFlags[idx];

(size and scalable flag updated separately)

AFTER

vecTileSizesAndFlags[idx] = innerVecTileSizesAndFlags[idx];

(size and scalable flag updated in one stmt)

The ultimate goal is to "hide" scalable flags for folks working on
targets that don't require those while preserving enough flexibility for
targets that do need to track this extra info. It should also simplify
further work (and review process) for future patches similar to #16350.

This is a follow-up for iree-org#16350 and is meant to simplify how tile sizes are updated. In particular, a new wrapper for tile sizes is added, `SizesAndScalableFlagsTuple`, that enables the following simplification: **BEFORE** ```cpp vecTileSizes[idx] = innerVecTileSizes[idx]; vecScalableTileFlags[idx] = innerVecScalableTileFlags[idx]; ``` (size and scalable flag updated separately) **AFTER** ```cpp vecTileSizesAndFlags[idx] = innerVecTileSizesAndFlags[idx]; ``` (size and scalable flag updated in one stmt) The ultimate goal is to "hide" scalable flags for folks working on targets that don't require those while preserving enough flexibility for targets that do need to track this extra info. It should also simplify further work (and review process) for future patches similar to iree-org#16350.

pzread · 2024-02-29T19:26:17Z

compiler/src/iree/compiler/Codegen/Common/TileSizeSelection.h

@@ -15,6 +15,75 @@ namespace mlir::iree_compiler {
 using SizesAndScalableFlags =
    std::pair<SmallVector<int64_t>, SmallVector<bool>>;

+using SizeAndScalableFlag = std::tuple<int64_t &, bool &>;


I'm about the make a refactor on this part and found this PR.

Is there a reason not to define a struct like:

struct ScalableTileSize { int64_t tileSize; bool scalableFlag; };

Having field name seems to be clearer and easier to access them?

What we really need is a pair of references (rather than a plain int64_t and a bool) to elements of 2 vectors. That’s because we do this in KernelDispatch.cpp

SmallVector<int64_t> commonVecTileSizes = parallelVecTileSizes; SmallVector<bool> commonVecScalableTileFlags = parallelVecScalableTileSizes;

Since one of the vectors contains bool, things are not as straightforward.

Btw, this alias is not really used (sorry, I was meant to remove it). But I appreciate that this code is not pretty regardless.

pzread · 2024-02-29T19:35:07Z

compiler/src/iree/compiler/Codegen/Common/TileSizeSelection.h

+    return std::pair<SmallVector<int64_t>, SmallVector<bool>>(sizes, flags);
+  }
+
+  ReferencePair operator[](size_t index) {


This allows users to get and store a reference, which seems to easily create a bug that keeps using the reference pair while the SizesAndScalableFlagsTuple has been destroyed. Especially if users don't understand the internal state of ReferencePair

Thanks for pointing this out!

Not sure there’s a way around it, TBH. Well, other than replacing (“pair of vectors”)

SmallVector<int64_t> commonVecTileSizes = parallelVecTileSizes; SmallVector<bool> commonVecScalableTileFlags = parallelVecScalableTileSizes;

with (“vector of pairs”):

SmallVector<ScalableTileSize> commonVecTileSizesAndFlags = parallelVecTileSizesAndFlags

That would be more intrusive, but would allow safer code. Having said that, ReferencePair is not intended for general use anyway. WDYT?

Yes, I was thinking about replacing it with vector of pair SmallVector<ScalableTileSize> everywhere in KernelDispatch.cpp. And in the end convert it into two separate vectors tileSizes and scalableTileFlags for the lowering_config attribute here:

https://github.com/openxla/iree/blob/9eff8615c2b2615db9f2ba6c79211452329aa872/compiler/src/iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenAttrs.h#L135-L152

So the new setOpConfigAndEntryPointFnTranslation will be:

inline LogicalResult setOpConfigAndEntryPointFnTranslation( mlir::FunctionOpInterface entryPointFn, Operation *op, SmallVector<SmallVector<ScalableTileSize>> scalableTileSizesList, IREE::Codegen::DispatchLoweringPassPipeline passPipeline, ArrayRef<int64_t> workgroupSize = {}, std::optional<int64_t> subgroupSize = {}, DictionaryAttr pipelineConfig = DictionaryAttr()) { ... }

Ideally we can also have [{tileSize, scalableFlag}, ...] on lowering_config attribute, but that will be a larger refactor and can be done separately. WDYT?

+1

I am really glad that you also see this as a better design.

One small nit, I would call it TileSizeAndScalableFlag rather than ScalableTileSize. The former indicates that 2 values are encapsulated - that would be always true. The latter suggests that tile sizes are always scalable - not always true - tile sizes can be both fixed width or scalable. Naming is hard.

I think we can have a C++ class which maintain all the information about tile sizes and scalable flags. Then we won't be bothered by pair and vector things. The class generates SmallVector<SmallVector<ScalableTileSize>> at the end, and we use it to set the strategy and configuration. It changes the whole flow, and it requires a broader review from folks. I proposed the idea internally (in AMD) and collected some feedback from @pzread and @dcaballe, but I did not get a chance to push it forward. I'll try to find some cycles and can share with you offline.

Please share when possible. If there's appetite for a larger refactor then perhaps that's the path that we should be taking. Is it something to raise in the mai-tai call?

It was in collecting initial feedback stage. I will revamp it and share with you.

MaheshRavishankar

Fly by comment. I see this is in Common/ . While we are cleaning this up can we move this to Common/CPU?
Thanks for the cleanup here though. I prefer structs to tuples as well. Less verbose.

hanhanW · 2024-03-05T21:08:38Z

compiler/src/iree/compiler/Codegen/Common/TileSizeSelection.h

+        flags(SmallVector<bool>(numElements, false)) {}
+
+  SizesAndScalableFlags get() {
+    return std::pair<SmallVector<int64_t>, SmallVector<bool>>(sizes, flags);


can this be std::make_pair(sizes, flags)?

hanhanW · 2024-03-05T21:09:41Z

compiler/src/iree/compiler/Codegen/Common/TileSizeSelection.h

+    return std::pair<SmallVector<int64_t>, SmallVector<bool>>(sizes, flags);
+  }
+
+  ReferencePair operator[](size_t index) {


I think we can have a C++ class which maintain all the information about tile sizes and scalable flags. Then we won't be bothered by pair and vector things. The class generates SmallVector<SmallVector<ScalableTileSize>> at the end, and we use it to set the strategy and configuration. It changes the whole flow, and it requires a broader review from folks. I proposed the idea internally (in AMD) and collected some feedback from @pzread and @dcaballe, but I did not get a chance to push it forward. I'll try to find some cycles and can share with you offline.

hanhanW · 2024-03-05T21:10:33Z

compiler/src/iree/compiler/Codegen/LLVMCPU/KernelDispatch.cpp

+  SizesAndScalableFlagsTuple commanVecTileSizesAndFlags = {
+      parallelVecTileSizes, parallelVecScalableTileSizes};


I think we can directly use the constructor here? E.g.,

Suggested change

SizesAndScalableFlagsTuple commanVecTileSizesAndFlags = {

parallelVecTileSizes, parallelVecScalableTileSizes};

SizesAndScalableFlagsTuple commanVecTileSizesAndFlags(parallelVecTileSizes, parallelVecScalableTileSizes);

hanhanW · 2024-03-05T21:14:00Z

compiler/src/iree/compiler/Codegen/Common/TileSizeSelection.h

+    }
+  };
+
+  SizesAndScalableFlagsTuple(SmallVector<int64_t> s, SmallVector<bool> f)


can this be ArrayRef?

hanhanW · 2024-03-05T21:18:18Z

compiler/src/iree/compiler/Codegen/Common/TileSizeSelection.h

+  SizesAndScalableFlagsTuple(size_t numElements)
+      : sizes(SmallVector<int64_t>(numElements, 0)),
+        flags(SmallVector<bool>(numElements, false)) {}


I think we don't need SmallVector<...> to construct it. Does the below snippet work?

Suggested change

SizesAndScalableFlagsTuple(size_t numElements)

: sizes(SmallVector<int64_t>(numElements, 0)),

flags(SmallVector<bool>(numElements, false)) {}

SizesAndScalableFlagsTuple(size_t numElements)

: sizes(numElements, 0),

flags(numElements, false) {}

banach-space requested review from dcaballe, hanhanW and MaheshRavishankar as code owners February 16, 2024 15:49

banach-space requested a review from MacDue February 16, 2024 15:49

banach-space force-pushed the andrzej/better_logic_for_tile_size_update branch from 8c63a4c to 52cb8af Compare February 17, 2024 08:45

pzread reviewed Feb 29, 2024

View reviewed changes

MaheshRavishankar reviewed Mar 2, 2024

View reviewed changes

hanhanW reviewed Mar 5, 2024

View reviewed changes

banach-space mentioned this pull request Mar 17, 2024

[CPU] Enable DT for SVE for linalg.matmul #16162

Open

3 tasks

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[CPU] Simplify how tile sizes are updated #16435

[CPU] Simplify how tile sizes are updated #16435

banach-space commented Feb 16, 2024

pzread Feb 29, 2024

banach-space Feb 29, 2024

pzread Feb 29, 2024 •

edited

Loading

banach-space Feb 29, 2024

pzread Feb 29, 2024 •

edited

Loading

banach-space Feb 29, 2024

hanhanW Mar 5, 2024

banach-space Mar 6, 2024

hanhanW Mar 6, 2024

MaheshRavishankar left a comment

hanhanW Mar 5, 2024

hanhanW Mar 5, 2024

hanhanW Mar 5, 2024

hanhanW Mar 5, 2024

hanhanW Mar 5, 2024

		SizesAndScalableFlagsTuple commanVecTileSizesAndFlags = {
		parallelVecTileSizes, parallelVecScalableTileSizes};

[CPU] Simplify how tile sizes are updated #16435

Are you sure you want to change the base?

[CPU] Simplify how tile sizes are updated #16435

Conversation

banach-space commented Feb 16, 2024

Choose a reason for hiding this comment

Choose a reason for hiding this comment

pzread Feb 29, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

pzread Feb 29, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

MaheshRavishankar left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

pzread Feb 29, 2024 •

edited

Loading

pzread Feb 29, 2024 •

edited

Loading