[Mlir-commits] [mlir] [mlir][vector] Add a check to ensure input vector rank equals target shape rank (PR #127706)

[Andrzej Warzyński]

================
@@ -188,6 +188,20 @@ func.func @vector_fma(%a: vector<4x4xf32>, %b: vector<4x4xf32>, %c: vector<4x4xf
 //   CHECK-LABEL: func @vector_fma
 // CHECK-COUNT-4: vector.fma %{{.+}}, %{{.+}}, %{{.+}} : vector<2x2xf32>
 
+func.func @higher_rank_unroll() {
+  %cst_25 = arith.constant dense<3.718400e+04> : vector<4x2x2xf16>
+  %cst_26 = arith.constant dense<1.000000e+00> : vector<24x2x2xf32>
+  %47 = vector.fma %cst_26, %cst_26, %cst_26 : vector<24x2x2xf32>
+  %818 = scf.execute_region -> vector<24x2x2xf32> {
+      scf.yield %47 : vector<24x2x2xf32>
+    }
+  %823 = vector.extract_strided_slice %cst_25 {offsets = [2], sizes = [1], strides = [1]} : vector<4x2x2xf16> to vector<1x2x2xf16>
+  return
+}
+
+// CHECK-LABEL: func @higher_rank_unroll
+//       CHECK: return 
----------------
banach-space wrote:

Thanks! I have some suggestion regarding the formatting of this. I am actually in the process of extending the MLIR guidelines, but that's not been merged yet:
* https://github.com/llvm/mlir-www/pull/216

So:
* Use meaningful variable names (`%cst_25` is quite random).
* Avoid unnecessary Ops - are `arith.constant` + `scf.execute_region` + `vector.extract_strided_slice` required to trigger the issue that you are fixing? I am confident that `%c25` and `%c26` can be replaced with function arguments.
* Could this be written as a variation of `%vector_fma` above? For example, `%vector_fma_3d`?
* This is a negative test (i.e. the unrolling is meant to fail to apply). Let's make it clear: `@negative_vector_fma_3d`.
* You can also add `CHECK-NOT: vector.extract_strided_slice` and `CHECK: %47 = vector.fma %cst_26, %cst_26, %cst_26 : vector<24x2x2xf32>` (with some meaningful LIT variable names) to show that no unrolling took place.

Hope these make sense. If not, do ask and I will elaborate :) 

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