[flang][HLFIR] add skip_rebox option to hlfir.declare

[jeanPerier]

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.

[llvmbot]

@llvm/pr-subscribers-flang-fir-hlfir

Author: None (jeanPerier)

Changes

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.

---

Full diff: https://github.com/llvm/llvm-project/pull/162305.diff

5 Files Affected:

• (modified) flang/include/flang/Optimizer/HLFIR/HLFIROps.td (+7-1)
• (modified) flang/lib/Optimizer/HLFIR/IR/HLFIROps.cpp (+4-1)
• (modified) flang/lib/Optimizer/HLFIR/Transforms/ConvertToFIR.cpp (+8-3)
• (modified) flang/test/HLFIR/declare-codegen.fir (+9)
• (modified) flang/test/HLFIR/invalid.fir (+9)

diff --git a/flang/include/flang/Optimizer/HLFIR/HLFIROps.td b/flang/include/flang/Optimizer/HLFIR/HLFIROps.td
index 218435a44c24f..b7563a2f752f0 100644
--- a/flang/include/flang/Optimizer/HLFIR/HLFIROps.td
+++ b/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,7 +102,8 @@ 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);
 
@@ -106,6 +111,7 @@ def hlfir_DeclareOp
     $memref (`(` $shape^ `)`)? (`typeparams` $typeparams^)?
     (`dummy_scope` $dummy_scope^)?
     (`storage` `(` $storage^ `[` $storage_offset `]` `)`)?
+    (`skip_rebox` $skip_rebox^)?
     attr-dict `:` functional-type(operands, results)
   }];
 
diff --git a/flang/lib/Optimizer/HLFIR/IR/HLFIROps.cpp b/flang/lib/Optimizer/HLFIR/IR/HLFIROps.cpp
index 0cc65f939723e..1332dc57fb086 100644
--- a/flang/lib/Optimizer/HLFIR/IR/HLFIROps.cpp
+++ b/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 =
diff --git a/flang/lib/Optimizer/HLFIR/Transforms/ConvertToFIR.cpp b/flang/lib/Optimizer/HLFIR/Transforms/ConvertToFIR.cpp
index 8104e53920c27..6a57bf2ae6fec 100644
--- a/flang/lib/Optimizer/HLFIR/Transforms/ConvertToFIR.cpp
+++ b/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(
diff --git a/flang/test/HLFIR/declare-codegen.fir b/flang/test/HLFIR/declare-codegen.fir
index b3f0b73158603..04c2ddcece4a7 100644
--- a/flang/test/HLFIR/declare-codegen.fir
+++ b/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
diff --git a/flang/test/HLFIR/invalid.fir b/flang/test/HLFIR/invalid.fir
index 887113959429c..91e9ed6c2e4c4 100644
--- a/flang/test/HLFIR/invalid.fir
+++ b/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,?>'}}

[vzakhari]

LGTM! Thank you!