Require expected rank in isInVnniLayout

[chelini]

Until we have a better way to express the VNNI layout (see: #563), it is up to the caller to specify the expected rank in the VNNI layout as the rank depends on the operations we are dealing with. The current check for rank < 3 is not correct; this is still ugly, but at least more explicit, as there are not magic assumptions on rank.

[chelini]

I don't have good solution at the moment for VNNI. Current checks are not correct. The VNNI ranks depends on the operation we are dealing with (i.e., GEMM may have VNNI layout when memref has rank 3, while BRGEMM 4 for the inputs while 3 for the output).

[rengolin]

I don't have good solution at the moment for VNNI. Current checks are not correct. The VNNI ranks depends on the operation we are dealing with (i.e., GEMM may have VNNI layout when memref has rank 3, while BRGEMM 4 for the inputs while 3 for the output).

The problem with this is that we're passing down to the caller the responsibility to "expect" things. I'd use some enum to hide that like:

enum class VNNISource {
  PACKED = 3,
  GEMM = 3,
  BRGEMM = 4,
  FUSED_BRGEMM = ?
};
bool isInVnniLayout(VNNISource source, MemRefType memref);

Then use isVNNILayout(VNNISource::BRGEMM, memref).