[Mlir-commits] [mlir] [mlir][ArmNeon] Update `LowerContractionToSMMLAPattern` to support proper unrolling for k dimension (PR #88591)
Diego Caballero
llvmlistbot at llvm.org
Tue Apr 16 08:57:06 PDT 2024
================
@@ -133,8 +133,12 @@ class LowerContractionToSMMLAPattern
smmlaShape.insert(smmlaShape.begin(), isVecmat ? 1 : 2);
loopOrder.push_back(2);
}
+
+ // Keep track of the previous accumulator when tiling over K
+ Value kAcc;
for (SmallVector<int64_t> offsets :
StaticTileOffsetRange(unrolledSize, smmlaShape, loopOrder)) {
+ auto kTileIndex = offsets[offsets.size() - 1] / 8;
----------------
dcaballe wrote:
if we only need to know if it's the first iteration of `k`, we can avoid the division and do something like:
```
bool initializeResult = offsets.back() == 0;
```
and then write the logic below using this boolean flag.
https://github.com/llvm/llvm-project/pull/88591
More information about the Mlir-commits
mailing list