diff --git a/flang/lib/Optimizer/HLFIR/Transforms/LowerHLFIROrderedAssignments.cpp b/flang/lib/Optimizer/HLFIR/Transforms/LowerHLFIROrderedAssignments.cpp --- a/flang/lib/Optimizer/HLFIR/Transforms/LowerHLFIROrderedAssignments.cpp +++ b/flang/lib/Optimizer/HLFIR/Transforms/LowerHLFIROrderedAssignments.cpp @@ -364,6 +364,45 @@ builder.setInsertionPointAfter(constructStack.pop_back_val()); } +/// Convert an entity to the type of a given mold. +/// This is intended to help with cases where hlfir entity is a value while +/// it must be used as a variable or vice-versa. These mismatches may occur +/// between the type of user defined assignment block arguments and the actual +/// argument that was lowered for them. The actual may be an in-memory copy +/// while the block argument expects an hlfir.expr. +static hlfir::Entity +convertToMoldType(mlir::Location loc, fir::FirOpBuilder &builder, + hlfir::Entity input, hlfir::Entity mold, + llvm::SmallVectorImpl &cleanups) { + if (input.getType() == mold.getType()) + return input; + fir::FirOpBuilder *b = &builder; + if (input.isVariable() && mold.isValue()) { + if (fir::isa_trivial(mold.getType())) { + // fir.ref to T. + mlir::Value load = builder.create(loc, input); + return hlfir::Entity{builder.createConvert(loc, mold.getType(), load)}; + } + // fir.ref to hlfir.expr. + mlir::Value asExpr = builder.create(loc, input); + if (asExpr.getType() != mold.getType()) + TODO(loc, "hlfir.expr conversion"); + cleanups.emplace_back([=]() { b->create(loc, asExpr); }); + return hlfir::Entity{asExpr}; + } + if (input.isValue() && mold.isVariable()) { + // T to fir.ref, or hlfir.expr to fir.ref. + hlfir::AssociateOp associate = hlfir::genAssociateExpr( + loc, builder, input, mold.getFortranElementType(), ".tmp.val2ref"); + cleanups.emplace_back( + [=]() { b->create(loc, associate); }); + return hlfir::Entity{associate.getBase()}; + } + // Variable to Variable mismatch (e.g., fir.heap vs fir.ref), or value + // to Value mismatch (e.g. i1 vs fir.logical<4>). + return hlfir::Entity{builder.createConvert(loc, mold.getType(), input)}; +} + void OrderedAssignmentRewriter::pre(hlfir::RegionAssignOp regionAssignOp) { mlir::Location loc = regionAssignOp.getLoc(); auto [rhs, oldRhsYield] = @@ -372,11 +411,45 @@ TODO(loc, "assignment to vector subscripted entity"); auto [lhs, oldLhsYield] = generateYieldedEntity(regionAssignOp.getLhsRegion()); - if (!regionAssignOp.getUserDefinedAssignment().empty()) - TODO(loc, "user defined assignment inside FORALL or WHERE"); - // TODO: preserve allocatable assignment aspects for forall once - // they are conveyed in hlfir.region_assign. - builder.create(loc, rhs, lhs); + if (!regionAssignOp.getUserDefinedAssignment().empty()) { + hlfir::Entity userAssignLhs{regionAssignOp.getUserAssignmentLhs()}; + hlfir::Entity userAssignRhs{regionAssignOp.getUserAssignmentRhs()}; + hlfir::Entity lhsEntity{lhs}; + hlfir::Entity rhsEntity{rhs}; + fir::DoLoopOp outerElementalLoop = nullptr; + if (lhsEntity.isArray() && userAssignLhs.isScalar()) { + // Elemental assignment with array argument (the RHS cannot be an array + // if the LHS is not). + mlir::Value shape = hlfir::genShape(loc, builder, lhsEntity); + hlfir::LoopNest elementalLoopNest = + hlfir::genLoopNest(loc, builder, shape); + outerElementalLoop = elementalLoopNest.outerLoop; + builder.setInsertionPointToStart(elementalLoopNest.innerLoop.getBody()); + lhsEntity = hlfir::getElementAt(loc, builder, lhsEntity, + elementalLoopNest.oneBasedIndices); + rhsEntity = hlfir::getElementAt(loc, builder, rhsEntity, + elementalLoopNest.oneBasedIndices); + } + + llvm::SmallVector argConversionCleanups; + lhsEntity = convertToMoldType(loc, builder, lhsEntity, userAssignLhs, + argConversionCleanups); + rhsEntity = convertToMoldType(loc, builder, rhsEntity, userAssignRhs, + argConversionCleanups); + mapper.map(userAssignLhs, lhsEntity); + mapper.map(userAssignRhs, rhsEntity); + for (auto &op : + regionAssignOp.getUserDefinedAssignment().front().without_terminator()) + (void)builder.clone(op, mapper); + for (auto &cleanupConversion : argConversionCleanups) + cleanupConversion(); + if (outerElementalLoop) + builder.setInsertionPointAfter(outerElementalLoop); + } else { + // TODO: preserve allocatable assignment aspects for forall once + // they are conveyed in hlfir.region_assign. + builder.create(loc, rhs, lhs); + } generateCleanupIfAny(oldRhsYield); generateCleanupIfAny(oldLhsYield); } @@ -995,12 +1068,12 @@ mlir::LogicalResult matchAndRewrite(hlfir::RegionAssignOp regionAssignOp, mlir::PatternRewriter &rewriter) const override { + auto root = mlir::cast( + regionAssignOp.getOperation()); if (!regionAssignOp.getUserDefinedAssignment().empty()) - TODO(regionAssignOp.getLoc(), "user defined assignment in HLFIR"); - else - TODO(regionAssignOp.getLoc(), - "assignments to vector subscripted entity in HLFIR"); - return mlir::failure(); + return ::rewrite(root, /*tryFusingAssignments=*/false, rewriter); + TODO(regionAssignOp.getLoc(), + "assignments to vector subscripted entity in HLFIR"); } }; diff --git a/flang/lib/Optimizer/HLFIR/Transforms/ScheduleOrderedAssignments.cpp b/flang/lib/Optimizer/HLFIR/Transforms/ScheduleOrderedAssignments.cpp --- a/flang/lib/Optimizer/HLFIR/Transforms/ScheduleOrderedAssignments.cpp +++ b/flang/lib/Optimizer/HLFIR/Transforms/ScheduleOrderedAssignments.cpp @@ -242,9 +242,21 @@ "assignment to vector subscripted entity in HLFIR"); assignEffects.emplace_back(mlir::MemoryEffects::Write::get(), assignedVar); - // TODO: gather the read/write effects of user defined assignments. - if (!regionAssign.getUserDefinedAssignment().empty()) - TODO(regionAssign.getLoc(), "user defined assignments"); + if (!regionAssign.getUserDefinedAssignment().empty()) { + // The write effect on the INTENT(OUT) LHS argument is already taken + // into account above. + // This side effects are "defensive" and could be improved. + // On top of the passed RHS argument, user defined assignments (even when + // pure) may also read host/used/common variable. Impure user defined + // assignments may write to host/used/common variables not passed via + // arguments. For now, simply assume the worst. Once fir.call side effects + // analysis is improved, it would best to let the call side effects be used + // directly. + if (userDefAssignmentMayOnlyWriteToAssignedVariable) + assignEffects.emplace_back(mlir::MemoryEffects::Read::get()); + else + assignEffects.emplace_back(mlir::MemoryEffects::Write::get()); + } } //===----------------------------------------------------------------------===// diff --git a/flang/test/HLFIR/order_assignments/user-defined-assignment.fir b/flang/test/HLFIR/order_assignments/user-defined-assignment.fir new file mode 100644 --- /dev/null +++ b/flang/test/HLFIR/order_assignments/user-defined-assignment.fir @@ -0,0 +1,182 @@ +// Test code generation of hlfir.region_assign with user defined +// assignment. +// RUN: fir-opt %s --lower-hlfir-ordered-assignments | FileCheck %s + +func.func @test_simple_scalar(%i: !fir.ref, %l: !fir.ref>) { + hlfir.region_assign { + hlfir.yield %l : !fir.ref> + } to { + hlfir.yield %i : !fir.ref + } user_defined_assign (%arg0: !fir.ref>) to (%arg1: !fir.ref) { + fir.call @logical_to_numeric(%arg1, %arg0) : (!fir.ref, !fir.ref>) -> () + } + return +} +// CHECK-LABEL: func.func @test_simple_scalar( +// CHECK-SAME: %[[VAL_0:.*]]: !fir.ref, +// CHECK-SAME: %[[VAL_1:.*]]: !fir.ref>) { +// CHECK: %[[VAL_2:.*]] = fir.load %[[VAL_1]] : !fir.ref> +// CHECK: %[[VAL_3:.*]]:3 = hlfir.associate %[[VAL_2]] {uniq_name = ".tmp.forall"} : (!fir.logical<4>) -> (!fir.ref>, !fir.ref>, i1) +// CHECK: fir.call @logical_to_numeric(%[[VAL_0]], %[[VAL_3]]#0) : (!fir.ref, !fir.ref>) -> () +// CHECK: hlfir.end_associate %[[VAL_3]]#1, %[[VAL_3]]#2 : !fir.ref>, i1 + +func.func @test_elemental_overlap(%i: !fir.ref>) { + %c0_i32 = arith.constant 0 : i32 + %c10 = arith.constant 10 : index + %shape = fir.shape %c10 : (index) -> !fir.shape<1> + hlfir.region_assign { + %cmp = hlfir.elemental %shape : (!fir.shape<1>) -> !hlfir.expr<10x!fir.logical<4>> { + ^bb0(%j: index): + %ielt = hlfir.designate %i (%j) : (!fir.ref>, index) -> !fir.ref + %ielt_val = fir.load %ielt : !fir.ref + %smaller = arith.cmpi slt, %ielt_val, %c0_i32 : i32 + %smaller_cast = fir.convert %smaller : (i1) -> !fir.logical<4> + hlfir.yield_element %smaller_cast : !fir.logical<4> + } + hlfir.yield %cmp : !hlfir.expr<10x!fir.logical<4>> + } to { + hlfir.yield %i : !fir.ref> + } user_defined_assign (%arg0: !fir.logical<4>) to (%arg1: !fir.ref) { + fir.call @logical_value_to_numeric(%arg1, %arg0) : (!fir.ref, !fir.logical<4>) -> () + } + return +} +// CHECK-LABEL: func.func @test_elemental_overlap( +// CHECK-SAME: %[[VAL_0:.*]]: !fir.ref>) { +// CHECK: %[[VAL_1:.*]] = arith.constant 0 : i32 +// CHECK: %[[VAL_2:.*]] = arith.constant 10 : index +// CHECK: %[[VAL_3:.*]] = fir.shape %[[VAL_2]] : (index) -> !fir.shape<1> +// CHECK: %[[VAL_4:.*]] = hlfir.elemental %[[VAL_3]] : (!fir.shape<1>) -> !hlfir.expr<10x!fir.logical<4>> { +// CHECK: ^bb0(%[[VAL_5:.*]]: index): +// CHECK: %[[VAL_6:.*]] = hlfir.designate %[[VAL_0]] (%[[VAL_5]]) : (!fir.ref>, index) -> !fir.ref +// CHECK: %[[VAL_7:.*]] = fir.load %[[VAL_6]] : !fir.ref +// CHECK: %[[VAL_8:.*]] = arith.cmpi slt, %[[VAL_7]], %[[VAL_1]] : i32 +// CHECK: %[[VAL_9:.*]] = fir.convert %[[VAL_8]] : (i1) -> !fir.logical<4> +// CHECK: hlfir.yield_element %[[VAL_9]] : !fir.logical<4> +// CHECK: } +// CHECK: %[[VAL_10:.*]]:3 = hlfir.associate %[[VAL_11:.*]](%[[VAL_3]]) {uniq_name = ".tmp.forall"} : (!hlfir.expr<10x!fir.logical<4>>, !fir.shape<1>) -> (!fir.ref>>, !fir.ref>>, i1) +// CHECK: %[[VAL_12:.*]] = arith.constant 10 : index +// CHECK: %[[VAL_13:.*]] = fir.shape %[[VAL_12]] : (index) -> !fir.shape<1> +// CHECK: %[[VAL_14:.*]] = arith.constant 1 : index +// CHECK: fir.do_loop %[[VAL_15:.*]] = %[[VAL_14]] to %[[VAL_12]] step %[[VAL_14]] { +// CHECK: %[[VAL_16:.*]] = hlfir.designate %[[VAL_0]] (%[[VAL_15]]) : (!fir.ref>, index) -> !fir.ref +// CHECK: %[[VAL_17:.*]] = hlfir.designate %[[VAL_10]]#0 (%[[VAL_15]]) : (!fir.ref>>, index) -> !fir.ref> +// CHECK: %[[VAL_18:.*]] = fir.load %[[VAL_17]] : !fir.ref> +// CHECK: fir.call @logical_value_to_numeric(%[[VAL_16]], %[[VAL_18]]) : (!fir.ref, !fir.logical<4>) -> () +// CHECK: } +// CHECK: hlfir.end_associate %[[VAL_10]]#1, %[[VAL_10]]#2 : !fir.ref>>, i1 + +func.func @test_array_overlap(%i: !fir.ref>) { + %c0_i32 = arith.constant 0 : i32 + %c10 = arith.constant 10 : index + %shape = fir.shape %c10 : (index) -> !fir.shape<1> + hlfir.region_assign { + %cmp = hlfir.elemental %shape : (!fir.shape<1>) -> !hlfir.expr<10x!fir.logical<4>> { + ^bb0(%j: index): + %ielt = hlfir.designate %i (%j) : (!fir.ref>, index) -> !fir.ref + %ielt_val = fir.load %ielt : !fir.ref + %smaller = arith.cmpi slt, %ielt_val, %c0_i32 : i32 + %smaller_cast = fir.convert %smaller : (i1) -> !fir.logical<4> + hlfir.yield_element %smaller_cast : !fir.logical<4> + } + hlfir.yield %cmp : !hlfir.expr<10x!fir.logical<4>> + } to { + hlfir.yield %i : !fir.ref> + } user_defined_assign (%arg0: !hlfir.expr<10x!fir.logical<4>>) to (%arg1: !fir.ref>) { + %1:3 = hlfir.associate %arg0(%shape) {uniq_name = "adapt.valuebyref"} : (!hlfir.expr<10x!fir.logical<4>>, !fir.shape<1>) -> (!fir.ref>>, !fir.ref>>, i1) + fir.call @logical_array_to_numeric(%arg1, %1#0) : (!fir.ref>, !fir.ref>>) -> () + hlfir.end_associate %1#1, %1#2 : !fir.ref>>, i1 + } + return +} +// CHECK-LABEL: func.func @test_array_overlap( +// CHECK-SAME: %[[VAL_0:.*]]: !fir.ref>) { +// CHECK: %[[VAL_1:.*]] = arith.constant 0 : i32 +// CHECK: %[[VAL_2:.*]] = arith.constant 10 : index +// CHECK: %[[VAL_3:.*]] = fir.shape %[[VAL_2]] : (index) -> !fir.shape<1> +// CHECK: %[[VAL_4:.*]] = hlfir.elemental %[[VAL_3]] : (!fir.shape<1>) -> !hlfir.expr<10x!fir.logical<4>> { +// CHECK: ^bb0(%[[VAL_5:.*]]: index): +// CHECK: %[[VAL_6:.*]] = hlfir.designate %[[VAL_0]] (%[[VAL_5]]) : (!fir.ref>, index) -> !fir.ref +// CHECK: %[[VAL_7:.*]] = fir.load %[[VAL_6]] : !fir.ref +// CHECK: %[[VAL_8:.*]] = arith.cmpi slt, %[[VAL_7]], %[[VAL_1]] : i32 +// CHECK: %[[VAL_9:.*]] = fir.convert %[[VAL_8]] : (i1) -> !fir.logical<4> +// CHECK: hlfir.yield_element %[[VAL_9]] : !fir.logical<4> +// CHECK: } +// CHECK: %[[VAL_10:.*]]:3 = hlfir.associate %[[VAL_11:.*]](%[[VAL_3]]) {uniq_name = ".tmp.forall"} : (!hlfir.expr<10x!fir.logical<4>>, !fir.shape<1>) -> (!fir.ref>>, !fir.ref>>, i1) +// CHECK: %[[VAL_12:.*]] = hlfir.as_expr %[[VAL_10]]#0 : (!fir.ref>>) -> !hlfir.expr<10x!fir.logical<4>> +// CHECK: %[[VAL_13:.*]]:3 = hlfir.associate %[[VAL_12]](%[[VAL_3]]) {uniq_name = "adapt.valuebyref"} : (!hlfir.expr<10x!fir.logical<4>>, !fir.shape<1>) -> (!fir.ref>>, !fir.ref>>, i1) +// CHECK: fir.call @logical_array_to_numeric(%[[VAL_0]], %[[VAL_13]]#0) : (!fir.ref>, !fir.ref>>) -> () +// CHECK: hlfir.end_associate %[[VAL_13]]#1, %[[VAL_13]]#2 : !fir.ref>>, i1 +// CHECK: hlfir.destroy %[[VAL_12]] : !hlfir.expr<10x!fir.logical<4>> +// CHECK: hlfir.end_associate %[[VAL_10]]#1, %[[VAL_10]]#2 : !fir.ref>>, i1 + +func.func @test_scalar_forall_overlap(%i: !fir.ref>) { + %c0_i32 = arith.constant 0 : i32 + %c1 = arith.constant 1 : index + %c10 = arith.constant 10 : index + %c11 = arith.constant 11 : index + hlfir.forall lb { + hlfir.yield %c1 : index + } ub { + hlfir.yield %c10 : index + } (%j: index) { + hlfir.region_assign { + %reverse_j = arith.subi %c11, %j : index + %ielt = hlfir.designate %i (%reverse_j) : (!fir.ref>, index) -> !fir.ref + %ielt_val = fir.load %ielt : !fir.ref + %smaller = arith.cmpi slt, %ielt_val, %c0_i32 : i32 + hlfir.yield %smaller : i1 + } to { + %ielt = hlfir.designate %i (%j) : (!fir.ref>, index) -> !fir.ref + hlfir.yield %ielt : !fir.ref + } user_defined_assign (%arg0: i1) to (%arg1: !fir.ref) { + %cast = fir.convert %arg0 : (i1) -> !fir.logical<4> + fir.call @logical_value_to_numeric(%arg1, %cast) : (!fir.ref, !fir.logical<4>) -> () + } + } + return +} +// CHECK-LABEL: func.func @test_scalar_forall_overlap( +// CHECK-SAME: %[[VAL_0:.*]]: !fir.ref>) { +// CHECK: %[[VAL_1:.*]] = fir.alloca index +// CHECK: %[[VAL_2:.*]] = arith.constant 0 : i32 +// CHECK: %[[VAL_3:.*]] = arith.constant 1 : index +// CHECK: %[[VAL_4:.*]] = arith.constant 10 : index +// CHECK: %[[VAL_5:.*]] = arith.constant 11 : index +// CHECK: %[[VAL_6:.*]] = arith.constant 1 : index +// CHECK: %[[VAL_7:.*]] = arith.constant 0 : index +// CHECK: %[[VAL_8:.*]] = arith.subi %[[VAL_4]], %[[VAL_3]] : index +// CHECK: %[[VAL_9:.*]] = arith.addi %[[VAL_8]], %[[VAL_6]] : index +// CHECK: %[[VAL_10:.*]] = arith.divsi %[[VAL_9]], %[[VAL_6]] : index +// CHECK: %[[VAL_11:.*]] = arith.cmpi sgt, %[[VAL_10]], %[[VAL_7]] : index +// CHECK: %[[VAL_12:.*]] = arith.select %[[VAL_11]], %[[VAL_10]], %[[VAL_7]] : index +// CHECK: %[[VAL_13:.*]] = arith.constant 1 : index +// CHECK: %[[VAL_14:.*]] = arith.constant 1 : index +// CHECK: fir.store %[[VAL_13]] to %[[VAL_1]] : !fir.ref +// CHECK: %[[VAL_15:.*]] = fir.allocmem !fir.array, %[[VAL_12]] {bindc_name = ".tmp.forall", uniq_name = ""} +// CHECK: %[[VAL_16:.*]] = fir.shape %[[VAL_12]] : (index) -> !fir.shape<1> +// CHECK: %[[VAL_17:.*]]:2 = hlfir.declare %[[VAL_15]](%[[VAL_16]]) {uniq_name = ".tmp.forall"} : (!fir.heap>, !fir.shape<1>) -> (!fir.box>, !fir.heap>) +// CHECK: fir.do_loop %[[VAL_18:.*]] = %[[VAL_3]] to %[[VAL_4]] step %[[VAL_6]] { +// CHECK: %[[VAL_19:.*]] = arith.subi %[[VAL_5]], %[[VAL_18]] : index +// CHECK: %[[VAL_20:.*]] = hlfir.designate %[[VAL_0]] (%[[VAL_19]]) : (!fir.ref>, index) -> !fir.ref +// CHECK: %[[VAL_21:.*]] = fir.load %[[VAL_20]] : !fir.ref +// CHECK: %[[VAL_22:.*]] = arith.cmpi slt, %[[VAL_21]], %[[VAL_2]] : i32 +// CHECK: %[[VAL_23:.*]] = fir.load %[[VAL_1]] : !fir.ref +// CHECK: %[[VAL_24:.*]] = arith.addi %[[VAL_23]], %[[VAL_14]] : index +// CHECK: fir.store %[[VAL_24]] to %[[VAL_1]] : !fir.ref +// CHECK: %[[VAL_25:.*]] = hlfir.designate %[[VAL_17]]#0 (%[[VAL_23]]) : (!fir.box>, index) -> !fir.ref +// CHECK: hlfir.assign %[[VAL_22]] to %[[VAL_25]] : i1, !fir.ref +// CHECK: } +// CHECK: %[[VAL_26:.*]] = arith.constant 1 : index +// CHECK: fir.store %[[VAL_13]] to %[[VAL_1]] : !fir.ref +// CHECK: fir.do_loop %[[VAL_27:.*]] = %[[VAL_3]] to %[[VAL_4]] step %[[VAL_26]] { +// CHECK: %[[VAL_28:.*]] = fir.load %[[VAL_1]] : !fir.ref +// CHECK: %[[VAL_29:.*]] = arith.addi %[[VAL_28]], %[[VAL_14]] : index +// CHECK: fir.store %[[VAL_29]] to %[[VAL_1]] : !fir.ref +// CHECK: %[[VAL_30:.*]] = hlfir.designate %[[VAL_17]]#0 (%[[VAL_28]]) : (!fir.box>, index) -> !fir.ref +// CHECK: %[[VAL_31:.*]] = fir.load %[[VAL_30]] : !fir.ref +// CHECK: %[[VAL_32:.*]] = hlfir.designate %[[VAL_0]] (%[[VAL_27]]) : (!fir.ref>, index) -> !fir.ref +// CHECK: %[[VAL_33:.*]] = fir.convert %[[VAL_31]] : (i1) -> !fir.logical<4> +// CHECK: fir.call @logical_value_to_numeric(%[[VAL_32]], %[[VAL_33]]) : (!fir.ref, !fir.logical<4>) -> () +// CHECK: } +// CHECK: fir.freemem %[[VAL_15]] : !fir.heap>