[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