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

[Akash Banerjee 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<
----------------
TIFitis wrote:

All mapping is done on the host side, so we don't need any mapping information, or mapping related directives such as `target data/enter/exit` to be present on the device IR. We only need to make sure the mapped data members are made available along with the target code to be executed of course.

We can safely get rid of these as well, to further cleanup the device IR.

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