diff --git a/mlir/include/mlir/Dialect/OpenACC/OpenACCUtilsCG.h b/mlir/include/mlir/Dialect/OpenACC/OpenACCUtilsCG.h index 8cfc4dac4ab36..9ab5365405d5d 100644 --- a/mlir/include/mlir/Dialect/OpenACC/OpenACCUtilsCG.h +++ b/mlir/include/mlir/Dialect/OpenACC/OpenACCUtilsCG.h @@ -95,6 +95,19 @@ void updateParDimsAttr(Operation *op, GPUParallelDimsAttr attr); /// Copy parallel dimensions from \p from to \p to. void copyParDimsAttr(Operation *from, Operation *to); +/// Create a gang dim 1 GPUParallelDimsAttr based on the mapping policy. +inline GPUParallelDimsAttr +getGangDim1ParDimsAttr(MLIRContext *ctx, ACCToGPUMappingPolicy &policy) { + return GPUParallelDimsAttr::get( + ctx, {policy.gangDim(ctx, acc::ParLevel::gang_dim1)}); +} + +/// Create a sequential GPUParallelDimsAttr based on the mapping policy. +inline GPUParallelDimsAttr getSeqParDimsAttr(MLIRContext *ctx, + ACCToGPUMappingPolicy &policy) { + return GPUParallelDimsAttr::get(ctx, {policy.seqDim(ctx)}); +} + /// Tracks aligned byte consumption against a configurable shared memory cap. class SharedMemoryBudget { public: diff --git a/mlir/lib/Dialect/OpenACC/Transforms/ACCRecipeMaterialization.cpp b/mlir/lib/Dialect/OpenACC/Transforms/ACCRecipeMaterialization.cpp index d2e8d77c168d1..e666bcd9a6242 100644 --- a/mlir/lib/Dialect/OpenACC/Transforms/ACCRecipeMaterialization.cpp +++ b/mlir/lib/Dialect/OpenACC/Transforms/ACCRecipeMaterialization.cpp @@ -383,25 +383,29 @@ LogicalResult ACCRecipeMaterialization::materialize( cloneRegionIntoAccRegion(&combinerRegion, &combineRegionOp.getRegion(), /*hasResult=*/false); - auto ctx = b.getContext(); + auto *ctx = b.getContext(); // For reductions that come from parallel constructs, explicitly set the // GPU parallel dimensions attribute to blockXDim since they will always be // gang private. GPU parallel dimensions cannot be determined for acc.loop // at this point. if constexpr (std::is_same_v) { - auto parDimsAttr = acc::GPUParallelDimsAttr::get( - ctx, {policy.gangDim(ctx, acc::ParLevel::gang_dim1)}); + acc::GPUParallelDimsAttr parDimsAttr; + if (accOp.isEffectivelySerial()) { + // If acc.serial has been lowered to a parallel op that is effectively + // sequential + parDimsAttr = acc::getSeqParDimsAttr(ctx, policy); + } else { + parDimsAttr = acc::getGangDim1ParDimsAttr(ctx, policy); + } acc::setParDimsAttr(reductionOp, parDimsAttr); acc::setParDimsAttr(combineRegionOp, parDimsAttr); } // Set sequential parallel dimensions attribute for loops in the recipe. - auto seqParDimsAttr = - acc::GPUParallelDimsAttr::get(ctx, {policy.seqDim(ctx)}); auto setSeqParDimsForRecipeLoops = [&](Region *r) { r->walk([&](LoopLikeOpInterface loopLike) { - acc::setParDimsAttr(loopLike, seqParDimsAttr); + acc::setParDimsAttr(loopLike, acc::getSeqParDimsAttr(ctx, policy)); }); }; setSeqParDimsForRecipeLoops(&reductionOp.getRegion()); diff --git a/mlir/test/Dialect/OpenACC/acc-recipe-materialization-reduction.mlir b/mlir/test/Dialect/OpenACC/acc-recipe-materialization-reduction.mlir index 6f2bafb2b6165..e5dc884d37d20 100644 --- a/mlir/test/Dialect/OpenACC/acc-recipe-materialization-reduction.mlir +++ b/mlir/test/Dialect/OpenACC/acc-recipe-materialization-reduction.mlir @@ -51,3 +51,35 @@ func.func @par_reduction_clause_(%arg0: memref) { } return } + +// CHECK-LABEL: func.func @par_reduction_clause_serial +// CHECK: acc.parallel {{.*}} { +// CHECK: [[PRIVATE:%.*]] = acc.reduction_init {{.*}} +// CHECK-NEXT: [[ZERO:%.*]] = arith.constant 0.000000e+00 : f64 +// CHECK-NEXT: [[ALLOCA:%.*]] = memref.alloca() : memref +// CHECK-NEXT: memref.store [[ZERO]], [[ALLOCA]][] +// CHECK-NEXT: acc.yield {{.*}} +// CHECK: } {{.*}}acc.par_dims = #acc, acc.var_name = #acc.var_name<"tmp"> +// CHECK: memref.load [[PRIVATE]][] +// CHECK: memref.store {{.*}}, [[PRIVATE]][] +// CHECK: acc.reduction_combine_region [[PRIVATE]] into [[REDUCVAR:%.*]] : memref { +// CHECK: [[LOADVAR:%.*]] = memref.load [[REDUCVAR]][] +// CHECK-NEXT: [[LOADPRIV:%.*]] = memref.load [[PRIVATE]][] +// CHECK-NEXT: [[COMBINE:%.*]] = arith.addf [[LOADVAR]], [[LOADPRIV]] +// CHECK-NEXT: memref.store [[COMBINE]], [[REDUCVAR]][] +// CHECK-NEXT: } {acc.par_dims = #acc} +// CHECK-NEXT: memref.dealloc [[PRIVATE]] : memref +// CHECK: acc.yield + +func.func @par_reduction_clause_serial(%arg0: memref) { + %c1_i32 = arith.constant 1 : i32 + %cst = arith.constant 1.000000e+00 : f64 + %red = acc.reduction varPtr(%arg0 : memref) recipe(@reduction_add_memref_f64) -> memref {name = "tmp"} + acc.parallel num_gangs({%c1_i32 : i32}) num_workers(%c1_i32 : i32) vector_length(%c1_i32 : i32) reduction(%red : memref) { + %3 = memref.load %red[] : memref + %4 = arith.addf %3, %cst fastmath : f64 + memref.store %4, %red[] : memref + acc.yield + } + return +}