[mlir] [MLIR][Linalg] Introduce broadcast/transpose semantic to 'linalg.batc… (PR #122275)

Andrzej Warzyński llvmlistbot at llvm.org
Mon Jan 20 12:09:29 PST 2025


================
@@ -3450,6 +3467,95 @@ static LogicalResult verifyExtendedMatmulSemantic(MatmulOp matmulOp,
   return success();
 }
 
+/// Checks if the given AffineMap represents a valid batch dimension.
+/// It checks if the first result dimension is a function of the first
+/// dimension.
+static bool isValidBatchDim(AffineMap bcastMap) {
+  AffineExpr exp = bcastMap.getResult(0);
+  return exp.isFunctionOfDim(0);
+}
+
+/// Checks if the given AffineMap's result dimensions are valid output result
+/// dimensions.
+static bool isValidOutputResultDim(AffineMap outputMap) {
+  enum Indices { batchPos, mPos, nPos };
+  AffineExpr exp0 = outputMap.getResult(batchPos);
+  AffineExpr exp1 = outputMap.getResult(mPos);
+  AffineExpr exp2 = outputMap.getResult(nPos);
+  return exp0.isFunctionOfDim(batchPos) && exp1.isFunctionOfDim(mPos) &&
+         exp2.isFunctionOfDim(nPos);
+}
+
+// Check general validity of input indexing map.
+static LogicalResult verifyInputMaps(BatchMatmulOp batchMatmulOp,
+                                     AffineMap opIndexingMap,
+                                     AffineMap defaultIndexingMap, bool isLHS) {
+  // Check the result dims are valid.
+  if (!isValidResultDimExprs(opIndexingMap, defaultIndexingMap))
----------------
banach-space wrote:

Please rename `isValidResultDimExprs` as `areResultExprsASubsetOf`. And then, in https://github.com/llvm/llvm-project/blob/3b001db4f9668cfa29572e5f1911ec7cef8b0ac2/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp#L3409-L3418

```cpp
 // Returns true if the result expression of `subMap` are a subset of `fullMap`.
static bool areResultExprsASubsetOf(AffineMap subMap, AffineMap fullMap):
```

Hope I am reading this correctly. Other suggestions are welcome, but `isValidResultDimExprs` is too generic (while the condition being checked is very specific).

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


More information about the Mlir-commits mailing list