diff --git a/flang/include/flang/Optimizer/HLFIR/HLFIROps.td b/flang/include/flang/Optimizer/HLFIR/HLFIROps.td --- a/flang/include/flang/Optimizer/HLFIR/HLFIROps.td +++ b/flang/include/flang/Optimizer/HLFIR/HLFIROps.td @@ -331,5 +331,76 @@ let assemblyFormat = "$val attr-dict `:` type($val)"; } +def hlfir_ElementalOp : hlfir_Op<"elemental", []> { + let summary = "elemental expression"; + let description = [{ + Represent an elemental expression as a function of the indices. + This operation contain a region whose block arguments are one + based indices iterating over the elemental expression shape. + Given these indices, the element value for the given iteration + can be computed in the region and yielded with the hlfir.yield_element + operation. + + The shape and typeparams operands represent the extents and type + parameters of the resulting array value. + + + Example: Y + X, with Integer :: X(10, 20), Y(10,20) + ``` + %0 = fir.shape %c10, %c20 : (index, index) -> !fir.shape<2> + %5 = hlfir.elemental %0 : (!fir.shape<2>) -> !hlfir.expr<10x20xi32> { + ^bb0(%i: index, %j: index): + %6 = hlfir.designate %x (%i, %j) : (!fir.ref>, index, index) -> !fir.ref + %7 = hlfir.designate %y (%i, %j) : (!fir.ref>, index, index) -> !fir.ref + %8 = fir.load %6 : !fir.ref + %9 = fir.load %7 : !fir.ref + %10 = arith.addi %8, %9 : i32 + hlfir.yield_element %10 : i32 + } + ``` + }]; + + let arguments = (ins + AnyShapeType:$shape, + Variadic:$typeparams + ); + + let results = (outs hlfir_ExprType); + let regions = (region SizedRegion<1>:$region); + + let assemblyFormat = [{ + $shape (`typeparams` $typeparams^)? + attr-dict `:` functional-type(operands, results) + $region + }]; + + let extraClassDeclaration = [{ + mlir::Block *getBody() { return &getRegion().front(); } + + // Get the indices iterating over the shape. + mlir::Block::BlockArgListType getIndices() { + return getBody()->getArguments(); + } + }]; + + let skipDefaultBuilders = 1; + let builders = [ + OpBuilder<(ins "mlir::Type":$result_type, "mlir::Value":$shape, + CArg<"mlir::ValueRange", "{}">:$typeparams)>]; + +} + +def hlfir_YieldElementOp : hlfir_Op<"yield_element", [Terminator, HasParent<"ElementalOp">]> { + let summary = "Yield the elemental value in an ElementalOp"; + let description = [{ + Yield the element value of the current elemental expression iteration + in an hlfir.elemental region. See hlfir.elemental description for an + example. + }]; + + let arguments = (ins AnyType:$element_value); + + let assemblyFormat = "$element_value attr-dict `:` type($element_value)"; +} #endif // FORTRAN_DIALECT_HLFIR_OPS diff --git a/flang/lib/Optimizer/HLFIR/IR/HLFIROps.cpp b/flang/lib/Optimizer/HLFIR/IR/HLFIROps.cpp --- a/flang/lib/Optimizer/HLFIR/IR/HLFIROps.cpp +++ b/flang/lib/Optimizer/HLFIR/IR/HLFIROps.cpp @@ -433,5 +433,26 @@ return build(builder, result, resultType, var); } +//===----------------------------------------------------------------------===// +// ElementalOp +//===----------------------------------------------------------------------===// + +void hlfir::ElementalOp::build(mlir::OpBuilder &builder, + mlir::OperationState &odsState, + mlir::Type resultType, mlir::Value shape, + mlir::ValueRange typeparams) { + odsState.addOperands(shape); + odsState.addOperands(typeparams); + odsState.addTypes(resultType); + mlir::Region *bodyRegion = odsState.addRegion(); + bodyRegion->push_back(new mlir::Block{}); + if (auto exprType = resultType.dyn_cast()) { + unsigned dim = exprType.getRank(); + mlir::Type indexType = builder.getIndexType(); + for (unsigned d = 0; d < dim; ++d) + bodyRegion->front().addArgument(indexType, odsState.location); + } +} + #define GET_OP_CLASSES #include "flang/Optimizer/HLFIR/HLFIROps.cpp.inc" diff --git a/flang/test/HLFIR/elemental.fir b/flang/test/HLFIR/elemental.fir new file mode 100644 --- /dev/null +++ b/flang/test/HLFIR/elemental.fir @@ -0,0 +1,78 @@ +// Test hlfir.elemental and hlfir.yield_element operation parse, verify +// (no errors), and unparse. + +// RUN: fir-opt %s | fir-opt | FileCheck %s + +func.func @numeric_type(%x: !fir.ref>, %y: !fir.ref>) { + %c10 = arith.constant 10 : index + %c20 = arith.constant 20 : index + %0 = fir.shape %c10, %c20 : (index, index) -> !fir.shape<2> + %3 = hlfir.elemental %0 : (!fir.shape<2>) -> !hlfir.expr<10x20xi32> { + ^bb0(%i: index, %j: index): + %4 = hlfir.designate %x (%i, %j) : (!fir.ref>, index, index) -> !fir.ref + %5 = hlfir.designate %y (%i, %j) : (!fir.ref>, index, index) -> !fir.ref + %6 = fir.load %4 : !fir.ref + %7 = fir.load %5 : !fir.ref + %8 = arith.addi %6, %7 : i32 + hlfir.yield_element %8 : i32 + } + return +} +// CHECK-LABEL: func.func @numeric_type( +// CHECK-SAME: %[[VAL_0:[^:]*]]: !fir.ref>, +// CHECK-SAME: %[[VAL_1:[^:]*]]: !fir.ref> +// CHECK: %[[VAL_4:.*]] = fir.shape +// CHECK: %[[VAL_5:.*]] = hlfir.elemental %[[VAL_4]] : (!fir.shape<2>) -> !hlfir.expr<10x20xi32> { +// CHECK: ^bb0(%[[VAL_6:.*]]: index, %[[VAL_7:.*]]: index): +// CHECK: %[[VAL_8:.*]] = hlfir.designate %[[VAL_0]] (%[[VAL_6]], %[[VAL_7]]) : (!fir.ref>, index, index) -> !fir.ref +// CHECK: %[[VAL_9:.*]] = hlfir.designate %[[VAL_1]] (%[[VAL_6]], %[[VAL_7]]) : (!fir.ref>, index, index) -> !fir.ref +// CHECK: %[[VAL_10:.*]] = fir.load %[[VAL_8]] : !fir.ref +// CHECK: %[[VAL_11:.*]] = fir.load %[[VAL_9]] : !fir.ref +// CHECK: %[[VAL_12:.*]] = arith.addi %[[VAL_10]], %[[VAL_11]] : i32 +// CHECK: hlfir.yield_element %[[VAL_12]] : i32 +// CHECK: } + +func.func @char_type(%x: !fir.box>>, %n: index, %l : index, %l0 :index) { + %0 = fir.shape %n : (index) -> !fir.shape<1> + %3 = hlfir.elemental %0 typeparams %l : (!fir.shape<1>, index) -> !hlfir.expr> { + ^bb0(%i: index): + %4 = hlfir.designate %x (%i) typeparams %l0 : (!fir.box>>, index, index) -> !fir.boxchar<1> + %6 = hlfir.concat %4, %4 len %l : (!fir.boxchar<1>, !fir.boxchar<1>, index) -> !hlfir.expr> + hlfir.yield_element %6 : !hlfir.expr> + } + return +} +// CHECK-LABEL: func.func @char_type( +// CHECK-SAME: %[[VAL_0:[^:]*]]: !fir.box>>, +// CHECK-SAME: %[[VAL_1:[^:]*]]: index, +// CHECK-SAME: %[[VAL_2:[^:]*]]: index, +// CHECK-SAME: %[[VAL_3:[^:]*]]: index) { +// CHECK: %[[VAL_4:.*]] = fir.shape %[[VAL_1]] : (index) -> !fir.shape<1> +// CHECK: %[[VAL_5:.*]] = hlfir.elemental %[[VAL_4]] typeparams %[[VAL_2]] : (!fir.shape<1>, index) -> !hlfir.expr> { +// CHECK: ^bb0(%[[VAL_6:.*]]: index): +// CHECK: %[[VAL_7:.*]] = hlfir.designate %[[VAL_0]] (%[[VAL_6]]) typeparams %[[VAL_3]] : (!fir.box>>, index, index) -> !fir.boxchar<1> +// CHECK: %[[VAL_8:.*]] = hlfir.concat %[[VAL_7]], %[[VAL_7]] len %[[VAL_2]] : (!fir.boxchar<1>, !fir.boxchar<1>, index) -> !hlfir.expr> +// CHECK: hlfir.yield_element %[[VAL_8]] : !hlfir.expr> +// CHECK: } + +!pdt = !fir.type +func.func @parametrized_derived_transpose(%x: !fir.box>, %n: index, %m: index, %l: i32) { + %0 = fir.shape %m, %n : (index, index) -> !fir.shape<2> + %3 = hlfir.elemental %0 typeparams %l : (!fir.shape<2>, i32) -> !hlfir.expr { + ^bb0(%j: index, %i: index): + %4 = hlfir.designate %x (%j, %i) typeparams %l : (!fir.box>, index, index, i32) -> !fir.box + %5 = hlfir.as_expr %4 : (!fir.box) -> !hlfir.expr + hlfir.yield_element %5 : !hlfir.expr + } + return +} +// CHECK-LABEL: func.func @parametrized_derived_transpose( +// CHECK-SAME: %[[VAL_0:.*]]: !fir.box +// CHECK-SAME: %[[VAL_3:[^:]*]]: i32 +// CHECK: %[[VAL_4:.*]] = fir.shape +// CHECK: %[[VAL_5:.*]] = hlfir.elemental %[[VAL_4]] typeparams %[[VAL_3]] : (!fir.shape<2>, i32) -> !hlfir.expr> { +// CHECK: ^bb0(%[[VAL_6:.*]]: index, %[[VAL_7:.*]]: index): +// CHECK: %[[VAL_8:.*]] = hlfir.designate %[[VAL_0]] (%[[VAL_6]], %[[VAL_7]]) typeparams %[[VAL_3]] : (!fir.box>>, index, index, i32) -> !fir.box> +// CHECK: %[[VAL_9:.*]] = hlfir.as_expr %[[VAL_8]] : (!fir.box>) -> !hlfir.expr> +// CHECK: hlfir.yield_element %[[VAL_9]] : !hlfir.expr> +// CHECK: }