297 lines · plain
1; RUN: llc -mtriple=aarch64-linux-gnu -mattr=+sme2 -aarch64-new-sme-abi < %s | FileCheck %s2 3; This test case was generated by lowering mlir/test/Integration/Dialect/Linalg/CPU/ArmSME/matmul.mlir to LLVM IR.4; The actual contents of the function are not that important. The main interesting quality here is that many blocks5; don't directly use ZA. The only blocks that require ZA are the MOPA (and load/stores) in the inner loop, and the6;`printMemrefF32()` call in the exit block.7;8; If ZA states are not propagated in the MachineSMEABIPass block %48 (which is within the outer loop), will9; have an edge to block %226 (the exit block), which requires ZA in the "saved" state, and an edge to block %5110; (which has no preference on ZA state). This means block %48 will also end up in the locally saved state.11; This is not really what we want, as it means we will save/restore ZA in the outer loop. We can fix this by12; propagating the "active" state from the inner loop through basic blocks with no preference, to ensure the outer13; loop is in the "active" state too.14;15; If done correctly, the only ZA save/restore should be in the exit block (with all other blocks in the active state).16 17define void @matmul(ptr %0, ptr %1, i64 %2, i64 %3, i64 %4, i64 %5, i64 %6, ptr %7, ptr %8, i64 %9, i64 %10, i64 %11, i64 %12, i64 %13, ptr %14, ptr %15, i64 %16, i64 %17, i64 %18, i64 %19, i64 %20) #0 {18; Check for a ZA zero in the entry block, then no uses of TPIDR2_EL0 (for ZA saves/restore)19; until the exit block (which contains the call to printMemrefF32).20;21; CHECK-LABEL: matmul:22; CHECK: zero {za}23; CHECK-NOT: TPIDR2_EL024; CHECK: msr TPIDR2_EL0, x{{.*}}25; CHECK-NOT: .LBB{{.*}}26; CHECK: bl printMemrefF3227 %22 = insertvalue { ptr, ptr, i64, [2 x i64], [2 x i64] } poison, ptr %14, 028 %23 = insertvalue { ptr, ptr, i64, [2 x i64], [2 x i64] } %22, ptr %15, 129 %24 = insertvalue { ptr, ptr, i64, [2 x i64], [2 x i64] } %23, i64 %16, 230 %25 = insertvalue { ptr, ptr, i64, [2 x i64], [2 x i64] } %24, i64 %17, 3, 031 %26 = insertvalue { ptr, ptr, i64, [2 x i64], [2 x i64] } %25, i64 %19, 4, 032 %27 = insertvalue { ptr, ptr, i64, [2 x i64], [2 x i64] } %26, i64 %18, 3, 133 %28 = insertvalue { ptr, ptr, i64, [2 x i64], [2 x i64] } %27, i64 %20, 4, 134 %29 = insertvalue { ptr, ptr, i64, [2 x i64], [2 x i64] } poison, ptr %7, 035 %30 = insertvalue { ptr, ptr, i64, [2 x i64], [2 x i64] } %29, ptr %8, 136 %31 = insertvalue { ptr, ptr, i64, [2 x i64], [2 x i64] } %30, i64 %9, 237 %32 = insertvalue { ptr, ptr, i64, [2 x i64], [2 x i64] } %31, i64 %10, 3, 038 %33 = insertvalue { ptr, ptr, i64, [2 x i64], [2 x i64] } %32, i64 %12, 4, 039 %34 = insertvalue { ptr, ptr, i64, [2 x i64], [2 x i64] } %33, i64 %11, 3, 140 %35 = insertvalue { ptr, ptr, i64, [2 x i64], [2 x i64] } %34, i64 %13, 4, 141 %36 = insertvalue { ptr, ptr, i64, [2 x i64], [2 x i64] } poison, ptr %0, 042 %37 = insertvalue { ptr, ptr, i64, [2 x i64], [2 x i64] } %36, ptr %1, 143 %38 = insertvalue { ptr, ptr, i64, [2 x i64], [2 x i64] } %37, i64 %2, 244 %39 = insertvalue { ptr, ptr, i64, [2 x i64], [2 x i64] } %38, i64 %3, 3, 045 %40 = insertvalue { ptr, ptr, i64, [2 x i64], [2 x i64] } %39, i64 %5, 4, 046 %41 = insertvalue { ptr, ptr, i64, [2 x i64], [2 x i64] } %40, i64 %4, 3, 147 %42 = insertvalue { ptr, ptr, i64, [2 x i64], [2 x i64] } %41, i64 %6, 4, 148 %43 = extractvalue { ptr, ptr, i64, [2 x i64], [2 x i64] } %42, 3, 049 %44 = extractvalue { ptr, ptr, i64, [2 x i64], [2 x i64] } %42, 3, 150 %45 = extractvalue { ptr, ptr, i64, [2 x i64], [2 x i64] } %35, 3, 151 %46 = call i64 @llvm.vscale.i64()52 %47 = mul i64 %46, 453 br label %4854 5548: ; preds = %224, %2156 %49 = phi i64 [ %225, %224 ], [ 0, %21 ]57 %50 = icmp slt i64 %49, %4358 br i1 %50, label %51, label %22659 6051: ; preds = %4861 %52 = sub i64 %43, %4962 %53 = call i64 @llvm.smin.i64(i64 %47, i64 %52)63 %54 = call <vscale x 4 x i32> @llvm.stepvector.nxv4i32()64 %55 = trunc i64 %53 to i3265 %56 = insertelement <vscale x 4 x i32> poison, i32 %55, i32 066 %57 = shufflevector <vscale x 4 x i32> %56, <vscale x 4 x i32> poison, <vscale x 4 x i32> zeroinitializer67 %58 = icmp slt <vscale x 4 x i32> %54, %5768 br label %5969 7059: ; preds = %222, %5171 %60 = phi i64 [ %223, %222 ], [ 0, %51 ]72 %61 = icmp slt i64 %60, %4573 br i1 %61, label %62, label %22474 7562: ; preds = %5976 %63 = sub i64 %45, %6077 %64 = call i64 @llvm.smin.i64(i64 %47, i64 %63)78 %65 = extractvalue { ptr, ptr, i64, [2 x i64], [2 x i64] } %28, 079 %66 = extractvalue { ptr, ptr, i64, [2 x i64], [2 x i64] } %28, 180 %67 = insertvalue { ptr, ptr, i64 } poison, ptr %65, 081 %68 = insertvalue { ptr, ptr, i64 } %67, ptr %66, 182 %69 = insertvalue { ptr, ptr, i64 } %68, i64 0, 283 %70 = extractvalue { ptr, ptr, i64, [2 x i64], [2 x i64] } %28, 284 %71 = extractvalue { ptr, ptr, i64, [2 x i64], [2 x i64] } %28, 3, 085 %72 = extractvalue { ptr, ptr, i64, [2 x i64], [2 x i64] } %28, 3, 186 %73 = extractvalue { ptr, ptr, i64, [2 x i64], [2 x i64] } %28, 4, 087 %74 = extractvalue { ptr, ptr, i64, [2 x i64], [2 x i64] } %28, 4, 188 %75 = mul nsw i64 %49, %7389 %76 = add i64 %70, %7590 %77 = mul nsw i64 %60, %7491 %78 = add i64 %76, %7792 %79 = extractvalue { ptr, ptr, i64 } %69, 093 %80 = extractvalue { ptr, ptr, i64 } %69, 194 %81 = insertvalue { ptr, ptr, i64, [2 x i64], [2 x i64] } poison, ptr %79, 095 %82 = insertvalue { ptr, ptr, i64, [2 x i64], [2 x i64] } %81, ptr %80, 196 %83 = insertvalue { ptr, ptr, i64, [2 x i64], [2 x i64] } %82, i64 %78, 297 %84 = insertvalue { ptr, ptr, i64, [2 x i64], [2 x i64] } %83, i64 %53, 3, 098 %85 = insertvalue { ptr, ptr, i64, [2 x i64], [2 x i64] } %84, i64 %73, 4, 099 %86 = insertvalue { ptr, ptr, i64, [2 x i64], [2 x i64] } %85, i64 %64, 3, 1100 %87 = insertvalue { ptr, ptr, i64, [2 x i64], [2 x i64] } %86, i64 %74, 4, 1101 %88 = call <vscale x 4 x i32> @llvm.stepvector.nxv4i32()102 %89 = trunc i64 %64 to i32103 %90 = insertelement <vscale x 4 x i32> poison, i32 %89, i32 0104 %91 = shufflevector <vscale x 4 x i32> %90, <vscale x 4 x i32> poison, <vscale x 4 x i32> zeroinitializer105 %92 = icmp slt <vscale x 4 x i32> %88, %91106 br label %93107 10893: ; preds = %220, %62109 %94 = phi i64 [ %221, %220 ], [ 0, %62 ]110 %95 = icmp slt i64 %94, %44111 br i1 %95, label %96, label %222112 11396: ; preds = %93114 %97 = extractvalue { ptr, ptr, i64, [2 x i64], [2 x i64] } %42, 0115 %98 = extractvalue { ptr, ptr, i64, [2 x i64], [2 x i64] } %42, 1116 %99 = insertvalue { ptr, ptr, i64 } poison, ptr %97, 0117 %100 = insertvalue { ptr, ptr, i64 } %99, ptr %98, 1118 %101 = insertvalue { ptr, ptr, i64 } %100, i64 0, 2119 %102 = extractvalue { ptr, ptr, i64, [2 x i64], [2 x i64] } %42, 2120 %103 = extractvalue { ptr, ptr, i64, [2 x i64], [2 x i64] } %42, 3, 0121 %104 = extractvalue { ptr, ptr, i64, [2 x i64], [2 x i64] } %42, 3, 1122 %105 = extractvalue { ptr, ptr, i64, [2 x i64], [2 x i64] } %42, 4, 0123 %106 = extractvalue { ptr, ptr, i64, [2 x i64], [2 x i64] } %42, 4, 1124 %107 = mul nsw i64 %49, %105125 %108 = add i64 %102, %107126 %109 = mul nsw i64 %94, %106127 %110 = add i64 %108, %109128 %111 = extractvalue { ptr, ptr, i64 } %101, 0129 %112 = extractvalue { ptr, ptr, i64 } %101, 1130 %113 = insertvalue { ptr, ptr, i64, [1 x i64], [1 x i64] } poison, ptr %111, 0131 %114 = insertvalue { ptr, ptr, i64, [1 x i64], [1 x i64] } %113, ptr %112, 1132 %115 = insertvalue { ptr, ptr, i64, [1 x i64], [1 x i64] } %114, i64 %110, 2133 %116 = insertvalue { ptr, ptr, i64, [1 x i64], [1 x i64] } %115, i64 %53, 3, 0134 %117 = insertvalue { ptr, ptr, i64, [1 x i64], [1 x i64] } %116, i64 %105, 4, 0135 br label %118136 137118: ; preds = %133, %96138 %119 = phi i64 [ %135, %133 ], [ 0, %96 ]139 %120 = phi <vscale x 4 x float> [ %134, %133 ], [ poison, %96 ]140 %121 = icmp slt i64 %119, %47141 br i1 %121, label %122, label %136142 143122: ; preds = %118144 %123 = extractelement <vscale x 4 x i1> %58, i64 %119145 br i1 %123, label %124, label %133146 147124: ; preds = %122148 %125 = extractvalue { ptr, ptr, i64, [1 x i64], [1 x i64] } %117, 1149 %126 = extractvalue { ptr, ptr, i64, [1 x i64], [1 x i64] } %117, 2150 %127 = getelementptr float, ptr %125, i64 %126151 %128 = extractvalue { ptr, ptr, i64, [1 x i64], [1 x i64] } %117, 4, 0152 %129 = mul nuw nsw i64 %119, %128153 %130 = getelementptr inbounds nuw float, ptr %127, i64 %129154 %131 = load float, ptr %130, align 4155 %132 = insertelement <vscale x 4 x float> %120, float %131, i64 %119156 br label %133157 158133: ; preds = %124, %122159 %134 = phi <vscale x 4 x float> [ %132, %124 ], [ %120, %122 ]160 %135 = add i64 %119, 1161 br label %118162 163136: ; preds = %118164 %137 = extractvalue { ptr, ptr, i64, [2 x i64], [2 x i64] } %35, 0165 %138 = extractvalue { ptr, ptr, i64, [2 x i64], [2 x i64] } %35, 1166 %139 = insertvalue { ptr, ptr, i64 } poison, ptr %137, 0167 %140 = insertvalue { ptr, ptr, i64 } %139, ptr %138, 1168 %141 = insertvalue { ptr, ptr, i64 } %140, i64 0, 2169 %142 = extractvalue { ptr, ptr, i64, [2 x i64], [2 x i64] } %35, 2170 %143 = extractvalue { ptr, ptr, i64, [2 x i64], [2 x i64] } %35, 3, 0171 %144 = extractvalue { ptr, ptr, i64, [2 x i64], [2 x i64] } %35, 3, 1172 %145 = extractvalue { ptr, ptr, i64, [2 x i64], [2 x i64] } %35, 4, 0173 %146 = extractvalue { ptr, ptr, i64, [2 x i64], [2 x i64] } %35, 4, 1174 %147 = mul nsw i64 %94, %145175 %148 = add i64 %142, %147176 %149 = mul nsw i64 %60, %146177 %150 = add i64 %148, %149178 %151 = extractvalue { ptr, ptr, i64 } %141, 0179 %152 = extractvalue { ptr, ptr, i64 } %141, 1180 %153 = insertvalue { ptr, ptr, i64, [1 x i64], [1 x i64] } poison, ptr %151, 0181 %154 = insertvalue { ptr, ptr, i64, [1 x i64], [1 x i64] } %153, ptr %152, 1182 %155 = insertvalue { ptr, ptr, i64, [1 x i64], [1 x i64] } %154, i64 %150, 2183 %156 = insertvalue { ptr, ptr, i64, [1 x i64], [1 x i64] } %155, i64 %64, 3, 0184 %157 = insertvalue { ptr, ptr, i64, [1 x i64], [1 x i64] } %156, i64 %146, 4, 0185 br label %158186 187158: ; preds = %173, %136188 %159 = phi i64 [ %175, %173 ], [ 0, %136 ]189 %160 = phi <vscale x 4 x float> [ %174, %173 ], [ poison, %136 ]190 %161 = icmp slt i64 %159, %47191 br i1 %161, label %162, label %176192 193162: ; preds = %158194 %163 = extractelement <vscale x 4 x i1> %92, i64 %159195 br i1 %163, label %164, label %173196 197164: ; preds = %162198 %165 = extractvalue { ptr, ptr, i64, [1 x i64], [1 x i64] } %157, 1199 %166 = extractvalue { ptr, ptr, i64, [1 x i64], [1 x i64] } %157, 2200 %167 = getelementptr float, ptr %165, i64 %166201 %168 = extractvalue { ptr, ptr, i64, [1 x i64], [1 x i64] } %157, 4, 0202 %169 = mul nuw nsw i64 %159, %168203 %170 = getelementptr inbounds nuw float, ptr %167, i64 %169204 %171 = load float, ptr %170, align 4205 %172 = insertelement <vscale x 4 x float> %160, float %171, i64 %159206 br label %173207 208173: ; preds = %164, %162209 %174 = phi <vscale x 4 x float> [ %172, %164 ], [ %160, %162 ]210 %175 = add i64 %159, 1211 br label %158212 213176: ; preds = %158214 %177 = trunc i64 %64 to i32215 br label %178216 217178: ; preds = %181, %176218 %179 = phi i64 [ %202, %181 ], [ 0, %176 ]219 %180 = icmp slt i64 %179, %47220 br i1 %180, label %181, label %203221 222181: ; preds = %178223 %182 = icmp ult i64 %179, %53224 %183 = sext i1 %182 to i32225 %184 = and i32 %183, %177226 %185 = sext i32 %184 to i64227 %186 = call <vscale x 4 x i32> @llvm.stepvector.nxv4i32()228 %187 = trunc i64 %185 to i32229 %188 = insertelement <vscale x 4 x i32> poison, i32 %187, i32 0230 %189 = shufflevector <vscale x 4 x i32> %188, <vscale x 4 x i32> poison, <vscale x 4 x i32> zeroinitializer231 %190 = icmp slt <vscale x 4 x i32> %186, %189232 %191 = extractvalue { ptr, ptr, i64, [2 x i64], [2 x i64] } %87, 1233 %192 = extractvalue { ptr, ptr, i64, [2 x i64], [2 x i64] } %87, 2234 %193 = getelementptr float, ptr %191, i64 %192235 %194 = extractvalue { ptr, ptr, i64, [2 x i64], [2 x i64] } %87, 4, 0236 %195 = mul i64 %179, %194237 %196 = extractvalue { ptr, ptr, i64, [2 x i64], [2 x i64] } %87, 4, 1238 %197 = mul i64 0, %196239 %198 = add i64 %195, %197240 %199 = getelementptr float, ptr %193, i64 %198241 %200 = call <vscale x 4 x float> @llvm.masked.load.nxv4f32.p0(ptr %199, i32 4, <vscale x 4 x i1> %190, <vscale x 4 x float> poison)242 %201 = trunc i64 %179 to i32243 call void @llvm.aarch64.sme.write.horiz.nxv4f32(i32 0, i32 %201, <vscale x 4 x i1> splat (i1 true), <vscale x 4 x float> %200)244 %202 = add i64 %179, 1245 br label %178246 247203: ; preds = %178248 call void @llvm.aarch64.sme.mopa.nxv4f32(i32 0, <vscale x 4 x i1> %58, <vscale x 4 x i1> %92, <vscale x 4 x float> %120, <vscale x 4 x float> %160)249 %204 = call i64 @llvm.smin.i64(i64 %53, i64 %47)250 br label %205251 252205: ; preds = %208, %203253 %206 = phi i64 [ %219, %208 ], [ 0, %203 ]254 %207 = icmp slt i64 %206, %204255 br i1 %207, label %208, label %220256 257208: ; preds = %205258 %209 = extractvalue { ptr, ptr, i64, [2 x i64], [2 x i64] } %87, 1259 %210 = extractvalue { ptr, ptr, i64, [2 x i64], [2 x i64] } %87, 2260 %211 = getelementptr float, ptr %209, i64 %210261 %212 = extractvalue { ptr, ptr, i64, [2 x i64], [2 x i64] } %87, 4, 0262 %213 = mul i64 %206, %212263 %214 = extractvalue { ptr, ptr, i64, [2 x i64], [2 x i64] } %87, 4, 1264 %215 = mul i64 0, %214265 %216 = add i64 %213, %215266 %217 = getelementptr float, ptr %211, i64 %216267 %218 = trunc i64 %206 to i32268 call void @llvm.aarch64.sme.st1w.horiz(<vscale x 4 x i1> %92, ptr %217, i32 0, i32 %218)269 %219 = add i64 %206, 1270 br label %205271 272220: ; preds = %205273 %221 = add i64 %94, 1274 br label %93275 276222: ; preds = %93277 %223 = add i64 %60, %47278 br label %59279 280224: ; preds = %59281 %225 = add i64 %49, %47282 br label %48283 284226: ; preds = %48285 %227 = alloca { ptr, ptr, i64, [2 x i64], [2 x i64] }, i64 1, align 8286 store { ptr, ptr, i64, [2 x i64], [2 x i64] } %28, ptr %227, align 8287 %228 = insertvalue { i64, ptr } { i64 2, ptr poison }, ptr %227, 1288 %229 = extractvalue { i64, ptr } %228, 0289 %230 = extractvalue { i64, ptr } %228, 1290 call void @printMemrefF32(i64 %229, ptr %230)291 ret void292}293 294declare void @printMemrefF32(i64, ptr)295 296attributes #0 = { "aarch64_new_za" "aarch64_pstate_sm_body" }297