[openacc] Update acc.loop to expose data operands (llvm#70954)

The compute and data constructs implement getNumDataOperands and
getDataOperand. The acc.loop operation similarly has multiple data
operands - thus it makes sense to expose them the same way.

For loop, only private and reduction operands are exposed this way.
Technically, acc.loop also holds cache operands - but these are hints
not a data attribute.

Author: razvanlupusoru

flang/lib/Lower/OpenACC.cpp

@@ -1592,9 +1592,9 @@ createLoopOp(Fortran::lower::AbstractConverter &converter,
   addOperand(operands, operandSegments, workerNum);
   addOperand(operands, operandSegments, vectorNum);
   addOperands(operands, operandSegments, tileOperands);
+  addOperands(operands, operandSegments, cacheOperands);
   addOperands(operands, operandSegments, privateOperands);
   addOperands(operands, operandSegments, reductionOperands);
-  addOperands(operands, operandSegments, cacheOperands);
 
   auto loopOp = createRegionOp<mlir::acc::LoopOp, mlir::acc::YieldOp>(
       builder, currentLocation, eval, operands, operandSegments);

mlir/include/mlir/Dialect/OpenACC/OpenACC.h

@@ -48,12 +48,16 @@
       mlir::acc::CacheOp
 #define ACC_COMPUTE_CONSTRUCT_OPS                                              \
   mlir::acc::ParallelOp, mlir::acc::KernelsOp, mlir::acc::SerialOp
+#define ACC_COMPUTE_CONSTRUCT_AND_LOOP_OPS                                     \
+  ACC_COMPUTE_CONSTRUCT_OPS, mlir::acc::LoopOp
 #define ACC_DATA_CONSTRUCT_OPS                                                 \
   mlir::acc::DataOp, mlir::acc::EnterDataOp, mlir::acc::ExitDataOp,            \
       mlir::acc::UpdateOp, mlir::acc::HostDataOp, mlir::acc::DeclareEnterOp,   \
       mlir::acc::DeclareExitOp, mlir::acc::DeclareOp
 #define ACC_COMPUTE_AND_DATA_CONSTRUCT_OPS                                     \
   ACC_COMPUTE_CONSTRUCT_OPS, ACC_DATA_CONSTRUCT_OPS
+#define ACC_COMPUTE_LOOP_AND_DATA_CONSTRUCT_OPS                                \
+  ACC_COMPUTE_CONSTRUCT_AND_LOOP_OPS, ACC_DATA_CONSTRUCT_OPS
 
 namespace mlir {
 namespace acc {

mlir/include/mlir/Dialect/OpenACC/OpenACCOps.td

@@ -1196,11 +1196,12 @@ def OpenACC_LoopOp : OpenACC_Op<"loop",
       UnitAttr:$hasWorker,
       UnitAttr:$hasVector,
       Variadic<IntOrIndex>:$tileOperands,
+      Variadic<OpenACC_PointerLikeTypeInterface>:$cacheOperands,
       Variadic<OpenACC_PointerLikeTypeInterface>:$privateOperands,
       OptionalAttr<SymbolRefArrayAttr>:$privatizations,
       Variadic<AnyType>:$reductionOperands,
-      OptionalAttr<SymbolRefArrayAttr>:$reductionRecipes,
-      Variadic<OpenACC_PointerLikeTypeInterface>:$cacheOperands);
+      OptionalAttr<SymbolRefArrayAttr>:$reductionRecipes
+      );
 
   let results = (outs Variadic<AnyType>:$results);
 
@@ -1211,6 +1212,12 @@ def OpenACC_LoopOp : OpenACC_Op<"loop",
     static StringRef getGangNumKeyword() { return "num"; }
     static StringRef getGangDimKeyword() { return "dim"; }
     static StringRef getGangStaticKeyword() { return "static"; }
+
+    /// The number of private and reduction operands.
+    unsigned getNumDataOperands();
+
+    /// The i-th data operand passed.
+    Value getDataOperand(unsigned i);
   }];
 
   let hasCustomAssemblyFormat = 1;

mlir/lib/Dialect/OpenACC/IR/OpenACC.cpp

@@ -878,6 +878,22 @@ LogicalResult acc::LoopOp::verify() {
   return success();
 }
 
+unsigned LoopOp::getNumDataOperands() {
+  return getReductionOperands().size() + getPrivateOperands().size();
+}
+
+Value LoopOp::getDataOperand(unsigned i) {
+  unsigned numOptional = getGangNum() ? 1 : 0;
+  numOptional += getGangDim() ? 1 : 0;
+  numOptional += getGangStatic() ? 1 : 0;
+  numOptional += getVectorLength() ? 1 : 0;
+  numOptional += getWorkerNum() ? 1 : 0;
+  numOptional += getVectorLength() ? 1 : 0;
+  numOptional += getTileOperands().size();
+  numOptional += getCacheOperands().size();
+  return getOperand(numOptional + i);
+}
+
 //===----------------------------------------------------------------------===//
 // DataOp
 //===----------------------------------------------------------------------===//