119 lines · plain
1// RUN: fir-opt --lower-workdistribute %s | FileCheck %s2// Test lowering of workdistribute after fission on host device.3 4// CHECK-LABEL: func.func @x(5// CHECK: %[[VAL_0:.*]] = fir.alloca index {bindc_name = "lb"}6// CHECK: fir.store %[[ARG0:.*]] to %[[VAL_0]] : !fir.ref<index>7// CHECK: %[[VAL_1:.*]] = fir.alloca index {bindc_name = "ub"}8// CHECK: fir.store %[[ARG1:.*]] to %[[VAL_1]] : !fir.ref<index>9// CHECK: %[[VAL_2:.*]] = fir.alloca index {bindc_name = "step"}10// CHECK: fir.store %[[ARG2:.*]] to %[[VAL_2]] : !fir.ref<index>11// CHECK: %[[VAL_3:.*]] = omp.map.info var_ptr(%[[VAL_0]] : !fir.ref<index>, index) map_clauses(to) capture(ByRef) -> !fir.ref<index> {name = "lb"}12// CHECK: %[[VAL_4:.*]] = omp.map.info var_ptr(%[[VAL_1]] : !fir.ref<index>, index) map_clauses(to) capture(ByRef) -> !fir.ref<index> {name = "ub"}13// CHECK: %[[VAL_5:.*]] = omp.map.info var_ptr(%[[VAL_2]] : !fir.ref<index>, index) map_clauses(to) capture(ByRef) -> !fir.ref<index> {name = "step"}14// CHECK: %[[VAL_6:.*]] = omp.map.info var_ptr(%[[ARG3:.*]] : !fir.ref<index>, index) map_clauses(tofrom) capture(ByRef) -> !fir.ref<index> {name = "addr"}15// CHECK: %[[VAL_7:.*]] = omp.map.info var_ptr(%[[VAL_0]] : !fir.ref<index>, index) map_clauses(storage) capture(ByRef) -> !fir.ref<index> {name = "lb"}16// CHECK: %[[VAL_8:.*]] = omp.map.info var_ptr(%[[VAL_1]] : !fir.ref<index>, index) map_clauses(storage) capture(ByRef) -> !fir.ref<index> {name = "ub"}17// CHECK: %[[VAL_9:.*]] = omp.map.info var_ptr(%[[VAL_2]] : !fir.ref<index>, index) map_clauses(storage) capture(ByRef) -> !fir.ref<index> {name = "step"}18// CHECK: %[[VAL_10:.*]] = omp.map.info var_ptr(%[[ARG3]] : !fir.ref<index>, index) map_clauses(storage) capture(ByRef) -> !fir.ref<index> {name = "addr"}19// CHECK: omp.target_data map_entries(%[[VAL_3]], %[[VAL_4]], %[[VAL_5]], %[[VAL_6]] : !fir.ref<index>, !fir.ref<index>, !fir.ref<index>, !fir.ref<index>) {20// CHECK: %[[VAL_11:.*]] = fir.alloca index21// CHECK: %[[VAL_12:.*]] = omp.map.info var_ptr(%[[VAL_11]] : !fir.ref<index>, index) map_clauses(from) capture(ByRef) -> !fir.ref<index> {name = "__flang_workdistribute_from"}22// CHECK: %[[VAL_13:.*]] = omp.map.info var_ptr(%[[VAL_11]] : !fir.ref<index>, index) map_clauses(to) capture(ByRef) -> !fir.ref<index> {name = "__flang_workdistribute_to"}23// CHECK: %[[VAL_14:.*]] = fir.alloca index24// CHECK: %[[VAL_15:.*]] = omp.map.info var_ptr(%[[VAL_14]] : !fir.ref<index>, index) map_clauses(from) capture(ByRef) -> !fir.ref<index> {name = "__flang_workdistribute_from"}25// CHECK: %[[VAL_16:.*]] = omp.map.info var_ptr(%[[VAL_14]] : !fir.ref<index>, index) map_clauses(to) capture(ByRef) -> !fir.ref<index> {name = "__flang_workdistribute_to"}26// CHECK: %[[VAL_17:.*]] = fir.alloca index27// CHECK: %[[VAL_18:.*]] = omp.map.info var_ptr(%[[VAL_17]] : !fir.ref<index>, index) map_clauses(from) capture(ByRef) -> !fir.ref<index> {name = "__flang_workdistribute_from"}28// CHECK: %[[VAL_19:.*]] = omp.map.info var_ptr(%[[VAL_17]] : !fir.ref<index>, index) map_clauses(to) capture(ByRef) -> !fir.ref<index> {name = "__flang_workdistribute_to"}29// CHECK: %[[VAL_20:.*]] = fir.alloca !fir.heap<index>30// CHECK: %[[VAL_21:.*]] = omp.map.info var_ptr(%[[VAL_20]] : !fir.ref<!fir.heap<index>>, !fir.heap<index>) map_clauses(from) capture(ByRef) -> !fir.ref<!fir.heap<index>> {name = "__flang_workdistribute_from"}31// CHECK: %[[VAL_22:.*]] = omp.map.info var_ptr(%[[VAL_20]] : !fir.ref<!fir.heap<index>>, !fir.heap<index>) map_clauses(to) capture(ByRef) -> !fir.ref<!fir.heap<index>> {name = "__flang_workdistribute_to"}32// CHECK: %[[VAL_23:.*]] = llvm.mlir.constant(0 : i32) : i3233// CHECK: %[[VAL_24:.*]] = fir.load %[[VAL_0]] : !fir.ref<index>34// CHECK: %[[VAL_25:.*]] = fir.load %[[VAL_1]] : !fir.ref<index>35// CHECK: %[[VAL_26:.*]] = fir.load %[[VAL_2]] : !fir.ref<index>36// CHECK: %[[VAL_27:.*]] = arith.constant 1 : index37// CHECK: %[[VAL_28:.*]] = arith.addi %[[VAL_25]], %[[VAL_25]] : index38// CHECK: %[[VAL_29:.*]] = omp.target_allocmem %[[VAL_23]] : i32, index, %[[VAL_27]] {uniq_name = "dev_buf"}39// CHECK: %[[VAL_30:.*]] = fir.convert %[[VAL_29]] : (i64) -> !fir.heap<index>40// CHECK: fir.store %[[VAL_24]] to %[[VAL_11]] : !fir.ref<index>41// CHECK: fir.store %[[VAL_25]] to %[[VAL_14]] : !fir.ref<index>42// CHECK: fir.store %[[VAL_26]] to %[[VAL_17]] : !fir.ref<index>43// CHECK: fir.store %[[VAL_30]] to %[[VAL_20]] : !fir.ref<!fir.heap<index>>44// CHECK: omp.target map_entries(%[[VAL_7]] -> %[[VAL_31:.*]], %[[VAL_8]] -> %[[VAL_32:.*]], %[[VAL_9]] -> %[[VAL_33:.*]], %[[VAL_10]] -> %[[VAL_34:.*]], %[[VAL_13]] -> %[[VAL_35:.*]], %[[VAL_16]] -> %[[VAL_36:.*]], %[[VAL_19]] -> %[[VAL_37:.*]], %[[VAL_22]] -> %[[VAL_38:.*]] : !fir.ref<index>, !fir.ref<index>, !fir.ref<index>, !fir.ref<index>, !fir.ref<index>, !fir.ref<index>, !fir.ref<index>, !fir.ref<!fir.heap<index>>) {45// CHECK: %[[VAL_39:.*]] = fir.load %[[VAL_35]] : !fir.ref<index>46// CHECK: %[[VAL_40:.*]] = fir.load %[[VAL_36]] : !fir.ref<index>47// CHECK: %[[VAL_41:.*]] = fir.load %[[VAL_37]] : !fir.ref<index>48// CHECK: %[[VAL_42:.*]] = fir.load %[[VAL_38]] : !fir.ref<!fir.heap<index>>49// CHECK: %[[VAL_43:.*]] = arith.addi %[[VAL_40]], %[[VAL_40]] : index50// CHECK: omp.teams {51// CHECK: omp.parallel {52// CHECK: omp.distribute {53// CHECK: omp.wsloop {54// CHECK: omp.loop_nest (%[[VAL_44:.*]]) : index = (%[[VAL_39]]) to (%[[VAL_40]]) inclusive step (%[[VAL_41]]) {55// CHECK: fir.store %[[VAL_43]] to %[[VAL_42]] : !fir.heap<index>56// CHECK: omp.yield57// CHECK: }58// CHECK: } {omp.composite}59// CHECK: } {omp.composite}60// CHECK: omp.terminator61// CHECK: } {omp.composite}62// CHECK: omp.terminator63// CHECK: }64// CHECK: omp.terminator65// CHECK: }66// CHECK: %[[VAL_45:.*]] = llvm.mlir.constant(0 : i32) : i3267// CHECK: %[[VAL_46:.*]] = fir.load %[[VAL_11]] : !fir.ref<index>68// CHECK: %[[VAL_47:.*]] = fir.load %[[VAL_14]] : !fir.ref<index>69// CHECK: %[[VAL_48:.*]] = fir.load %[[VAL_17]] : !fir.ref<index>70// CHECK: %[[VAL_49:.*]] = fir.load %[[VAL_20]] : !fir.ref<!fir.heap<index>>71// CHECK: %[[VAL_50:.*]] = arith.addi %[[VAL_47]], %[[VAL_47]] : index72// CHECK: fir.store %[[VAL_46]] to %[[VAL_49]] : !fir.heap<index>73// CHECK: %[[VAL_51:.*]] = fir.convert %[[VAL_49]] : (!fir.heap<index>) -> i6474// CHECK: omp.target_freemem %[[VAL_45]], %[[VAL_51]] : i32, i6475// CHECK: omp.terminator76// CHECK: }77// CHECK: return78// CHECK: }79 80 81module attributes {llvm.target_triple = "amdgcn-amd-amdhsa", omp.is_gpu = true, omp.is_target_device = true} {82func.func @x(%lb : index, %ub : index, %step : index, %addr : !fir.ref<index>) {83 %lb_ref = fir.alloca index {bindc_name = "lb"}84 fir.store %lb to %lb_ref : !fir.ref<index>85 %ub_ref = fir.alloca index {bindc_name = "ub"}86 fir.store %ub to %ub_ref : !fir.ref<index>87 %step_ref = fir.alloca index {bindc_name = "step"}88 fir.store %step to %step_ref : !fir.ref<index>89 90 %lb_map = omp.map.info var_ptr(%lb_ref : !fir.ref<index>, index) map_clauses(to) capture(ByRef) -> !fir.ref<index> {name = "lb"}91 %ub_map = omp.map.info var_ptr(%ub_ref : !fir.ref<index>, index) map_clauses(to) capture(ByRef) -> !fir.ref<index> {name = "ub"}92 %step_map = omp.map.info var_ptr(%step_ref : !fir.ref<index>, index) map_clauses(to) capture(ByRef) -> !fir.ref<index> {name = "step"}93 %addr_map = omp.map.info var_ptr(%addr : !fir.ref<index>, index) map_clauses(tofrom) capture(ByRef) -> !fir.ref<index> {name = "addr"}94 95 omp.target map_entries(%lb_map -> %ARG0, %ub_map -> %ARG1, %step_map -> %ARG2, %addr_map -> %ARG3 : !fir.ref<index>, !fir.ref<index>, !fir.ref<index>, !fir.ref<index>) {96 %lb_val = fir.load %ARG0 : !fir.ref<index>97 %ub_val = fir.load %ARG1 : !fir.ref<index>98 %step_val = fir.load %ARG2 : !fir.ref<index>99 %one = arith.constant 1 : index100 101 %20 = arith.addi %ub_val, %ub_val : index102 omp.teams {103 omp.workdistribute {104 %dev_mem = fir.allocmem index, %one {uniq_name = "dev_buf"}105 fir.do_loop %iv = %lb_val to %ub_val step %step_val unordered {106 fir.store %20 to %dev_mem : !fir.heap<index>107 }108 fir.store %lb_val to %dev_mem : !fir.heap<index>109 fir.freemem %dev_mem : !fir.heap<index>110 omp.terminator111 }112 omp.terminator113 }114 omp.terminator115 }116 return117}118}119