[Mlir-commits] [mlir] [mlir][Vector] Reject scalable outer dim in multi_reduction unrolling (PR #224174)

[Cheng-Chieh Liu]

================
@@ -64,6 +64,50 @@ func.func @vector_multi_reduction_masked(%arg0: vector<2x4xf32>, %acc: vector<2x
 
 // -----
 
+//       ALL-LABEL: @negative_multi_reduction_scalable_outer
+//        ALL-SAME: %[[INPUT:.+]]: vector<[2]x4xf32>, %[[ACC:.+]]: vector<[2]xf32>
+// INNER-REDUCTION: vector.multi_reduction <mul>, %[[INPUT]], %[[ACC]] [1] : vector<[2]x4xf32> to vector<[2]xf32>
+//  INNER-PARALLEL: vector.transpose %[[INPUT]], [1, 0] : vector<[2]x4xf32> to vector<4x[2]xf32>
+func.func @negative_multi_reduction_scalable_outer(%arg0: vector<[2]x4xf32>, %acc: vector<[2]xf32>) -> vector<[2]xf32> {
+    %0 = vector.multi_reduction <mul>, %arg0, %acc [1] : vector<[2]x4xf32> to vector<[2]xf32>
+    return %0 : vector<[2]xf32>
----------------
jay87406 wrote:

Agreed, the checks stopped at the transpose. The lowering strategy now checks the unrolled sequence through to the `return` (all four iterations for the unmasked tests; first iteration plus the final result for the masked ones, like `@vector_multi_reduction_masked`). Same for `INNER-REDUCTION` in the reduce variants. The bailing strategy now also checks that the original op is what gets returned.

https://github.com/llvm/llvm-project/pull/224174