Skip to content

Commit

Permalink
[flang][hlfir] Add hlfir.elemental_addr for vector subscripted assign…
Browse files Browse the repository at this point in the history 
…ment

See the operation description in HLFIROps.td.

Depends on D149442

Differential Revision: https://reviews.llvm.org/D149449
  • Loading branch information
@jeanPerier
jeanPerier committed May 3, 2023
1 parent 64b591a commit bdc9914
Show file tree
Hide file tree
Showing 4 changed files with 252 additions and 4 deletions.
66 changes: 65 additions & 1 deletion flang/include/flang/Optimizer/HLFIR/HLFIROps.td
View file
Original file line number Diff line number Diff line change
Expand Up @@ -868,7 +868,7 @@ def hlfir_RegionAssignOp : hlfir_Op<"region_assign", []> {
let hasVerifier = 1;
}

def hlfir_YieldOp : hlfir_Op<"yield", [Terminator, ParentOneOf<["RegionAssignOp"]>,
def hlfir_YieldOp : hlfir_Op<"yield", [Terminator, ParentOneOf<["RegionAssignOp", "ElementalAddrOp"]>,
SingleBlockImplicitTerminator<"fir::FirEndOp">]> {

let summary = "Yield a value or variable inside a forall, where or region assignment";
Expand Down Expand Up @@ -898,4 +898,68 @@ def hlfir_YieldOp : hlfir_Op<"yield", [Terminator, ParentOneOf<["RegionAssignOp"
let assemblyFormat = "$entity attr-dict `:` type($entity) custom<YieldOpCleanup>($cleanup)";
}

def hlfir_ElementalAddrOp : hlfir_Op<"elemental_addr", [Terminator, HasParent<"RegionAssignOp">, RecursiveMemoryEffects, RecursivelySpeculatable]> {
let summary = "Yield the address of a vector subscripted variable inside an hlfir.region_assign";
let description = [{
Special terminator node for the left-hand side region of an hlfir.region_assign
to a vector subscripted entity.

It represents how the address of an element of such entity is computed given
one based indices.

It is very similar to hlfir.elemental, except that it does not produce an SSA
value because there is no hlfir type to describe a vector subscripted entity
(the codegen of such type would be problematic). Hence, it is tightly linked
to an hlfir.region_assign by its terminator property.

An optional cleanup region may be provided if any of the subscript expressions
of the designator require a cleanup.
This allows documenting cleanups that cannot be generated after the vector
subscripted designator usage (that has not been materizaled yet). The cleanups
will be evaluated after the assignment once the related
hlfir.region_assign is lowered.

Example: "X(VECTOR) = Y"

```
hlfir.region_assign {
hlfir.yield %y : !fir.ref<!fir.array<20xf32>>
} to {
hlfir.elemental_addr %vector_shape : !fir.shape<1> {
^bb0(%i: index):
%0 = hlfir.designate %vector (%i) : (!fir.ref<!fir.array<20xi32>>, index) -> !fir.ref<i32>
%1 = fir.load %0 : !fir.ref<i32>
%x_element_addr = hlfir.designate %x (%1) : (!fir.ref<!fir.array<100xf32>>, i32) -> !fir.ref<f32>
hlfir.yield %x_element_addr : !fir.ref<f32>
}
}
```
}];

let arguments = (ins
fir_ShapeType:$shape,
Variadic<AnyIntegerType>:$typeparams);

let regions = (region SizedRegion<1>:$body,
MaxSizedRegion<1>:$cleanup);

let builders = [
OpBuilder<(ins "mlir::Value":$shape)>
];

let assemblyFormat = [{
$shape (`typeparams` $typeparams^)?
attr-dict `:` type(operands) $body
custom<YieldOpCleanup>($cleanup)}];

let extraClassDeclaration = [{
mlir::Region::BlockArgListType getIndices() {
return getBody().getArguments();
}
}];

let hasVerifier = 1;
}


#endif // FORTRAN_DIALECT_HLFIR_OPS
41 changes: 38 additions & 3 deletions flang/lib/Optimizer/HLFIR/IR/HLFIROps.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1032,8 +1032,8 @@ mlir::LogicalResult hlfir::RegionAssignOp::verify() {
if (!mlir::isa_and_nonnull<hlfir::YieldOp>(getTerminator(getRhsRegion())))
return emitOpError(
"right-hand side region must be terminated by an hlfir.yield");
// TODO: allow hlfir.elemental_addr.
if (!mlir::isa_and_nonnull<hlfir::YieldOp>(getTerminator(getLhsRegion())))
if (!mlir::isa_and_nonnull<hlfir::YieldOp, hlfir::ElementalAddrOp>(
getTerminator(getLhsRegion())))
return emitOpError("left-hand side region must be terminated by an "
"hlfir.yield or hlfir.elemental_addr");
return mlir::success();
Expand All @@ -1052,7 +1052,6 @@ static mlir::ParseResult parseYieldOpCleanup(mlir::OpAsmParser &parser,
hlfir::YieldOp::ensureTerminator(cleanup, parser.getBuilder(),
parser.getBuilder().getUnknownLoc());
}

return mlir::success();
}

Expand All @@ -1066,5 +1065,41 @@ static void printYieldOpCleanup(mlir::OpAsmPrinter &p, YieldOp yieldOp,
}
}

//===----------------------------------------------------------------------===//
// ElementalAddrOp
//===----------------------------------------------------------------------===//

void hlfir::ElementalAddrOp::build(mlir::OpBuilder &builder,
mlir::OperationState &odsState,
mlir::Value shape) {
odsState.addOperands(shape);
mlir::Region *bodyRegion = odsState.addRegion();
bodyRegion->push_back(new mlir::Block{});
if (auto shapeType = shape.getType().dyn_cast<fir::ShapeType>()) {
unsigned dim = shapeType.getRank();
mlir::Type indexType = builder.getIndexType();
for (unsigned d = 0; d < dim; ++d)
bodyRegion->front().addArgument(indexType, odsState.location);
}
// Push cleanUp region.
odsState.addRegion();
}

mlir::LogicalResult hlfir::ElementalAddrOp::verify() {
hlfir::YieldOp yieldOp =
mlir::dyn_cast_or_null<hlfir::YieldOp>(getTerminator(getBody()));
if (!yieldOp)
return emitOpError("body region must be terminated by an hlfir.yield");
mlir::Type elementAddrType = yieldOp.getEntity().getType();
if (!hlfir::isFortranVariableType(elementAddrType) ||
hlfir::getFortranElementOrSequenceType(elementAddrType)
.isa<fir::SequenceType>())
return emitOpError("body must compute the address of a scalar entity");
unsigned shapeRank = getShape().getType().cast<fir::ShapeType>().getRank();
if (shapeRank != getIndices().size())
return emitOpError("body number of indices must match shape rank");
return mlir::success();
}

#define GET_OP_CLASSES
#include "flang/Optimizer/HLFIR/HLFIROps.cpp.inc"
83 changes: 83 additions & 0 deletions flang/test/HLFIR/element-addr.fir
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,83 @@
// Test hlfir.element_addr operation parse, verify (no errors), and unparse.
// RUN: fir-opt %s | fir-opt | FileCheck %s

func.func @test_element_addr(%x: !fir.ref<!fir.array<100xf32>>, %vector: !fir.ref<!fir.array<20xi32>>, %y: !fir.ref<!fir.array<20xf32>>) {
%c20 = arith.constant 20 : index
%vector_shape = fir.shape %c20 : (index) -> !fir.shape<1>
hlfir.region_assign {
hlfir.yield %y : !fir.ref<!fir.array<20xf32>>
} to {
hlfir.elemental_addr %vector_shape : !fir.shape<1> {
^bb0(%i: index):
%0 = hlfir.designate %vector (%i) : (!fir.ref<!fir.array<20xi32>>, index) -> !fir.ref<i32>
%1 = fir.load %0 : !fir.ref<i32>
%x_element_addr = hlfir.designate %x (%1) : (!fir.ref<!fir.array<100xf32>>, i32) -> !fir.ref<f32>
hlfir.yield %x_element_addr : !fir.ref<f32>
}
}
return
}
// CHECK-LABEL: func.func @test_element_addr(
// CHECK-SAME: %[[VAL_0:.*]]: !fir.ref<!fir.array<100xf32>>,
// CHECK-SAME: %[[VAL_1:.*]]: !fir.ref<!fir.array<20xi32>>,
// CHECK-SAME: %[[VAL_2:.*]]: !fir.ref<!fir.array<20xf32>>) {
// CHECK: %[[VAL_3:.*]] = arith.constant 20 : index
// CHECK: %[[VAL_4:.*]] = fir.shape %[[VAL_3]] : (index) -> !fir.shape<1>
// CHECK: hlfir.region_assign {
// CHECK: hlfir.yield %[[VAL_2]] : !fir.ref<!fir.array<20xf32>>
// CHECK: } to {
// CHECK: hlfir.elemental_addr %[[VAL_4]] : !fir.shape<1> {
// CHECK: ^bb0(%[[VAL_5:.*]]: index):
// CHECK: %[[VAL_6:.*]] = hlfir.designate %[[VAL_1]] (%[[VAL_5]]) : (!fir.ref<!fir.array<20xi32>>, index) -> !fir.ref<i32>
// CHECK: %[[VAL_7:.*]] = fir.load %[[VAL_6]] : !fir.ref<i32>
// CHECK: %[[VAL_8:.*]] = hlfir.designate %[[VAL_0]] (%[[VAL_7]]) : (!fir.ref<!fir.array<100xf32>>, i32) -> !fir.ref<f32>
// CHECK: hlfir.yield %[[VAL_8]] : !fir.ref<f32>
// CHECK: }
// CHECK: }



func.func @test_element_addr_cleanup(%x: !fir.box<!fir.array<?x!fir.char<1,?>>>, %y: !fir.box<!fir.array<?x!fir.char<1,?>>>) {
hlfir.region_assign {
hlfir.yield %y : !fir.box<!fir.array<?x!fir.char<1,?>>>
} to {
%c0 = arith.constant 0 : index
%len = fir.box_elesize %x : (!fir.box<!fir.array<?x!fir.char<1,?>>>) -> index
%vector = fir.call @returns_alloc() : () -> !fir.box<!fir.heap<!fir.array<?xi32>>>
%vector_dim:3 = fir.box_dims %x, %c0 : (!fir.box<!fir.array<?x!fir.char<1,?>>>, index) -> (index, index, index)
%vector_shape = fir.shape %vector_dim#1 : (index) -> !fir.shape<1>
hlfir.elemental_addr %vector_shape typeparams %len : !fir.shape<1>, index {
^bb0(%i: index):
%0 = hlfir.designate %vector (%i) : (!fir.box<!fir.heap<!fir.array<?xi32>>>, index) -> !fir.ref<i32>
%1 = fir.load %0 : !fir.ref<i32>
%x_element_addr = hlfir.designate %x (%1) typeparams %len : (!fir.box<!fir.array<?x!fir.char<1,?>>>, i32, index) -> !fir.boxchar<1>
hlfir.yield %x_element_addr : !fir.boxchar<1>
} cleanup {
%addr = fir.box_addr %vector : (!fir.box<!fir.heap<!fir.array<?xi32>>>) -> !fir.heap<!fir.array<?xi32>>
fir.freemem %addr : !fir.heap<!fir.array<?xi32>>
}
}
return
}
// CHECK-LABEL: func.func @test_element_addr_cleanup(
// CHECK-SAME: %[[VAL_0:[^:]*]]: !fir.box<!fir.array<?x!fir.char<1,?>>>,
// CHECK-SAME: %[[VAL_1:.*]]: !fir.box<!fir.array<?x!fir.char<1,?>>>) {
// CHECK: hlfir.region_assign {
// CHECK: hlfir.yield %[[VAL_1]] : !fir.box<!fir.array<?x!fir.char<1,?>>>
// CHECK: } to {
// CHECK: %[[VAL_2:.*]] = arith.constant 0 : index
// CHECK: %[[VAL_3:.*]] = fir.box_elesize %[[VAL_0]] : (!fir.box<!fir.array<?x!fir.char<1,?>>>) -> index
// CHECK: %[[VAL_4:.*]] = fir.call @returns_alloc() : () -> !fir.box<!fir.heap<!fir.array<?xi32>>>
// CHECK: %[[VAL_5:.*]]:3 = fir.box_dims %[[VAL_0]], %[[VAL_2]] : (!fir.box<!fir.array<?x!fir.char<1,?>>>, index) -> (index, index, index)
// CHECK: %[[VAL_6:.*]] = fir.shape %[[VAL_5]]#1 : (index) -> !fir.shape<1>
// CHECK: hlfir.elemental_addr %[[VAL_6]] typeparams %[[VAL_3]] : !fir.shape<1>, index {
// CHECK: ^bb0(%[[VAL_7:.*]]: index):
// CHECK: %[[VAL_8:.*]] = hlfir.designate %[[VAL_4]] (%[[VAL_7]]) : (!fir.box<!fir.heap<!fir.array<?xi32>>>, index) -> !fir.ref<i32>
// CHECK: %[[VAL_9:.*]] = fir.load %[[VAL_8]] : !fir.ref<i32>
// CHECK: %[[VAL_10:.*]] = hlfir.designate %[[VAL_0]] (%[[VAL_9]]) typeparams %[[VAL_3]] : (!fir.box<!fir.array<?x!fir.char<1,?>>>, i32, index) -> !fir.boxchar<1>
// CHECK: hlfir.yield %[[VAL_10]] : !fir.boxchar<1>
// CHECK: } cleanup {
// CHECK: %[[VAL_11:.*]] = fir.box_addr %[[VAL_4]] : (!fir.box<!fir.heap<!fir.array<?xi32>>>) -> !fir.heap<!fir.array<?xi32>>
// CHECK: fir.freemem %[[VAL_11]] : !fir.heap<!fir.array<?xi32>>
// CHECK: }
// CHECK: }
66 changes: 66 additions & 0 deletions flang/test/HLFIR/invalid.fir
View file
Original file line number Diff line number Diff line change
Expand Up @@ -541,3 +541,69 @@ func.func @bad_region_assign_2(%x: !fir.box<!fir.array<?xf32>>) {
}
return
}

// -----
func.func @bad_element_addr_1(%x: !fir.ref<!fir.array<20xf32>>) {
%c20 = arith.constant 20 : index
%vector_shape = fir.shape %c20 : (index) -> !fir.shape<1>
hlfir.region_assign {
hlfir.yield %x : !fir.ref<!fir.array<20xf32>>
} to {
// expected-error@+1 {{'hlfir.elemental_addr' op body must compute the address of a scalar entity}}
hlfir.elemental_addr %vector_shape : !fir.shape<1> {
^bb0(%i: index):
%c42 = arith.constant 42.0 : f32
hlfir.yield %c42 : f32
}
}
return
}

// -----
func.func @bad_element_addr_2(%x: !fir.ref<!fir.array<20xf32>>) {
%c20 = arith.constant 20 : index
%vector_shape = fir.shape %c20 : (index) -> !fir.shape<1>
hlfir.region_assign {
hlfir.yield %x : !fir.ref<!fir.array<20xf32>>
} to {
// expected-error@+1 {{'hlfir.elemental_addr' op body must compute the address of a scalar entity}}
hlfir.elemental_addr %vector_shape : !fir.shape<1> {
^bb0(%i: index):
hlfir.yield %x : !fir.ref<!fir.array<20xf32>>
}
}
return
}

// -----
func.func @bad_element_addr_3(%x: !fir.ref<!fir.array<20xf32>>) {
%c20 = arith.constant 20 : index
%vector_shape = fir.shape %c20 : (index) -> !fir.shape<1>
hlfir.region_assign {
hlfir.yield %x : !fir.ref<!fir.array<20xf32>>
} to {
// expected-error@+1 {{'hlfir.elemental_addr' op body region must be terminated by an hlfir.yield}}
hlfir.elemental_addr %vector_shape : !fir.shape<1> {
^bb0(%i: index):
%c42 = arith.constant 42.0 : f32
}
}
return
}

// -----
func.func @bad_element_addr_4(%x: !fir.ref<!fir.array<20xf32>>, %y: !fir.ref<!fir.array<20x20xf32>>) {
%c20 = arith.constant 20 : index
%vector_shape = fir.shape %c20 : (index) -> !fir.shape<1>
hlfir.region_assign {
hlfir.yield %x : !fir.ref<!fir.array<20xf32>>
} to {
// expected-error@+1 {{'hlfir.elemental_addr' op body number of indices must match shape rank}}
hlfir.elemental_addr %vector_shape : !fir.shape<1> {
^bb0(%i: index, %j: index):
%elt = hlfir.designate %y(%i, %j) : (!fir.ref<!fir.array<20x20xf32>>, index, index) -> !fir.ref<f32>
hlfir.yield %elt : !fir.ref<f32>
}
}
return
}

0 comments on commit bdc9914

Please sign in to comment.