[Mlir-commits] [mlir] [mlir][tosa] Support variadic matmul sharding (PR #226180)

[Frank Schlimbach]

================
@@ -73,6 +73,9 @@ checkOperandAffineExprRecursively(AffineExpr expr,
     seenIds[position] = true;
     return success();
   }
+  case AffineExprKind::Constant:
+    // A constant indexing expression does not reference a loop dimension.
+    return success();
----------------
fschlimb wrote:

This simple `case` accepts additional cases that are not tested or supported anywhere. For example, expressions like "d0 + 3" will be accepted and incorrectly propagated as "d0".

To correctly accept pure constants only (which is what this PR seems to rely on) you could just add an early-exit-with-success-check in checkOperandAffineExpr below.

```c++
static FailureOr<llvm::SmallSet<unsigned, 2>>
checkOperandAffineExpr(AffineExpr expr, unsigned numDims) {
  // A constant result (e.g. a broadcast dimension) references no loop.
  if (isa<AffineConstantExpr>(expr))
    return llvm::SmallSet<unsigned, 2>();
  SmallVector<bool> seenIds(numDims, false);
  ...
```

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