[Mlir-commits] [mlir] [mlir][SPIRV] Handle vector stores when unifying aliased resources (PR #232185)

[Anutosh Bhat]

================
@@ -550,15 +551,75 @@ struct ConvertStore : public ConvertAliasResource<spirv::StoreOp> {
         cast<spirv::PointerType>(storeOp.getPtr().getType()).getPointeeType();
     auto dstElemType =
         cast<spirv::PointerType>(adaptor.getPtr().getType()).getPointeeType();
-    if (!srcElemType.isIntOrFloat() || !dstElemType.isIntOrFloat())
-      return rewriter.notifyMatchFailure(storeOp, "not scalar type");
-    if (!areSameBitwidthScalarType(srcElemType, dstElemType))
-      return rewriter.notifyMatchFailure(storeOp, "different bitwidth");
-
     Location loc = storeOp.getLoc();
     Value value = adaptor.getValue();
-    if (srcElemType != dstElemType)
+    auto srcVector = dyn_cast<VectorType>(srcElemType);
+    auto dstVector = dyn_cast<VectorType>(dstElemType);
+    if (srcVector && dstVector && srcVector != dstVector) {
+      if (srcVector.getElementTypeBitWidth() !=
+              dstVector.getElementTypeBitWidth() ||
+          srcVector.getNumElements() % dstVector.getNumElements() != 0)
+        return rewriter.notifyMatchFailure(storeOp,
+                                           "incompatible vector types");
+      auto acOp = adaptor.getPtr().getDefiningOp<spirv::AccessChainOp>();
+      if (!acOp)
+        return rewriter.notifyMatchFailure(storeOp,
+                                           "ptr not spirv.AccessChain");
+      if (srcVector.getElementType() != dstVector.getElementType())
+        value = spirv::BitcastOp::create(
+            rewriter, loc,
+            VectorType::get(srcVector.getShape(), dstVector.getElementType()),
+            value);
+      unsigned count = dstVector.getNumElements();
+      unsigned ratio = srcVector.getNumElements() / count;
+      unsigned bytes = count * dstVector.getElementTypeBitWidth() / 8;
+
+      // Split the value into vectors matching the canonical resource and store
+      // them at consecutive vector indices. Adjust each store's alignment for
+      // its byte offset from the original pointer.
+      auto getPart = [&](unsigned i) -> Value {
+        if (ratio == 1)
+          return value;
+        SmallVector<int32_t> components;
+        for (unsigned j = 0; j < count; ++j)
+          components.push_back(i * count + j);
+        return spirv::VectorShuffleOp::create(
+            rewriter, loc, dstVector, value, value,
+            rewriter.getI32ArrayAttr(components));
+      };
+      auto firstStore = spirv::StoreOp::create(
+          rewriter, loc, adaptor.getPtr(), getPart(0),
+          storeOp.getMemoryAccessAttr(), storeOp.getAlignmentAttr());
+      firstStore->setDiscardableAttrs(storeOp->getDiscardableAttrDictionary());
+      for (unsigned i = 1; i < ratio; ++i) {
+        Value part = getPart(i);
+        auto indices = llvm::to_vector(acOp.getIndices());
----------------
anutosh491 wrote:

Done.

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