From 683e17fa35dd30236b74fa3251d6b716e9cf757d Mon Sep 17 00:00:00 2001
From: Thomas Raoux <thomasraoux@google.com>
Date: Thu, 7 Apr 2022 04:03:57 +0000
Subject: [PATCH] Add attribute control for splitK

---
 .../Dialect/Flow/Transforms/SplitReduction.cpp       | 12 ++++++++----
 1 file changed, 8 insertions(+), 4 deletions(-)

diff --git a/iree/compiler/Dialect/Flow/Transforms/SplitReduction.cpp b/iree/compiler/Dialect/Flow/Transforms/SplitReduction.cpp
index d9667ddb9099..0a0386dc35ba 100644
--- a/iree/compiler/Dialect/Flow/Transforms/SplitReduction.cpp
+++ b/iree/compiler/Dialect/Flow/Transforms/SplitReduction.cpp
@@ -22,6 +22,8 @@ namespace iree_compiler {
 namespace IREE {
 namespace Flow {
 
+static const char kSplitKAttr[] = "iree_flow_split_k";
+
 // TODO(thomasraoux): Move to attributes.
 static llvm::cl::opt<int64_t> splitReductionRatio(
     "iree-flow-split-matmul-reduction", llvm::cl::desc("split ratio"),
@@ -75,16 +77,18 @@ struct SplitReductionPass : public SplitReductionBase<SplitReductionPass> {
   }
 
   void runOnOperation() override {
-    if (splitReductionRatio <= 1) return;
-
     RewritePatternSet patterns(&getContext());
     patterns.add<LinalgSplitReduction>(
         &getContext(),
-        [&](linalg::LinalgOp op) {
+        [](linalg::LinalgOp op) {
+          int64_t ratio = splitReductionRatio;
+          if (auto attr = op->getAttrOfType<IntegerAttr>(kSplitKAttr))
+            ratio = attr.getInt();
+          if (ratio <= 1) return std::make_pair(int64_t(0), 0);
           // For matmul make the new parallel dimension first so that it looks
           // like a batch_matmul and can follow the same codegen.
           if (isa<linalg::MatmulOp>(op))
-            return std::make_pair(int64_t(splitReductionRatio), 0);
+            return std::make_pair(int64_t(ratio), 0);
           // Currently disable spliting reduction for non-matmul op. This will
           // get enabled after once tests are ready.
           return std::make_pair(int64_t(0), 0);