[Mlir-commits] [mlir] [mlir][x86] Fix VNNI operand load offsets (PR #230205)

[Adam Siemieniuk]

================
@@ -98,6 +98,30 @@ static bool isReadSrcMemref(Value operand) {
   return srcBuff && isa<MemRefType>(srcBuff.getType());
 }
 
+// Check if the memref read by a vector.contract operand has the VNNI factor as
+// its static innermost dim, so that the collapsed tile rows are contiguous.
+static bool hasVnniInnerDim(Value operand, int64_t vnni) {
+  Value srcBuff;
+  llvm::TypeSwitch<Operation *>(operand.getDefiningOp())
+      .Case<TransferReadOp, LoadOp>(
+          [&](auto readOp) { srcBuff = readOp.getOperand(0); });
+
+  if (!srcBuff)
+    return false;
+  auto srcType = dyn_cast<MemRefType>(srcBuff.getType());
+  return srcType && srcType.getRank() > 0 && srcType.getShape().back() == vnni;
+}
+
+// Replaces the indices of the two innermost dims of a VNNI operand by the
+// index of the dim they collapse into.
+static void collapseVnniIndices(OpBuilder &rewriter, Location loc,
+                                SmallVectorImpl<Value> &indices, int64_t vnni) {
+  Value inner = indices.pop_back_val();
+  Value cVnni = arith::ConstantIndexOp::create(rewriter, loc, vnni);
+  Value scaled = arith::MulIOp::create(rewriter, loc, indices.back(), cVnni);
+  indices.back() = arith::AddIOp::create(rewriter, loc, scaled, inner);
----------------
adam-smnk wrote:

Nope, it must be zero.
I emit the add anyway for completness as that's what I'd expect in collapse computation. It gets folded later.

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