1// RUN: mlir-opt %s -allow-unregistered-dialect -test-decompose-affine-ops -split-input-file -loop-invariant-code-motion -cse | FileCheck %s 2 3// CHECK-DAG: #[[$c42:.*]] = affine_map<() -> (42)> 4// CHECK-DAG: #[[$div32mod4:.*]] = affine_map<()[s0] -> ((s0 floordiv 32) mod 4)> 5// CHECK-DAG: #[[$add:.*]] = affine_map<()[s0, s1] -> (s0 + s1)> 6 7// CHECK-LABEL: func.func @simple_test_1 8// CHECK-SAME: %[[I0:[0-9a-zA-Z]+]]: index, 9// CHECK-SAME: %[[I1:[0-9a-zA-Z]+]]: index, 10// CHECK-SAME: %[[I2:[0-9a-zA-Z]+]]: index, 11// CHECK-SAME: %[[LB:[0-9a-zA-Z]+]]: index, 12// CHECK-SAME: %[[UB:[0-9a-zA-Z]+]]: index, 13// CHECK-SAME: %[[STEP:[0-9a-zA-Z]+]]: index 14func.func @simple_test_1(%0: index, %1: index, %2: index, %lb: index, %ub: index, %step: index) { 15 // CHECK: %[[c42:.*]] = affine.apply #[[$c42]]() 16 // CHECK: %[[R1:.*]] = affine.apply #[[$div32mod4]]()[%[[I1]]] 17 // CHECK: %[[a:.*]] = affine.apply #[[$add]]()[%[[c42]], %[[R1]]] 18 %a = affine.apply affine_map<(d0) -> ((d0 floordiv 32) mod 4 + 42)>(%1) 19 20 // CHECK: "some_side_effecting_consumer"(%[[a]]) : (index) -> () 21 "some_side_effecting_consumer"(%a) : (index) -> () 22 return 23} 24 25// ----- 26 27// CHECK-DAG: #[[$c42:.*]] = affine_map<() -> (42)> 28// CHECK-DAG: #[[$id:.*]] = affine_map<()[s0] -> (s0)> 29// CHECK-DAG: #[[$add:.*]] = affine_map<()[s0, s1] -> (s0 + s1)> 30// CHECK-DAG: #[[$div32div4timesm4:.*]] = affine_map<()[s0] -> (((s0 floordiv 32) floordiv 4) * -4)> 31// CHECK-DAG: #[[$div32:.*]] = affine_map<()[s0] -> (s0 floordiv 32)> 32 33// CHECK-LABEL: func.func @simple_test_2 34// CHECK-SAME: %[[I0:[0-9a-zA-Z]+]]: index, 35// CHECK-SAME: %[[I1:[0-9a-zA-Z]+]]: index, 36// CHECK-SAME: %[[I2:[0-9a-zA-Z]+]]: index, 37// CHECK-SAME: %[[LB:[0-9a-zA-Z]+]]: index, 38// CHECK-SAME: %[[UB:[0-9a-zA-Z]+]]: index, 39// CHECK-SAME: %[[STEP:[0-9a-zA-Z]+]]: index 40func.func @simple_test_2(%0: index, %1: index, %2: index, %lb: index, %ub: index, %step: index) { 41 // CHECK: %[[c42:.*]] = affine.apply #[[$c42]]() 42 // CHECK: scf.for %[[i:.*]] = 43 scf.for %i = %lb to %ub step %step { 44 // CHECK: %[[R1:.*]] = affine.apply #[[$id]]()[%[[i]]] 45 // CHECK: %[[R2:.*]] = affine.apply #[[$add]]()[%[[c42]], %[[R1]]] 46 // CHECK: scf.for %[[j:.*]] = 47 scf.for %j = %lb to %ub step %step { 48 // CHECK: %[[R3:.*]] = affine.apply #[[$div32div4timesm4]]()[%[[j]]] 49 // CHECK: %[[R4:.*]] = affine.apply #[[$add]]()[%[[R2]], %[[R3]]] 50 // CHECK: %[[R5:.*]] = affine.apply #[[$div32]]()[%[[j]]] 51 // CHECK: %[[a:.*]] = affine.apply #[[$add]]()[%[[R4]], %[[R5]]] 52 %a = affine.apply affine_map<(d0)[s0] -> ((d0 floordiv 32) mod 4 + s0 + 42)>(%j)[%i] 53 54 // CHECK: "some_side_effecting_consumer"(%[[a]]) : (index) -> () 55 "some_side_effecting_consumer"(%a) : (index) -> () 56 } 57 } 58 return 59} 60 61// ----- 62 63// CHECK-DAG: #[[$div4:.*]] = affine_map<()[s0] -> (s0 floordiv 4)> 64// CHECK-DAG: #[[$times32:.*]] = affine_map<()[s0] -> (s0 * 32)> 65// CHECK-DAG: #[[$times16:.*]] = affine_map<()[s0] -> (s0 * 16)> 66// CHECK-DAG: #[[$add:.*]] = affine_map<()[s0, s1] -> (s0 + s1)> 67// CHECK-DAG: #[[$div4timesm32:.*]] = affine_map<()[s0] -> ((s0 floordiv 4) * -32)> 68// CHECK-DAG: #[[$times8:.*]] = affine_map<()[s0] -> (s0 * 8)> 69// CHECK-DAG: #[[$id:.*]] = affine_map<()[s0] -> (s0)> 70// CHECK-DAG: #[[$div32div4timesm4:.*]] = affine_map<()[s0] -> (((s0 floordiv 32) floordiv 4) * -4)> 71// CHECK-DAG: #[[$div32:.*]] = affine_map<()[s0] -> (s0 floordiv 32)> 72 73// CHECK-LABEL: func.func @larger_test 74// CHECK-SAME: %[[I0:[0-9a-zA-Z]+]]: index, 75// CHECK-SAME: %[[I1:[0-9a-zA-Z]+]]: index, 76// CHECK-SAME: %[[I2:[0-9a-zA-Z]+]]: index, 77// CHECK-SAME: %[[LB:[0-9a-zA-Z]+]]: index, 78// CHECK-SAME: %[[UB:[0-9a-zA-Z]+]]: index, 79// CHECK-SAME: %[[STEP:[0-9a-zA-Z]+]]: index 80func.func @larger_test(%0: index, %1: index, %2: index, %lb: index, %ub: index, %step: index) { 81 %c2 = arith.constant 2 : index 82 %c6 = arith.constant 6 : index 83 84 // CHECK: %[[R0:.*]] = affine.apply #[[$div4]]()[%[[I0]]] 85 // CHECK-NEXT: %[[R1:.*]] = affine.apply #[[$times16]]()[%[[I1]]] 86 // CHECK-NEXT: %[[R2:.*]] = affine.apply #[[$add]]()[%[[R0]], %[[R1]]] 87 // CHECK-NEXT: %[[R3:.*]] = affine.apply #[[$times32]]()[%[[I2]]] 88 89 // I1 * 16 + I2 * 32 + I0 floordiv 4 90 // CHECK-NEXT: %[[b:.*]] = affine.apply #[[$add]]()[%[[R2]], %[[R3]]] 91 92 // (I0 floordiv 4) * 32 93 // CHECK-NEXT: %[[R5:.*]] = affine.apply #[[$div4timesm32]]()[%[[I0]]] 94 // 8 * I0 95 // CHECK-NEXT: %[[R6:.*]] = affine.apply #[[$times8]]()[%[[I0]]] 96 // 8 * I0 + (I0 floordiv 4) * 32 97 // CHECK-NEXT: %[[c:.*]] = affine.apply #[[$add]]()[%[[R5]], %[[R6]]] 98 99 // CHECK-NEXT: scf.for %[[i:.*]] = 100 scf.for %i = %lb to %ub step %step { 101 // remainder from %a not hoisted above %i. 102 // CHECK-NEXT: %[[R8:.*]] = affine.apply #[[$times32]]()[%[[i]]] 103 // CHECK-NEXT: %[[a:.*]] = affine.apply #[[$add]]()[%[[b]], %[[R8]]] 104 105 // CHECK-NEXT: scf.for %[[j:.*]] = 106 scf.for %j = %lb to %ub step %step { 107 // Gets hoisted partially to i and rest outermost. 108 // The hoisted part is %b. 109 %a = affine.apply affine_map<()[s0, s1, s2, s3] -> (s1 * 16 + s2 * 32 + s3 * 32 + s0 floordiv 4)>()[%0, %1, %2, %i] 110 111 // Gets completely hoisted 112 %b = affine.apply affine_map<()[s0, s1, s2] -> (s1 * 16 + s2 * 32 + s0 floordiv 4)>()[%0, %1, %2] 113 114 // Gets completely hoisted 115 %c = affine.apply affine_map<()[s0] -> (s0 * 8 - (s0 floordiv 4) * 32)>()[%0] 116 117 // 32 * %j + %c remains here, the rest is hoisted. 118 // CHECK-DAG: %[[R10:.*]] = affine.apply #[[$times32]]()[%[[j]]] 119 // CHECK-DAG: %[[d:.*]] = affine.apply #[[$add]]()[%[[c]], %[[R10]]] 120 %d = affine.apply affine_map<()[s0, s1] -> (s0 * 8 + s1 * 32 - (s0 floordiv 4) * 32)>()[%0, %j] 121 122 // CHECK-DAG: %[[idj:.*]] = affine.apply #[[$id]]()[%[[j]]] 123 // CHECK-NEXT: scf.for %[[k:.*]] = 124 scf.for %k = %lb to %ub step %step { 125 // CHECK-NEXT: %[[idk:.*]] = affine.apply #[[$id]]()[%[[k]]] 126 // CHECK-NEXT: %[[e:.*]] = affine.apply #[[$add]]()[%[[c]], %[[idk]]] 127 %e = affine.apply affine_map<()[s0, s1] -> (s0 + s1 * 8 - (s1 floordiv 4) * 32)>()[%k, %0] 128 129 // CHECK-NEXT: %[[R15:.*]] = affine.apply #[[$div32div4timesm4]]()[%[[k]]] 130 // CHECK-NEXT: %[[R16:.*]] = affine.apply #[[$add]]()[%[[idj]], %[[R15]]] 131 // CHECK-NEXT: %[[R17:.*]] = affine.apply #[[$div32]]()[%[[k]]] 132 // CHECK-NEXT: %[[f:.*]] = affine.apply #[[$add]]()[%[[R16]], %[[R17]]] 133 %f = affine.apply affine_map<(d0)[s0] -> ((d0 floordiv 32) mod 4 + s0)>(%k)[%j] 134 135 // CHECK-NEXT: %[[g:.*]] = affine.apply #[[$add]]()[%[[b]], %[[idk]]] 136 %g = affine.apply affine_map<()[s0, s1, s2, s3] -> (s0 + s2 * 16 + s3 * 32 + s1 floordiv 4)>()[%k, %0, %1, %2] 137 138 // CHECK-NEXT: "some_side_effecting_consumer"(%[[a]]) : (index) -> () 139 "some_side_effecting_consumer"(%a) : (index) -> () 140 // CHECK-NEXT: "some_side_effecting_consumer"(%[[b]]) : (index) -> () 141 "some_side_effecting_consumer"(%b) : (index) -> () 142 // CHECK-NEXT: "some_side_effecting_consumer"(%[[c]]) : (index) -> () 143 "some_side_effecting_consumer"(%c) : (index) -> () 144 // CHECK-NEXT: "some_side_effecting_consumer"(%[[d]]) : (index) -> () 145 "some_side_effecting_consumer"(%d) : (index) -> () 146 // CHECK-NEXT: "some_side_effecting_consumer"(%[[e]]) : (index) -> () 147 "some_side_effecting_consumer"(%e) : (index) -> () 148 // CHECK-NEXT: "some_side_effecting_consumer"(%[[f]]) : (index) -> () 149 "some_side_effecting_consumer"(%f) : (index) -> () 150 // CHECK-NEXT: "some_side_effecting_consumer"(%[[g]]) : (index) -> () 151 "some_side_effecting_consumer"(%g) : (index) -> () 152 } 153 } 154 } 155 return 156} 157