[Mlir-commits] [mlir] [mlir][x86vector] Lower BF16 vector.contract to FMA using AVX2 BF16 packed ops. (PR #170267)

[Adam Siemieniuk]

================
@@ -24,6 +26,63 @@ using namespace mlir;
 using namespace mlir::vector;
 using namespace mlir::x86vector;
 
+static bool validateVectorProdOp(Value prodOp) {
+  Operation *defOp = prodOp.getDefiningOp();
+  if (!defOp)
+    return false;
+
+  // If the LHS/RHS op is transfer_read return false if:
+  // (1) - It has false in-bounds
+  // (2) - The permutation map is not identical
+  if (auto readOp = prodOp.getDefiningOp<mlir::vector::TransferReadOp>()) {
+    ArrayAttr inBoundsAttr = readOp.getInBoundsAttr();
+    if (inBoundsAttr) {
+
+      for (Attribute attr : inBoundsAttr) {
+        auto boolAttr = llvm::dyn_cast<BoolAttr>(attr);
+        if (!boolAttr || !boolAttr.getValue()) {
+          return false;
+        }
+      }
+    }
+
+    if (!readOp.getPermutationMap().isIdentity())
+      return false;
+  }
+
+  Value srcBuff;
+  SmallVector<OpFoldResult> indexVals;
+  llvm::TypeSwitch<Operation *>(defOp).Case<TransferReadOp, LoadOp>(
+      [&](auto readOp) {
+        srcBuff = readOp.getOperand(0);
+        indexVals = SmallVector<OpFoldResult>(readOp.getIndices().begin(),
+                                              readOp.getIndices().end());
+      });
+
+  if (!srcBuff)
+    return false;
+
+  // Return false, if the source is not a memref type
+  Type srcType = srcBuff.getType();
+  if (!llvm::isa<MemRefType>(srcType))
+    return false;
+
+  // Return false, if the innermost stride of the memref is not 1.
----------------
adam-smnk wrote:

Btw, I wonder if in case of packed memref source, shouldn't both VNNI and one dim be both contiguous?

As in, for a packed source type `memref<1x1x32x2xbf16>`, we probably need to ensure that not only these `2` packed elements are contiguous but also there are enough contiguous packed tuples for `x86vector.avx.cvt.packed.even.indexed_to_f32` to readout.

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