diff --git a/flang/include/flang/Optimizer/Dialect/FIROps.td b/flang/include/flang/Optimizer/Dialect/FIROps.td --- a/flang/include/flang/Optimizer/Dialect/FIROps.td +++ b/flang/include/flang/Optimizer/Dialect/FIROps.td @@ -1484,6 +1484,67 @@ let verifier = "return ::verify(*this);"; } +def fir_ArrayModifyOp : fir_Op<"array_modify", [AttrSizedOperandSegments, + NoSideEffect]> { + let summary = "Get an address for an array value to modify it."; + + let description = [{ + Modify the value of an element in an array value through actions done + on the returned address. A new array value is also + returned where all element values of the input array are identical except + for the selected element which is the value after the modification done + on the element address. + + ```fortran + real :: a(n) + ... + ! Elemental user defined assignment from type(SomeType) to real. + a = value_of_some_type + ``` + + One can use `fir.array_modify` to update the (implied) value of `a(i)` + in an array expression as shown above. + + ```mlir + %s = fir.shape %n : (index) -> !fir.shape<1> + // Load the entire array 'a'. + %v = fir.array_load %a(%s) : (!fir.ref>, !fir.shape<1>) + -> !fir.array + // Update the value of one of the array value's elements with a user + // defined assignment from %rhs. + %new = fir.do_loop %i = ... (%inner = %v) { + %rhs = ... + %addr, %r = fir.array_modify %inner, %i, %j : (!fir.array, + index) -> fir.ref, !fir.array + fir.call @user_def_assign(%addr, %rhs) (fir.ref, + fir.ref>) -> () + fir.result %r : !fir.ref> + } + fir.array_merge_store %v, %new to %a : !fir.ref> + ``` + + An array value modification behaves as if a mapping function from the indices + to the new value has been added, replacing the previous mapping. These + mappings can be added to the ssa-value, but will not be materialized in + memory until the `fir.array_merge_store` is performed. + }]; + + let arguments = (ins + fir_SequenceType:$sequence, + Variadic:$indices, + Variadic:$typeparams + ); + + let results = (outs fir_ReferenceType, fir_SequenceType); + + let assemblyFormat = [{ + $sequence `,` $indices (`typeparams` $typeparams^)? attr-dict + `:` functional-type(operands, results) + }]; + + let verifier = [{ return ::verify(*this); }]; +} + def fir_ArrayMergeStoreOp : fir_Op<"array_merge_store", [AttrSizedOperandSegments]> { diff --git a/flang/lib/Optimizer/Dialect/FIROps.cpp b/flang/lib/Optimizer/Dialect/FIROps.cpp --- a/flang/lib/Optimizer/Dialect/FIROps.cpp +++ b/flang/lib/Optimizer/Dialect/FIROps.cpp @@ -496,6 +496,18 @@ return mlir::success(); } +//===----------------------------------------------------------------------===// +// ArrayModifyOp +//===----------------------------------------------------------------------===// + +static mlir::LogicalResult verify(fir::ArrayModifyOp op) { + auto arrTy = op.sequence().getType().cast(); + auto indSize = op.indices().size(); + if (indSize < arrTy.getDimension()) + return op.emitOpError("number of indices must match array dimension"); + return mlir::success(); +} + //===----------------------------------------------------------------------===// // BoxAddrOp //===----------------------------------------------------------------------===// diff --git a/flang/test/Fir/fir-ops.fir b/flang/test/Fir/fir-ops.fir --- a/flang/test/Fir/fir-ops.fir +++ b/flang/test/Fir/fir-ops.fir @@ -636,6 +636,16 @@ %av2 = fir.array_update %av1, %f, %i10, %j20 : (!fir.array, f32, index, index) -> !fir.array fir.array_merge_store %av1, %av2 to %arr1 : !fir.array, !fir.array, !fir.ref> + // CHECK: [[AV3:%.*]] = fir.array_load [[ARR1]]([[SHAPE]]) : (!fir.ref>, !fir.shapeshift<2>) -> !fir.array + // CHECK: [[FVAL2:%.*]] = fir.array_fetch [[AV3]], [[I10]], [[J20]] : (!fir.array, index, index) -> f32 + // CHECK: [[AV4:%.*]]:2 = fir.array_modify [[AV3]], [[I10]], [[J20]] : (!fir.array, index, index) -> (!fir.ref, !fir.array) + // CHECK: fir.store [[FVAL2]] to [[AV4]]#0 : !fir.ref + // CHECK: fir.array_merge_store [[AV3]], [[AV4]]#1 to [[ARR1]] : !fir.array, !fir.array, !fir.ref> + %av3 = fir.array_load %arr1(%s) : (!fir.ref>, !fir.shapeshift<2>) -> !fir.array + %f2 = fir.array_fetch %av3, %i10, %j20 : (!fir.array, index, index) -> f32 + %addr, %av4 = fir.array_modify %av3, %i10, %j20 : (!fir.array, index, index) -> (!fir.ref, !fir.array) + fir.store %f2 to %addr : !fir.ref + fir.array_merge_store %av3, %av4 to %arr1 : !fir.array, !fir.array, !fir.ref> return } diff --git a/flang/test/Fir/invalid.fir b/flang/test/Fir/invalid.fir --- a/flang/test/Fir/invalid.fir +++ b/flang/test/Fir/invalid.fir @@ -592,3 +592,17 @@ fir.array_merge_store %av2, %av2 to %arr1 : !fir.array, !fir.array, !fir.ref> return } + +// ----- + +func @bad_array_modify(%arr1 : !fir.ref>, %m : index, %n : index, %o : index, %p : index, %f : f32) { + %i10 = constant 10 : index + %j20 = constant 20 : index + %s = fir.shape_shift %m, %n, %o, %p : (index, index, index, index) -> !fir.shapeshift<2> + %av1 = fir.array_load %arr1(%s) : (!fir.ref>, !fir.shapeshift<2>) -> !fir.array + // expected-error@+1 {{'fir.array_modify' op number of indices must match array dimension}} + %addr, %av2 = fir.array_modify %av1, %i10 : (!fir.array, index) -> (!fir.ref, !fir.array) + fir.store %f to %addr : !fir.ref + fir.array_merge_store %av1, %av2 to %arr1 : !fir.array, !fir.array, !fir.ref> + return +}