[Mlir-commits] [mlir] a26ba87 - [OpenMP][MLIR] Prevent constant hoisting out of target regions

[Sergio Afonso]

Author: Sergio Afonso
Date: 2023-04-17T10:47:49+01:00
New Revision: a26ba8724f734fc85ef6378e3f82aef790e6e543

URL: https://github.com/llvm/llvm-project/commit/a26ba8724f734fc85ef6378e3f82aef790e6e543
DIFF: https://github.com/llvm/llvm-project/commit/a26ba8724f734fc85ef6378e3f82aef790e6e543.diff

LOG: [OpenMP][MLIR] Prevent constant hoisting out of target regions

This patch prevents constant operations defined inside `omp.target` from being
hoisted out and into their parent functions by canonicalization passes.

Differential Revision: https://reviews.llvm.org/D148349

Added: 
    

Modified: 
    mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
    mlir/test/Dialect/OpenMP/canonicalize.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp b/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
index 799ae2e0dc056..78c824c3deba1 100644
--- a/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
+++ b/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
@@ -16,6 +16,7 @@
 #include "mlir/IR/DialectImplementation.h"
 #include "mlir/IR/OpImplementation.h"
 #include "mlir/IR/OperationSupport.h"
+#include "mlir/Interfaces/FoldInterfaces.h"
 
 #include "llvm/ADT/BitVector.h"
 #include "llvm/ADT/SmallString.h"
@@ -43,6 +44,15 @@ struct PointerLikeModel
     return pointer.cast<T>().getElementType();
   }
 };
+
+struct OpenMPDialectFoldInterface : public DialectFoldInterface {
+  using DialectFoldInterface::DialectFoldInterface;
+
+  bool shouldMaterializeInto(Region *region) const final {
+    // Avoid folding constants across target regions
+    return isa<TargetOp>(region->getParentOp());
+  }
+};
 } // namespace
 
 void OpenMPDialect::initialize() {
@@ -55,6 +65,7 @@ void OpenMPDialect::initialize() {
 #include "mlir/Dialect/OpenMP/OpenMPOpsAttributes.cpp.inc"
       >();
 
+  addInterface<OpenMPDialectFoldInterface>();
   LLVM::LLVMPointerType::attachInterface<
       PointerLikeModel<LLVM::LLVMPointerType>>(*getContext());
   MemRefType::attachInterface<PointerLikeModel<MemRefType>>(*getContext());

diff  --git a/mlir/test/Dialect/OpenMP/canonicalize.mlir b/mlir/test/Dialect/OpenMP/canonicalize.mlir
index c5d18f39d1d58..68f5bacb1def1 100644
--- a/mlir/test/Dialect/OpenMP/canonicalize.mlir
+++ b/mlir/test/Dialect/OpenMP/canonicalize.mlir
@@ -126,3 +126,19 @@ func.func private @foo() -> ()
 // CHECK: omp.parallel
 // CHECK: func.call @foo() : () -> ()
 // CHECK: omp.terminator
+
+// -----
+
+func.func @constant_hoisting_target(%x : !llvm.ptr<i32>) {
+  omp.target {
+    %c1 = arith.constant 10 : i32
+    llvm.store %c1, %x : i32, !llvm.ptr<i32>
+    omp.terminator
+  }
+  return
+}
+
+// CHECK-LABEL: func.func @constant_hoisting_target
+// CHECK-NOT: arith.constant
+// CHECK: omp.target
+// CHECK-NEXT: arith.constant