[flang-commits] [flang] [Flang][MLIR] Teach AddAliasTags getFuncArgName to deal with omp::TargetOp's (PR #92036)

[via flang-commits]

https://github.com/agozillon created https://github.com/llvm/llvm-project/pull/92036

This PR attempts to fix the name retrieval of getFuncArgName inside of AddAliasTags when used in conjunction with an omp::TargetOp. While omp::TargetOp is function like it does not implement the FunctionOpInterface, so on the occasions where this pass invokes getFuncArgName on it, it will crash.

This PR attempts to address that by adding a special case for TargetOp where we find the map operand corresponding to the block argument and return the name attribute tied to the map operand.

This issue was found in a downstream test case that unfortunately won't run upstream yet as far as I am aware, the below is a variation of the offending block of code (attempted to reduce it a bit..):

```
!$omp target teams distribute parallel do &
!$omp collapse(2) private(i,ib,ic,i1)
    do i1=1,3
        do ic=1,3
            do ib=1,3
                do i=1,3
                    xgnn(1,1,1,1) = 1
                end do
            end do
        end do
    end do
!$omp end target teams distribute parallel do
```

When -O3 (or anything larger than default) is set, the AddAliasTags pass will be activated and when we encounter a construct like the above in a function or subroutine we end up ICE'ng the compiler.

>From edceddecc00df3326e995e477c8145793db4f522 Mon Sep 17 00:00:00 2001
From: agozillon <Andrew.Gozillon at amd.com>
Date: Mon, 13 May 2024 16:16:18 -0500
Subject: [PATCH] [Flang][MLIR] Teach AddAliasTags getFuncArgName to deal with
 omp::TargetOp's

This PR attempts to fix the name retrieval of getFuncArgName inside of AddAliasTags
when used in conjunction with an omp::TargetOp. While omp::TargetOp is function like
it does not implement the FunctionOpInterface, so on the occasions where this pass
invokes getFuncArgName on it, it will crash.

This PR attempts to address that by adding a special case for TargetOp where we find
the map operand corresponding to the block argument and return the name attribute
tied to the map operand.

This issue was found in a downstream test case that unfortunately won't run upstream
yet as far as I am aware, the below is the offending block of code:

!$omp target teams distribute parallel do &
!$omp collapse(2) private(i,ib,ic,i1)
    do i1=1,3
        do ic=1,3
            do ib=1,3
                do i=1,3
                    xgnn(1,1,1,1) = 1
                end do
            end do
        end do
    end do
!$omp end target teams distribute parallel do

When -O3 (or anything larger than default) is set, the AddAliasTags pass will be activated
and when we encounter a construct like the above in a function or subroutine we end up
ICE'ng the compiler.
---
 flang/lib/Optimizer/Transforms/AddAliasTags.cpp | 17 +++++++++++++++--
 1 file changed, 15 insertions(+), 2 deletions(-)

diff --git a/flang/lib/Optimizer/Transforms/AddAliasTags.cpp b/flang/lib/Optimizer/Transforms/AddAliasTags.cpp
index 684aa4462915e..026dfe93d6f11 100644
--- a/flang/lib/Optimizer/Transforms/AddAliasTags.cpp
+++ b/flang/lib/Optimizer/Transforms/AddAliasTags.cpp
@@ -105,8 +105,21 @@ static std::string getFuncArgName(mlir::Value arg) {
          "arg is a function argument");
   mlir::FunctionOpInterface func = mlir::dyn_cast<mlir::FunctionOpInterface>(
       blockArg.getOwner()->getParentOp());
-  mlir::StringAttr attr = func.getArgAttrOfType<mlir::StringAttr>(
-      blockArg.getArgNumber(), "fir.bindc_name");
+  mlir::StringAttr attr;
+  if (func) {
+    attr = func.getArgAttrOfType<mlir::StringAttr>(blockArg.getArgNumber(),
+                                                   "fir.bindc_name");
+  } else {
+    if (auto targetOp = mlir::dyn_cast<mlir::omp::TargetOp>(
+            blockArg.getOwner()->getParentOp())) {
+      if (auto mapOp = mlir::dyn_cast<mlir::omp::MapInfoOp>(
+              targetOp.getMapOperands()[blockArg.getArgNumber()]
+                  .getDefiningOp())) {
+        attr = mapOp.getNameAttr();
+      }
+    }
+  }
+
   if (!attr)
     return "";
   return attr.str();