[Mlir-commits] [mlir] [mlir][Vector] Improve support for vector.extract(broadcast) (PR #116234)

Andrzej WarzyƄski llvmlistbot at llvm.org
Thu Nov 14 08:52:57 PST 2024


================
@@ -1680,6 +1676,16 @@ static Value foldExtractFromBroadcast(ExtractOp extractOp) {
           broadcastVecType.getShape().take_back(extractResultRank))
     return Value();
 
+  // The dim-1 broadcast -> ExtractOp folder requires in place operation
----------------
banach-space wrote:

>   // The dim-1 broadcast -> ExtractOp folder requires in place operation

Does it? Or is just how it's implemented today? 

Also, given that "dim-1" broadcasting is only referred to (and identified) further down, you may want to introduce the concept here. Otherwise, the flow is a bit confusing:

```cpp
// If this is dim-1 broadcasting, than bail out. Btw, dim-1 broadcasting has not been defined or detected yet.
  if (extractOp.hasDynamicPosition())
    return Value();
    
// Now lets define, identify and deal with dim-1 broadcasting.
  auto broadcastOp = cast<vector::BroadcastOp>(defOp);
  int64_t broadcastDstRank = broadcastOp.getResultVectorType().getRank()
  
  llvm::SetVector<int64_t> broadcastedUnitDims = 
```

Perhaps it should/could be like this:
```cpp
// Now lets define, identify and deal with dim-1 broadcasting.
  auto broadcastOp = cast<vector::BroadcastOp>(defOp);
  int64_t broadcastDstRank = broadcastOp.getResultVectorType().getRank()
  
  llvm::SetVector<int64_t> broadcastedUnitDims = 
  ...
// If this is dim-1 broadcasting and some indices are dynamic, bail out. ... (the rest of your comment explaining "why")
  if (extractOp.hasDynamicPosition())
    return Value();
```

https://github.com/llvm/llvm-project/pull/116234


More information about the Mlir-commits mailing list