extend CanonicalizeReduction

Hardcode84 · Mar 19, 2021 · Mar 13, 2021 · Mar 13, 2021 · Mar 17, 2021 · Mar 17, 2021
commit 94cdca21ecb9541899d86c1c32bfc1e11caeec76
diff --git a/mlir-compiler/plier/CMakeLists.txt b/mlir-compiler/plier/CMakeLists.txt
@@ -22,10 +22,11 @@ set(SOURCES_LIST
     src/rewrites/cse.cpp
     src/rewrites/force_inline.cpp
     src/rewrites/index_type_propagation.cpp
-    include/plier/rewrites/if_rewrites.cpp
+    src/rewrites/if_rewrites.cpp
     src/rewrites/loop_rewrites.cpp
     src/rewrites/promote_to_parallel.cpp
     src/rewrites/type_conversion.cpp
+    src/transforms/block_utils.cpp
     src/transforms/cast_utils.cpp
     src/transforms/const_utils.cpp
     src/transforms/func_utils.cpp
@@ -49,6 +50,7 @@ set(HEADERS_LIST
     include/plier/rewrites/loop_rewrites.hpp
     include/plier/rewrites/promote_to_parallel.hpp
     include/plier/rewrites/type_conversion.hpp
+    include/plier/transforms/block_utils.hpp
     include/plier/transforms/cast_utils.hpp
     include/plier/transforms/const_utils.hpp
     include/plier/transforms/func_utils.hpp

diff --git a/mlir-compiler/plier/include/plier/transforms/block_utils.hpp b/mlir-compiler/plier/include/plier/transforms/block_utils.hpp
@@ -0,0 +1,20 @@
+#pragma once
+
+namespace mlir
+{
+class Operation;
+}
+
+namespace plier
+{
+
+enum class OpRelation
+{
+    Before,
+    After,
+    In,
+    Unknown
+};
+
+OpRelation relativeTo(mlir::Operation* op, mlir::Operation* relativeTo);
+}
diff --git a/mlir-compiler/plier/src/rewrites/canonicalize_reductions.cpp b/mlir-compiler/plier/src/rewrites/canonicalize_reductions.cpp
@@ -4,18 +4,39 @@
 #include <mlir/Dialect/StandardOps/IR/Ops.h>
 #include <mlir/IR/BlockAndValueMapping.h>
 
+#include "plier/transforms/block_utils.hpp"
+
 namespace
 {
 bool checkMemrefType(mlir::Value value)
 {
     if (auto type = value.getType().dyn_cast<mlir::MemRefType>())
     {
-        auto shape = type.getShape();
-        return shape.empty() || (1 == shape.size() && 1 == shape[0]);
+//        auto shape = type.getShape();
+//        return shape.empty() || (1 == shape.size() && 1 == shape[0]);
+        return true;
     }
     return false;
 }
 
+bool isOutsideBlock(mlir::ValueRange values, mlir::Block& block)
+{
+    auto blockArgs = block.getArguments();
+    for (auto val : values)
+    {
+        if (llvm::find(blockArgs, val) != blockArgs.end())
+        {
+            return false;
+        }
+        auto op = val.getDefiningOp();
+        if (op && block.findAncestorOpInBlock(*op))
+        {
+            return false;
+        }
+    }
+    return true;
+}
+
 bool checkForPotentialAliases(mlir::Value value, mlir::Operation* parent)
 {
     assert(parent->getRegions().size() == 1);
@@ -31,13 +52,53 @@ bool checkForPotentialAliases(mlir::Value value, mlir::Operation* parent)
     {
         return false;
     }
+
+    mlir::LoadOp load;
+    mlir::StoreOp store;
+    auto& parentBlock = parent->getRegions().front().front();
     for (auto user : value.getUsers())
     {
         if (mlir::isa<mlir::ViewLikeOpInterface>(user))
         {
             // TODO: very conservative
             return false;
         }
+        auto relation = plier::relativeTo(user, parent);
+        if (plier::OpRelation::Unknown == relation)
+        {
+            return false;
+        }
+        if (plier::OpRelation::In == relation)
+        {
+            if (auto effects = mlir::dyn_cast_or_null<mlir::MemoryEffectOpInterface>(user))
+            {
+                if (user->getBlock() != &parentBlock)
+                {
+                    return false;
+                }
+                if (effects.hasEffect<mlir::MemoryEffects::Read>())
+                {
+                    if (load || !mlir::isa<mlir::LoadOp>(user))
+                    {
+                        return false;
+                    }
+                    load = mlir::cast<mlir::LoadOp>(user);
+                }
+                if (effects.hasEffect<mlir::MemoryEffects::Write>())
+                {
+                    if (store || !mlir::isa<mlir::StoreOp>(user))
+                    {
+                        return false;
+                    }
+                    store = mlir::cast<mlir::StoreOp>(user);
+                }
+            }
+        }
+    }
+    if (!load || !store || !load->isBeforeInBlock(store) ||
+        load.indices() != store.indices() || !isOutsideBlock(load.indices(), parentBlock))
+    {
+        return false;
     }
     return true;
 }

diff --git a/...er/include/plier/rewrites/if_rewrites.cpp → ...mpiler/plier/src/rewrites/if_rewrites.cpp b/...er/include/plier/rewrites/if_rewrites.cpp → ...mpiler/plier/src/rewrites/if_rewrites.cpp
diff --git a/mlir-compiler/plier/src/transforms/block_utils.cpp b/mlir-compiler/plier/src/transforms/block_utils.cpp
@@ -0,0 +1,64 @@
+#include "plier/transforms/block_utils.hpp"
+
+#include <mlir/IR/Operation.h>
+
+namespace
+{
+auto collectParentOps(mlir::Operation* op)
+{
+    llvm::SmallVector<mlir::Operation*> ops;
+    while (true)
+    {
+        assert(op);
+        ops.emplace_back(op);
+        auto parent = op->getParentOp();
+        if (!parent)
+        {
+            break;
+        }
+        op = parent;
+    }
+    return ops;
+}
+}
+
+plier::OpRelation plier::relativeTo(mlir::Operation* op, mlir::Operation* relativeTo)
+{
+    assert(op);
+    assert(relativeTo);
+
+    for (auto& reg : relativeTo->getRegions())
+    {
+        for (auto& block : reg)
+        {
+            if (block.findAncestorOpInBlock(*op))
+            {
+                return OpRelation::In;
+            }
+        }
+    }
+
+    auto ops1 = collectParentOps(op);
+    auto ops2 = collectParentOps(relativeTo);
+
+    for (auto op1 : ops1)
+    {
+        assert(op1);
+        for (auto op2 : ops2)
+        {
+            assert(op2);
+            if (op1->getBlock() == op1->getBlock())
+            {
+                if (op1->isBeforeInBlock(op1))
+                {
+                    return OpRelation::Before;
+                }
+                else
+                {
+                    return OpRelation::After;
+                }
+            }
+        }
+    }
+    return OpRelation::Unknown;
+}