[flang][HLFIR] add skip_rebox option to hlfir.declare (#162305)

hlfir.declare with a fir.box input always lead to a fir.rebox being
created to ensure the lower bounds and attributes are set correctly in
the local descriptor for the entity. This is really needed for
hlfir.declare using fir.box function argument that do not come with any
guarantees with regards to the lower bounds.

Sometimes however, this fir.rebox just adds a lot of noise in the SSA
chain, especially at the LLVM level and it is known that the input
descriptor is already correct.

I am making this patch in the context of OpenACC where I want to remap
the variables inside the compute region, creating a fir.rebox on the
way. This fir.rebox cannot be optimized away by FIR because of the
OpenACC ops in the SSA chain.

This patch adds a flag to indicate the the fir.box is known to have the
correct lower bounds and attributes so that it can have a simpler code
generation to FIR.

Author: jeanPerier

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/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,?>'}}