[flang-commits] [flang] [flang] Expand SUM(DIM=CONSTANT) into an hlfir.elemental. (PR #118556)
Tom Eccles via flang-commits
flang-commits at lists.llvm.org
Wed Dec 4 03:12:32 PST 2024
================
@@ -0,0 +1,361 @@
+// RUN: fir-opt --simplify-hlfir-intrinsics %s | FileCheck %s
+
+// box with known extents
+func.func @sum_box_known_extents(%arg0: !fir.box<!fir.array<2x3xi32>>) {
+ %cst = arith.constant 2 : i32
+ %res = hlfir.sum %arg0 dim %cst : (!fir.box<!fir.array<2x3xi32>>, i32) -> !hlfir.expr<2xi32>
+ return
+}
+// CHECK-LABEL: func.func @sum_box_known_extents(
+// CHECK-SAME: %[[VAL_0:.*]]: !fir.box<!fir.array<2x3xi32>>) {
+// CHECK: %[[VAL_1:.*]] = arith.constant 2 : i32
+// CHECK: %[[VAL_2:.*]] = arith.constant 2 : index
+// CHECK: %[[VAL_3:.*]] = arith.constant 3 : index
+// CHECK: %[[VAL_4:.*]] = fir.shape %[[VAL_2]] : (index) -> !fir.shape<1>
+// CHECK: %[[VAL_5:.*]] = hlfir.elemental %[[VAL_4]] unordered : (!fir.shape<1>) -> !hlfir.expr<2xi32> {
+// CHECK: ^bb0(%[[VAL_6:.*]]: index):
+// CHECK: %[[VAL_7:.*]] = arith.constant 1 : index
+// CHECK: %[[VAL_8:.*]] = arith.constant 0 : i32
+// CHECK: %[[VAL_9:.*]] = fir.do_loop %[[VAL_10:.*]] = %[[VAL_7]] to %[[VAL_3]] step %[[VAL_7]] unordered iter_args(%[[VAL_11:.*]] = %[[VAL_8]]) -> (i32) {
+// CHECK: %[[VAL_12:.*]] = arith.constant 0 : index
+// CHECK: %[[VAL_13:.*]]:3 = fir.box_dims %[[VAL_0]], %[[VAL_12]] : (!fir.box<!fir.array<2x3xi32>>, index) -> (index, index, index)
+// CHECK: %[[VAL_14:.*]] = arith.constant 1 : index
+// CHECK: %[[VAL_15:.*]]:3 = fir.box_dims %[[VAL_0]], %[[VAL_14]] : (!fir.box<!fir.array<2x3xi32>>, index) -> (index, index, index)
+// CHECK: %[[VAL_16:.*]] = arith.constant 1 : index
+// CHECK: %[[VAL_17:.*]] = arith.subi %[[VAL_13]]#0, %[[VAL_16]] : index
+// CHECK: %[[VAL_18:.*]] = arith.addi %[[VAL_6]], %[[VAL_17]] : index
+// CHECK: %[[VAL_19:.*]] = arith.subi %[[VAL_15]]#0, %[[VAL_16]] : index
+// CHECK: %[[VAL_20:.*]] = arith.addi %[[VAL_10]], %[[VAL_19]] : index
+// CHECK: %[[VAL_21:.*]] = hlfir.designate %[[VAL_0]] (%[[VAL_18]], %[[VAL_20]]) : (!fir.box<!fir.array<2x3xi32>>, index, index) -> !fir.ref<i32>
+// CHECK: %[[VAL_22:.*]] = fir.load %[[VAL_21]] : !fir.ref<i32>
+// CHECK: %[[VAL_23:.*]] = arith.addi %[[VAL_11]], %[[VAL_22]] : i32
+// CHECK: fir.result %[[VAL_23]] : i32
+// CHECK: }
+// CHECK: hlfir.yield_element %[[VAL_9]] : i32
+// CHECK: }
+// CHECK: return
+// CHECK: }
+
+// expr with known extents
+func.func @sum_expr_known_extents(%arg0: !hlfir.expr<2x3xi32>) {
+ %cst = arith.constant 1 : i32
+ %res = hlfir.sum %arg0 dim %cst : (!hlfir.expr<2x3xi32>, i32) -> !hlfir.expr<3xi32>
+ return
+}
+// CHECK-LABEL: func.func @sum_expr_known_extents(
+// CHECK-SAME: %[[VAL_0:.*]]: !hlfir.expr<2x3xi32>) {
+// CHECK: %[[VAL_1:.*]] = arith.constant 1 : i32
+// CHECK: %[[VAL_2:.*]] = arith.constant 2 : index
+// CHECK: %[[VAL_3:.*]] = arith.constant 3 : index
+// CHECK: %[[VAL_4:.*]] = fir.shape %[[VAL_3]] : (index) -> !fir.shape<1>
+// CHECK: %[[VAL_5:.*]] = hlfir.elemental %[[VAL_4]] unordered : (!fir.shape<1>) -> !hlfir.expr<3xi32> {
+// CHECK: ^bb0(%[[VAL_6:.*]]: index):
+// CHECK: %[[VAL_7:.*]] = arith.constant 1 : index
+// CHECK: %[[VAL_8:.*]] = arith.constant 0 : i32
+// CHECK: %[[VAL_9:.*]] = fir.do_loop %[[VAL_10:.*]] = %[[VAL_7]] to %[[VAL_2]] step %[[VAL_7]] unordered iter_args(%[[VAL_11:.*]] = %[[VAL_8]]) -> (i32) {
+// CHECK: %[[VAL_12:.*]] = hlfir.apply %[[VAL_0]], %[[VAL_10]], %[[VAL_6]] : (!hlfir.expr<2x3xi32>, index, index) -> i32
+// CHECK: %[[VAL_13:.*]] = arith.addi %[[VAL_11]], %[[VAL_12]] : i32
+// CHECK: fir.result %[[VAL_13]] : i32
+// CHECK: }
+// CHECK: hlfir.yield_element %[[VAL_9]] : i32
+// CHECK: }
+// CHECK: return
+// CHECK: }
+
+// box with unknown extent
+func.func @sum_box_unknown_extent1(%arg0: !fir.box<!fir.array<?x3xcomplex<f64>>>) {
+ %cst = arith.constant 1 : i32
+ %res = hlfir.sum %arg0 dim %cst : (!fir.box<!fir.array<?x3xcomplex<f64>>>, i32) -> !hlfir.expr<3xcomplex<f64>>
+ return
+}
+// CHECK-LABEL: func.func @sum_box_unknown_extent1(
+// CHECK-SAME: %[[VAL_0:.*]]: !fir.box<!fir.array<?x3xcomplex<f64>>>) {
+// CHECK: %[[VAL_1:.*]] = arith.constant 1 : i32
+// CHECK: %[[VAL_2:.*]] = arith.constant 0 : index
+// CHECK: %[[VAL_3:.*]]:3 = fir.box_dims %[[VAL_0]], %[[VAL_2]] : (!fir.box<!fir.array<?x3xcomplex<f64>>>, index) -> (index, index, index)
+// CHECK: %[[VAL_4:.*]] = arith.constant 3 : index
+// CHECK: %[[VAL_5:.*]] = fir.shape %[[VAL_4]] : (index) -> !fir.shape<1>
+// CHECK: %[[VAL_6:.*]] = hlfir.elemental %[[VAL_5]] unordered : (!fir.shape<1>) -> !hlfir.expr<3xcomplex<f64>> {
+// CHECK: ^bb0(%[[VAL_7:.*]]: index):
+// CHECK: %[[VAL_8:.*]] = arith.constant 1 : index
+// CHECK: %[[VAL_9:.*]] = arith.constant 0.000000e+00 : f64
+// CHECK: %[[VAL_10:.*]] = fir.undefined complex<f64>
+// CHECK: %[[VAL_11:.*]] = fir.insert_value %[[VAL_10]], %[[VAL_9]], [0 : index] : (complex<f64>, f64) -> complex<f64>
+// CHECK: %[[VAL_12:.*]] = fir.insert_value %[[VAL_11]], %[[VAL_9]], [1 : index] : (complex<f64>, f64) -> complex<f64>
+// CHECK: %[[VAL_13:.*]] = fir.do_loop %[[VAL_14:.*]] = %[[VAL_8]] to %[[VAL_3]]#1 step %[[VAL_8]] iter_args(%[[VAL_15:.*]] = %[[VAL_12]]) -> (complex<f64>) {
+// CHECK: %[[VAL_16:.*]] = arith.constant 0 : index
+// CHECK: %[[VAL_17:.*]]:3 = fir.box_dims %[[VAL_0]], %[[VAL_16]] : (!fir.box<!fir.array<?x3xcomplex<f64>>>, index) -> (index, index, index)
+// CHECK: %[[VAL_18:.*]] = arith.constant 1 : index
+// CHECK: %[[VAL_19:.*]]:3 = fir.box_dims %[[VAL_0]], %[[VAL_18]] : (!fir.box<!fir.array<?x3xcomplex<f64>>>, index) -> (index, index, index)
+// CHECK: %[[VAL_20:.*]] = arith.constant 1 : index
+// CHECK: %[[VAL_21:.*]] = arith.subi %[[VAL_17]]#0, %[[VAL_20]] : index
+// CHECK: %[[VAL_22:.*]] = arith.addi %[[VAL_14]], %[[VAL_21]] : index
+// CHECK: %[[VAL_23:.*]] = arith.subi %[[VAL_19]]#0, %[[VAL_20]] : index
+// CHECK: %[[VAL_24:.*]] = arith.addi %[[VAL_7]], %[[VAL_23]] : index
+// CHECK: %[[VAL_25:.*]] = hlfir.designate %[[VAL_0]] (%[[VAL_22]], %[[VAL_24]]) : (!fir.box<!fir.array<?x3xcomplex<f64>>>, index, index) -> !fir.ref<complex<f64>>
+// CHECK: %[[VAL_26:.*]] = fir.load %[[VAL_25]] : !fir.ref<complex<f64>>
+// CHECK: %[[VAL_27:.*]] = fir.addc %[[VAL_15]], %[[VAL_26]] : complex<f64>
+// CHECK: fir.result %[[VAL_27]] : complex<f64>
+// CHECK: }
+// CHECK: hlfir.yield_element %[[VAL_13]] : complex<f64>
+// CHECK: }
+// CHECK: return
+// CHECK: }
+
+func.func @sum_box_unknown_extent2(%arg0: !fir.box<!fir.array<?x3xcomplex<f64>>>) {
+ %cst = arith.constant 2 : i32
+ %res = hlfir.sum %arg0 dim %cst : (!fir.box<!fir.array<?x3xcomplex<f64>>>, i32) -> !hlfir.expr<?xcomplex<f64>>
+ return
+}
+// CHECK-LABEL: func.func @sum_box_unknown_extent2(
+// CHECK-SAME: %[[VAL_0:.*]]: !fir.box<!fir.array<?x3xcomplex<f64>>>) {
+// CHECK: %[[VAL_1:.*]] = arith.constant 2 : i32
+// CHECK: %[[VAL_2:.*]] = arith.constant 0 : index
+// CHECK: %[[VAL_3:.*]]:3 = fir.box_dims %[[VAL_0]], %[[VAL_2]] : (!fir.box<!fir.array<?x3xcomplex<f64>>>, index) -> (index, index, index)
+// CHECK: %[[VAL_4:.*]] = arith.constant 3 : index
+// CHECK: %[[VAL_5:.*]] = fir.shape %[[VAL_3]]#1 : (index) -> !fir.shape<1>
+// CHECK: %[[VAL_6:.*]] = hlfir.elemental %[[VAL_5]] unordered : (!fir.shape<1>) -> !hlfir.expr<?xcomplex<f64>> {
+// CHECK: ^bb0(%[[VAL_7:.*]]: index):
+// CHECK: %[[VAL_8:.*]] = arith.constant 1 : index
+// CHECK: %[[VAL_9:.*]] = arith.constant 0.000000e+00 : f64
+// CHECK: %[[VAL_10:.*]] = fir.undefined complex<f64>
+// CHECK: %[[VAL_11:.*]] = fir.insert_value %[[VAL_10]], %[[VAL_9]], [0 : index] : (complex<f64>, f64) -> complex<f64>
+// CHECK: %[[VAL_12:.*]] = fir.insert_value %[[VAL_11]], %[[VAL_9]], [1 : index] : (complex<f64>, f64) -> complex<f64>
+// CHECK: %[[VAL_13:.*]] = fir.do_loop %[[VAL_14:.*]] = %[[VAL_8]] to %[[VAL_4]] step %[[VAL_8]] iter_args(%[[VAL_15:.*]] = %[[VAL_12]]) -> (complex<f64>) {
+// CHECK: %[[VAL_16:.*]] = arith.constant 0 : index
+// CHECK: %[[VAL_17:.*]]:3 = fir.box_dims %[[VAL_0]], %[[VAL_16]] : (!fir.box<!fir.array<?x3xcomplex<f64>>>, index) -> (index, index, index)
+// CHECK: %[[VAL_18:.*]] = arith.constant 1 : index
+// CHECK: %[[VAL_19:.*]]:3 = fir.box_dims %[[VAL_0]], %[[VAL_18]] : (!fir.box<!fir.array<?x3xcomplex<f64>>>, index) -> (index, index, index)
+// CHECK: %[[VAL_20:.*]] = arith.constant 1 : index
+// CHECK: %[[VAL_21:.*]] = arith.subi %[[VAL_17]]#0, %[[VAL_20]] : index
+// CHECK: %[[VAL_22:.*]] = arith.addi %[[VAL_7]], %[[VAL_21]] : index
+// CHECK: %[[VAL_23:.*]] = arith.subi %[[VAL_19]]#0, %[[VAL_20]] : index
+// CHECK: %[[VAL_24:.*]] = arith.addi %[[VAL_14]], %[[VAL_23]] : index
+// CHECK: %[[VAL_25:.*]] = hlfir.designate %[[VAL_0]] (%[[VAL_22]], %[[VAL_24]]) : (!fir.box<!fir.array<?x3xcomplex<f64>>>, index, index) -> !fir.ref<complex<f64>>
+// CHECK: %[[VAL_26:.*]] = fir.load %[[VAL_25]] : !fir.ref<complex<f64>>
+// CHECK: %[[VAL_27:.*]] = fir.addc %[[VAL_15]], %[[VAL_26]] : complex<f64>
+// CHECK: fir.result %[[VAL_27]] : complex<f64>
+// CHECK: }
+// CHECK: hlfir.yield_element %[[VAL_13]] : complex<f64>
+// CHECK: }
+// CHECK: return
+// CHECK: }
+
+// expr with unknown extent
+func.func @sum_expr_unkwnonw_extent1(%arg0: !hlfir.expr<?x3xf32>) {
----------------
tblah wrote:
```suggestion
func.func @sum_expr_unknown_extent1(%arg0: !hlfir.expr<?x3xf32>) {
```
And see below
https://github.com/llvm/llvm-project/pull/118556
More information about the flang-commits
mailing list