[flang][openacc] map data operand results to symbols inside compute region

flang/include/flang/Lower/AbstractConverter.h

@@ -123,6 +123,10 @@ class AbstractConverter {
   virtual Fortran::lower::SymMap::StorageDesc
   getSymbolStorage(SymbolRef sym) = 0;
 
+  /// Return the Symbol Map used to map semantics::Symbol to their SSA
+  /// values in the generated MLIR.
+  virtual Fortran::lower::SymMap &getSymbolMap() = 0;
+
   /// Override lowering of expression with pre-lowered values.
   /// Associate mlir::Value to evaluate::Expr. All subsequent call to
   /// genExprXXX() will replace any occurrence of an overridden

flang/include/flang/Optimizer/HLFIR/HLFIROps.td

@@ -73,6 +73,10 @@ def hlfir_DeclareOp
         local lower bound values. It is intended to be used when generating FIR
         from HLFIR in order to avoid descriptor creation for simple entities.
 
+    The attribute skip_rebox can be set to indicate that the second and first
+    result are known to be the same descriptors (the input descriptor is known
+    to already have the correct attributes and lower bounds).
+
     Example:
 
     CHARACTER(n) :: c(10:n, 20:n)
@@ -98,14 +102,16 @@ def hlfir_DeclareOp
       DefaultValuedAttr<UI64Attr, "0">:$storage_offset,
       Builtin_StringAttr:$uniq_name,
       OptionalAttr<fir_FortranVariableFlagsAttr>:$fortran_attrs,
-      OptionalAttr<cuf_DataAttributeAttr>:$data_attr);
+      OptionalAttr<cuf_DataAttributeAttr>:$data_attr,
+      OptionalAttr<UnitAttr>:$skip_rebox);
 
   let results = (outs AnyFortranVariable, AnyRefOrBoxLike);
 
   let assemblyFormat = [{
     $memref (`(` $shape^ `)`)? (`typeparams` $typeparams^)?
     (`dummy_scope` $dummy_scope^)?
     (`storage` `(` $storage^ `[` $storage_offset `]` `)`)?
+    (`skip_rebox` $skip_rebox^)?
     attr-dict `:` functional-type(operands, results)
   }];
 

flang/lib/Lower/Bridge.cpp

@@ -643,6 +643,8 @@ class FirConverter : public Fortran::lower::AbstractConverter {
     return localSymbols.lookupStorage(sym);
   }
 
