[libcxx-commits] [libcxx] [mlir] Add ContractionOpInterface utility functions for vector matrix multiplication (PR #68945)

[Jerry Wu via libcxx-commits]

================
@@ -240,4 +240,134 @@ TEST(isRowMajorBatchMatmul, FirstInputSwapped) {
   EXPECT_THAT(maps, Not(Truly(isRowMajorBatchMatmul)));
 }
 
+TEST(isVecmat, Simple) {
+  MLIRContext context;
+
+  AffineExpr k, n;
+  bindDims(&context, k, n);
+  auto mapA = AffineMapAttr::get(AffineMap::get(2, 0, {k}, &context));
+  auto mapB = AffineMapAttr::get(AffineMap::get(2, 0, {k, n}, &context));
+  auto mapC = AffineMapAttr::get(AffineMap::get(2, 0, {n}, &context));
+  auto maps = ArrayAttr::get(&context, {mapA, mapB, mapC});
+
+  EXPECT_THAT(maps, Truly(isVecmat));
+}
+
+TEST(isVecmat, BindingSwapped) {
+  MLIRContext context;
+
+  AffineExpr k, n;
+  bindDims(&context, k, n); // bind in different order
+  auto mapA = AffineMapAttr::get(AffineMap::get(2, 0, {k}, &context));
+  auto mapB = AffineMapAttr::get(AffineMap::get(2, 0, {k, n}, &context));
+  auto mapC = AffineMapAttr::get(AffineMap::get(2, 0, {n}, &context));
+  auto maps = ArrayAttr::get(&context, {mapA, mapB, mapC});
+
+  EXPECT_THAT(maps, Truly(isVecmat));
+}
+
+TEST(isVecmat, WrongDimOrderMatrix) {
+  MLIRContext context;
+
+  AffineExpr k, n;
+  bindDims(&context, k, n);
+  auto mapA = AffineMapAttr::get(AffineMap::get(2, 0, {k}, &context));
+  auto mapB = AffineMapAttr::get(AffineMap::get(2, 0, {n, k}, &context));
+  auto mapC = AffineMapAttr::get(AffineMap::get(2, 0, {n}, &context));
+  auto maps = ArrayAttr::get(&context, {mapA, mapB, mapC});
+
+  EXPECT_THAT(maps, Not(Truly(isVecmat)));
+}
+
+TEST(isMatvec, Simple) {
+  MLIRContext context;
+
+  AffineExpr k, n;
+  bindDims(&context, k, n);
+  auto mapA = AffineMapAttr::get(AffineMap::get(2, 0, {n, k}, &context));
+  auto mapB = AffineMapAttr::get(AffineMap::get(2, 0, {k}, &context));
+  auto mapC = AffineMapAttr::get(AffineMap::get(2, 0, {n}, &context));
+  auto maps = ArrayAttr::get(&context, {mapA, mapB, mapC});
+
+  EXPECT_THAT(maps, Truly(isMatvec));
+}
+
+TEST(isMatvec, BindingSwapped) {
+  MLIRContext context;
+
+  AffineExpr k, n;
+  bindDims(&context, k, n); // bind in different order
+  auto mapA = AffineMapAttr::get(AffineMap::get(2, 0, {n, k}, &context));
+  auto mapB = AffineMapAttr::get(AffineMap::get(2, 0, {k}, &context));
+  auto mapC = AffineMapAttr::get(AffineMap::get(2, 0, {n}, &context));
+  auto maps = ArrayAttr::get(&context, {mapA, mapB, mapC});
+
+  EXPECT_THAT(maps, Truly(isMatvec));
+}
+
+TEST(isMatvec, WrongDimOrderMatrix) {
+  MLIRContext context;
+
+  AffineExpr k, n;
+  bindDims(&context, k, n);
+  auto mapA = AffineMapAttr::get(AffineMap::get(2, 0, {k, n}, &context));
+  auto mapB = AffineMapAttr::get(AffineMap::get(2, 0, {k}, &context));
+  auto mapC = AffineMapAttr::get(AffineMap::get(2, 0, {n}, &context));
+  auto maps = ArrayAttr::get(&context, {mapA, mapB, mapC});
+
+  EXPECT_THAT(maps, Not(Truly(isMatvec)));
+}
+
+TEST(isBatchMatvec, Simple) {
+  MLIRContext context;
+
+  AffineExpr batch, k, n;
+  bindDims(&context, batch, k, n);
+  auto mapA = AffineMapAttr::get(AffineMap::get(3, 0, {batch, n, k}, &context));
+  auto mapB = AffineMapAttr::get(AffineMap::get(3, 0, {batch, k}, &context));
+  auto mapC = AffineMapAttr::get(AffineMap::get(3, 0, {batch, n}, &context));
+  auto maps = ArrayAttr::get(&context, {mapA, mapB, mapC});
+
+  EXPECT_THAT(maps, Truly(isBatchMatvec));
+}
+
+TEST(isBatchMatvec, BindingSwapped) {
+  MLIRContext context;
+
+  AffineExpr batch, k, n;
+  bindDims(&context, batch, k, n); // bind in different order
----------------
pzread wrote:

Same here

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