[flang-commits] [flang] [Flang][OpenMP] Minimize host ops remaining in device compilation (PR #137200)

[Sergio Afonso via flang-commits]

================
@@ -94,12 +195,359 @@ class FunctionFilteringPass
           funcOp.erase();
           return WalkResult::skip();
         }
+
+        if (failed(rewriteHostRegion(funcOp.getRegion()))) {
+          funcOp.emitOpError() << "could not be rewritten for target device";
+          return WalkResult::interrupt();
+        }
+
         if (declareTargetOp)
           declareTargetOp.setDeclareTarget(declareType,
                                            omp::DeclareTargetCaptureClause::to);
       }
       return WalkResult::advance();
     });
   }
+
+private:
+  /// Rewrite the given host device region belonging to a function that contains
+  /// \c omp.target operations, to remove host-only operations that are not used
+  /// by device codegen.
+  ///
+  /// It is based on the expected form of the MLIR module as produced by Flang
+  /// lowering and it performs the following mutations:
+  ///   - Replace all values returned by the function with \c fir.undefined.
+  ///   - Operations taking map-like clauses (e.g. \c omp.target,
+  ///     \c omp.target_data, etc) are moved to the end of the function. If they
+  ///     are nested inside of any other operations, they are hoisted out of
+  ///     them. If the region belongs to \c omp.target_data, these operations
+  ///     are hoisted to its top level, rather than to the parent function.
+  ///   - \c device, \c if and \c depend clauses are removed from these target
+  ///     functions. Values initializing other clauses are either replaced by
+  ///     placeholders as follows:
+  ///     - Values defined by block arguments are replaced by placeholders only
+  ///       if they are not attached to \c func.func or \c omp.target_data
+  ///       operations. In that case, they are kept unmodified.
+  ///     - \c arith.constant and \c fir.address_of are maintained.
+  ///     - Other values are replaced by a combination of an \c fir.alloca for a
+  ///       single bit and an \c fir.convert to the original type of the value.
+  ///       This can be done because the code eventually generated for these
+  ///       operations will be discarded, as they aren't runnable by the target
+  ///       device.
+  ///   - \c omp.map.info operations associated to these target regions are
+  ///     preserved. These are moved above all \c omp.target and sorted to
+  ///     satisfy dependencies among them.
+  ///   - \c bounds arguments are removed from \c omp.map.info operations.
+  ///   - \c var_ptr and \c var_ptr_ptr arguments of \c omp.map.info are
+  ///     handled as follows:
+  ///     - \c var_ptr_ptr is expected to be defined by a \c fir.box_offset
+  ///       operation which is preserved. Otherwise, the pass will fail.
+  ///     - \c var_ptr can be defined by an \c hlfir.declare which is also
+  ///       preserved. Its \c memref argument is replaced by a placeholder or
+  ///       maintained similarly to non-map clauses of target operations
+  ///       described above. If it has \c shape or \c typeparams arguments, they
+  ///       are replaced by applicable constants. \c dummy_scope arguments
+  ///       are discarded.
+  ///   - Every other operation not located inside of an \c omp.target is
+  ///     removed.
+  ///   - Whenever a value or operation that would otherwise be replaced with a
+  ///     placeholder is defined outside of the region being rewritten, it is
+  ///     added to the \c parentOpRewrites or \c parentValRewrites output
+  ///     argument, to be later handled by the caller. This is only intended to
+  ///     properly support nested \c omp.target_data and \c omp.target placed
+  ///     inside of \c omp.target_data. When called for the main function, these
+  ///     output arguments must not be set.
+  LogicalResult
+  rewriteHostRegion(Region &region,
+                    llvm::SetVector<Operation *> *parentOpRewrites = nullptr,
+                    llvm::SetVector<Value> *parentValRewrites = nullptr) {
+    // Extract parent op information.
+    auto [funcOp, targetDataOp] = [&region]() {
+      Operation *parent = region.getParentOp();
+      return std::make_tuple(dyn_cast<func::FuncOp>(parent),
+                             dyn_cast<omp::TargetDataOp>(parent));
+    }();
+    assert((bool)funcOp != (bool)targetDataOp &&
+           "region must be defined by either func.func or omp.target_data");
+    assert((bool)parentOpRewrites == (bool)targetDataOp &&
+           (bool)parentValRewrites == (bool)targetDataOp &&
+           "parent rewrites must be passed iff rewriting omp.target_data");
+
+    // Collect operations that have mapping information associated to them.
+    llvm::SmallVector<
----------------
skatrak wrote:

Thank you Akash for these suggestions, I was able to remove operations related to `target data`, `target enter data`, `target exit data` and `target update` as well. This also helped make things a little simpler as well, since a lot of the complexity was to keep `omp.target_data` ops around. Let me know if these changes work for you.

https://github.com/llvm/llvm-project/pull/137200