[MLIR][OpenMP] Fix the assertion failure for VariableCaptureKind::ByCopy

[TIFitis]

This patch removes some of the code instrumentation that was put in place to support non reference values for map operands.

#72444 removes such values from being mapped and thus, we no longer require to handle them.

Depends on #72444.

[llvmbot]

@llvm/pr-subscribers-mlir-llvm
@llvm/pr-subscribers-flang-openmp
@llvm/pr-subscribers-mlir-openmp

@llvm/pr-subscribers-mlir

Author: Akash Banerjee (TIFitis)

Changes

---

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

1 Files Affected:

• (modified) mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp (+3-16)

diff --git a/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp b/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
index ae13a745196c42d..c1bc9d42bed75e6 100644
--- a/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
+++ b/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
@@ -2202,16 +2202,11 @@ createDeviceArgumentAccessor(MapInfoData &mapData, llvm::Argument &arg,
       ompBuilder.M.getDataLayout().getProgramAddressSpace();
 
   // Create the alloca for the argument the current point.
-  llvm::Value *v =
-      builder.CreateAlloca(arg.getType()->isPointerTy()
-                               ? arg.getType()
-                               : llvm::Type::getInt64Ty(builder.getContext()),
-                           ompBuilder.M.getDataLayout().getAllocaAddrSpace());
+  llvm::Value *v = builder.CreateAlloca(arg.getType(), allocaAS);
 
-  if (allocaAS != defaultAS && arg.getType()->isPointerTy()) {
+  if (allocaAS != defaultAS && arg.getType()->isPointerTy())
     v = builder.CreatePointerBitCastOrAddrSpaceCast(
         v, arg.getType()->getPointerTo(defaultAS));
-  }
 
   builder.CreateStore(&arg, v);
 
@@ -2219,15 +2214,7 @@ createDeviceArgumentAccessor(MapInfoData &mapData, llvm::Argument &arg,
 
   switch (capture) {
   case mlir::omp::VariableCaptureKind::ByCopy: {
-    // RHS of || aims to ignore conversions like int -> uint, but further
-    // extension of this path must be implemented for the moment it'll fall
-    // through to the assert.
-    if (inputType->isPointerTy() || v->getType() == inputType->getPointerTo()) {
-      retVal = v;
-      return builder.saveIP();
-    }
-
-    assert(false && "Currently unsupported OMPTargetVarCaptureByCopy Type");
+    retVal = v;
     break;
   }
   case mlir::omp::VariableCaptureKind::ByRef: {

[agozillon]

This seems to be only one part of the problem when mapping the captured index variables unfortunately.

Take the following example that I believe will crash during verification with these changes:

program main
  INTEGER :: sp(10) = (/0,0,0,0,0,0,0,0,0,0/)
!$omp target map(tofrom:sp)
  do i = 1, 10
      sp(i) = i;
  end do
!$omp end target
end program

The verifier will complain due to the for loop in this case (if I haven't butchered this example at least), because of incompatible PHI node branches (as the index is now a pointer with this patch, but it successfully gets passed the previous hurdle).

So, there's two possible fixes that I've found to this:

1. We add a another CreateAlignedLoad to the generated alloca as a seperate path inside of the ByCopy portion of createDeviceArgumentAccessor, when the input (or ArgType, InputType is the underlying element type, but ArgType will be an opaque pointer in cases where the value isn't an index and is actually a pointer, so the CreateAlignedLoad type would be slightly different from the ByRef case) for example something like:

  case mlir::omp::VariableCaptureKind::ByCopy: {
    if (arg.getType()->isPointerTy())
       retVal = v;
    else 
      retVal = builder.CreateAlignedLoad(
        arg.getType(), v,
        ompBuilder.M.getDataLayout().getPrefTypeAlign(v->getType()));
    break;
  }

Or the second simpler fix, where we change the following line inside of the argAccessorCB lambda from:

if (!ompBuilder->Config.isTargetDevice()) {
to
if (!ompBuilder->Config.isTargetDevice() || !arg.getType()->isPointerTy()) {

And simply return the argument without any extra instructions for accessing for now, we might end up with some weird write back behavior of the index to the host, at the end of the loop despite it being a scalar, although, I'm not entirely sure without some testing. But it can be ironed out if we maintain the usage of non-pointer types long term, and for the moment we tend to pick these cases up as implicit captures that I don't think users directly access. I do also think the runtime has issues passing types that are bigger than i64's from my very basic understanding of some comments I've read, so we might need to be careful about that occurring but it seems quite unlikely (it could also be fixed and the comment I read is just out of date).

If you do opt for two, you can keep the useful simplification of the createDeviceArgumentAccessor function you have in the PR currently and perhaps also remove the other isPointerTy() checks for the time being as it'd now only handle pointers! The isPointerTy checks are there for the cases where a variable is ByCopy+i64 type which Clang seems to do in very special cases (and stores types other than i64 in it I think), but it does need fleshed out a fair bit more for this case I think, if it's something we want to carry over. I'd also remove the inputType variable if it's no longer used!

However, these are just two suggestions to the secondary issue I found, you don't need to go with either, there might be better solutions!

[kiranchandramohan]

Please add a test and a suitable description.

[TIFitis]

However, these are just two suggestions to the secondary issue I found, you don't need to go with either, there might be better solutions!

#72444 should fix scenarios.

[TIFitis]

Please add a test and a suitable description.

I've updated the description.

We already have tests for them in openmp/libomptarget/test/offloading/fortran. The build bots are missing proper configuration to run them for now, so we didn't catch that they were failing.

The lowering tests don't show these errors because with the -flang-deprecated-no-hlfir flag the error goes away.

[TIFitis]

I plan on merging this patch tomorrow along with #72444. Please let me know if you have any concerns, or if you'd like some more time.