+  Fortran::lower::SymMap &getSymbolMap() override final { return localSymbols; }
+
   void
   overrideExprValues(const Fortran::lower::ExprToValueMap *map) override final {
     exprValueOverrides = map;

flang/lib/Optimizer/HLFIR/IR/HLFIROps.cpp

@@ -281,7 +281,7 @@ void hlfir::DeclareOp::build(mlir::OpBuilder &builder,
       getDeclareOutputTypes(inputType, hasExplicitLbs);
   build(builder, result, {hlfirVariableType, firVarType}, memref, shape,
         typeparams, dummy_scope, storage, storage_offset, nameAttr,
-        fortran_attrs, data_attr);
+        fortran_attrs, data_attr, /*skip_rebox=*/mlir::UnitAttr{});
 }
 
 llvm::LogicalResult hlfir::DeclareOp::verify() {
@@ -294,6 +294,9 @@ llvm::LogicalResult hlfir::DeclareOp::verify() {
     return emitOpError("first result type is inconsistent with variable "
                        "properties: expected ")
            << hlfirVariableType;
+  if (getSkipRebox() && !llvm::isa<fir::BaseBoxType>(getMemref().getType()))
+    return emitOpError(
+        "skip_rebox attribute must only be set when the input is a box");
   // The rest of the argument verification is done by the
   // FortranVariableInterface verifier.
   auto fortranVar =

flang/lib/Optimizer/HLFIR/Transforms/ConvertToFIR.cpp

@@ -23,6 +23,7 @@
 #include "flang/Optimizer/HLFIR/HLFIROps.h"
 #include "flang/Optimizer/HLFIR/Passes.h"
 #include "mlir/Transforms/DialectConversion.h"
+#include "llvm/ADT/SmallSet.h"
 
 namespace hlfir {
 #define GEN_PASS_DEF_CONVERTHLFIRTOFIR
@@ -312,10 +313,12 @@ class DeclareOpConversion : public mlir::OpRewritePattern<hlfir::DeclareOp> {
     // Propagate other attributes from hlfir.declare to fir.declare.
     // OpenACC's acc.declare is one example. Right now, the propagation
     // is verbatim.
-    mlir::NamedAttrList elidedAttrs =
-        mlir::NamedAttrList{firDeclareOp->getAttrs()};
+    llvm::SmallSet<llvm::StringRef, 8> elidedAttrs;
+    for (const mlir::NamedAttribute &firAttr : firDeclareOp->getAttrs())
+      elidedAttrs.insert(firAttr.getName());
+    elidedAttrs.insert(declareOp.getSkipReboxAttrName());
     for (const mlir::NamedAttribute &attr : declareOp->getAttrs())
-      if (!elidedAttrs.get(attr.getName()))
+      if (!elidedAttrs.contains(attr.getName()))
         firDeclareOp->setAttr(attr.getName(), attr.getValue());
 
     auto firBase = firDeclareOp.getResult();
@@ -328,6 +331,8 @@ class DeclareOpConversion : public mlir::OpRewritePattern<hlfir::DeclareOp> {
       auto genHlfirBox = [&]() -> mlir::Value {
         if (auto baseBoxType =
                 mlir::dyn_cast<fir::BaseBoxType>(firBase.getType())) {
+          if (declareOp.getSkipRebox())
+            return firBase;
           // Rebox so that lower bounds and attributes are correct.
           if (baseBoxType.isAssumedRank())
             return fir::ReboxAssumedRankOp::create(

flang/test/HLFIR/declare-codegen.fir

@@ -124,6 +124,15 @@ func.func @array_declare_box_2(%arg0: !fir.box<!fir.array<?x?xf32>>) {
 // CHECK:         %[[VAL_1:.*]] = fir.declare %[[VAL_0]] {uniq_name = "x"} : (!fir.box<!fir.array<?x?xf32>>) -> !fir.box<!fir.array<?x?xf32>>
 // CHECK:         %[[VAL_2:.*]] = fir.rebox %[[VAL_1]] : (!fir.box<!fir.array<?x?xf32>>) -> !fir.box<!fir.array<?x?xf32>>
 
+func.func @array_declare_box_3(%arg0: !fir.box<!fir.array<?x?xf32>>) -> !fir.box<!fir.array<?x?xf32>> {
+  %0:2 = hlfir.declare %arg0 skip_rebox {uniq_name = "x"} : (!fir.box<!fir.array<?x?xf32>>) -> (!fir.box<!fir.array<?x?xf32>>, !fir.box<!fir.array<?x?xf32>>)
+  return %0#0 : !fir.box<!fir.array<?x?xf32>>
+}
+// CHECK-LABEL:   func.func @array_declare_box_3(
+// CHECK-SAME:    %[[VAL_0:.*]]: !fir.box<!fir.array<?x?xf32>>) -> !fir.box<!fir.array<?x?xf32>> {
+// CHECK:         %[[VAL_1:.*]] = fir.declare %[[VAL_0]] {uniq_name = "x"} : (!fir.box<!fir.array<?x?xf32>>) -> !fir.box<!fir.array<?x?xf32>>
+// CHECK:         return %[[VAL_1]]
+
 func.func @array_declare_char_box(%arg0: !fir.box<!fir.array<?x?x!fir.char<1,?>>>) {
   %0:2 = hlfir.declare %arg0 {uniq_name = "x"} : (!fir.box<!fir.array<?x?x!fir.char<1,?>>>) -> (!fir.box<!fir.array<?x?x!fir.char<1,?>>>, !fir.box<!fir.array<?x?x!fir.char<1,?>>>)
   return

flang/test/HLFIR/invalid.fir

@@ -36,6 +36,15 @@ func.func @bad_array_declare(%arg0: !fir.ref<!fir.array<?x?xf32>>) {
   return
 }
 
+// -----
+func.func @bad_declare_skip_rebox(%arg0: !fir.ref<f32>) {
+  // expected-error@+1 {{'hlfir.declare' op skip_rebox attribute must only be set when the input is a box}}
+  %0:2 = hlfir.declare %arg0 skip_rebox {uniq_name = "x"} : (!fir.ref<f32>) -> (!fir.ref<f32>, !fir.ref<f32>)
+  return
+}
+
+// -----
+
 // -----
 func.func @bad_assign_scalar_character(%arg0: !fir.boxchar<1>, %arg1: !fir.char<1,?>) {
   // expected-error@+1 {{'hlfir.assign' op operand #0 must be any Fortran value or variable type, but got '!fir.char<1,?>'}}