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

[skatrak]

Summary

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

[kiranchandramohan]

Nice find. LGTM.

[skatrak]

An alternative approach to solve this issue would be to mark the omp.target operation with the IsolatedFromAbove trait. I decided against it to allow implicit mapping of outside values to target regions, but this would be a simple change if that's a preferable approach.

[kiranchandramohan]

In D148349#4268762, @skatrak wrote:

An alternative approach to solve this issue would be to mark the omp.target operation with the IsolatedFromAbove trait. I decided against it to allow implicit mapping of outside values to target regions, but this would be a simple change if that's a preferable approach.

Yes, Isolated from Above would mean there are no live-ins into the region. So this looks good to me.

[skatrak]

In D148349#4268807, @kiranchandramohan wrote:

In D148349#4268762, @skatrak wrote:

An alternative approach to solve this issue would be to mark the omp.target operation with the IsolatedFromAbove trait. I decided against it to allow implicit mapping of outside values to target regions, but this would be a simple change if that's a preferable approach.

Yes, Isolated from Above would mean there are no live-ins into the region. So this looks good to me.

Thank you for the quick review, Kiran. I'll wait until Monday to land this patch, in case someone else has any further suggestions.

mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp

 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/OpenMP/OpenMPDialect.h"
 #include "mlir/Dialect/LLVMIR/LLVMTypes.h"
 #include "mlir/IR/Attributes.h"
 #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"
 #include "llvm/ADT/StringExtras.h"
 #include "llvm/ADT/StringRef.h"
 #include "llvm/ADT/TypeSwitch.h"
 #include "llvm/Frontend/OpenMP/OMPConstants.h"
 #include <cstddef>

 /// Model for pointer-like types that already provide a `getElementType` method.
 template <typename T>
 struct PointerLikeModel
     : public PointerLikeType::ExternalModel<PointerLikeModel<T>, T> {
   Type getElementType(Type pointer) const {
     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() {
   addOperations<
 #define GET_OP_LIST
 #include "mlir/Dialect/OpenMP/OpenMPOps.cpp.inc"
       >();
   addAttributes<
 #define GET_ATTRDEF_LIST
 #include "mlir/Dialect/OpenMP/OpenMPOpsAttributes.cpp.inc"
       >();
 
+  addInterface<OpenMPDialectFoldInterface>();
   LLVM::LLVMPointerType::attachInterface<
       PointerLikeModel<LLVM::LLVMPointerType>>(*getContext());
   MemRefType::attachInterface<PointerLikeModel<MemRefType>>(*getContext());
   LLVM::LLVMPointerType::attachInterface<
       PointerLikeModel<LLVM::LLVMPointerType>>(*getContext());
   mlir::ModuleOp::attachInterface<mlir::omp::OffloadModuleDefaultModel>(
       *getContext());
 }

mlir/test/Dialect/OpenMP/canonicalize.mlir

func.func @parallel_maybe_side_effects(%a: i32, %b: i32) {
   return
 }
 
 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