[Mlir-commits] [mlir] [mlir][SPIR-V][VectorToSPIRV] Add conversion patterns for vector.gather/scatter to SPIR-V (PR #193422)

[Jakub Kuderski]

================
@@ -853,6 +853,175 @@ struct VectorStoreOpConverter final
   }
 };
 
+struct VectorGatherOpConverter final
+    : public OpConversionPattern<vector::GatherOp> {
+  using Base::Base;
+
+  LogicalResult
+  matchAndRewrite(vector::GatherOp gatherOp, OpAdaptor adaptor,
+                  ConversionPatternRewriter &rewriter) const override {
+    // Only support 1-D result vectors.
+    auto vectorType = gatherOp.getVectorType();
+    if (vectorType.getRank() != 1)
+      return rewriter.notifyMatchFailure(gatherOp,
+                                         "only 1-D vectors supported");
+
+    // Only support memref base (not tensor).
+    auto memrefType = dyn_cast<MemRefType>(gatherOp.getBaseType());
+    if (!memrefType)
+      return rewriter.notifyMatchFailure(gatherOp,
+                                         "only memref base supported");
+
+    auto attr =
+        dyn_cast_or_null<spirv::StorageClassAttr>(memrefType.getMemorySpace());
+    if (!attr)
+      return rewriter.notifyMatchFailure(gatherOp,
+                                         "expected spirv.storage_class");
+
+    const auto &typeConverter = *getTypeConverter<SPIRVTypeConverter>();
+    if (!typeConverter.getTargetEnv().allows(
+            spirv::Extension::SPV_INTEL_masked_gather_scatter))
+      return rewriter.notifyMatchFailure(gatherOp,
+                                         "target environment does not enable "
+                                         "SPV_INTEL_masked_gather_scatter");
+    auto loc = gatherOp.getLoc();
+
+    // Compute base element pointer from memref + offsets.
+    Value basePtr =
+        spirv::getElementPtr(typeConverter, memrefType, adaptor.getBase(),
+                             adaptor.getOffsets(), loc, rewriter);
+    if (!basePtr)
+      return rewriter.notifyMatchFailure(gatherOp,
+                                         "failed to get element pointer");
+
+    // Convert element type and construct pointer vector type.
+    auto storageClass = attr.getValue();
+    Type elementType = typeConverter.convertType(memrefType.getElementType());
+    if (!elementType)
+      return rewriter.notifyMatchFailure(gatherOp, "unsupported element type");
+    auto ptrType = spirv::PointerType::get(elementType, storageClass);
+    int64_t numElements = vectorType.getDimSize(0);
+    auto ptrVectorType = VectorType::get({numElements}, ptrType);
+
+    // Build pointer vector: for each index, compute ptr via PtrAccessChain.
+    auto indexType = typeConverter.getIndexType();
+    SmallVector<Value> pointers;
+    for (int64_t i = 0; i < numElements; ++i) {
+      auto i32Type = rewriter.getI32Type();
+      Value idx = spirv::ConstantOp::create(rewriter, loc, i32Type,
+                                            rewriter.getI32IntegerAttr(i));
+      Value scalarIndex = spirv::VectorExtractDynamicOp::create(
+          rewriter, loc, adaptor.getIndices(), idx);
+      // Cast index to the SPIR-V index type if needed.
+      if (scalarIndex.getType() != indexType)
+        scalarIndex =
+            spirv::SConvertOp::create(rewriter, loc, indexType, scalarIndex);
+      Value ptr = spirv::PtrAccessChainOp::create(rewriter, loc, basePtr,
+                                                  scalarIndex, /*indices=*/{});
+      pointers.push_back(ptr);
----------------
kuhar wrote:

Have you confirmed this produces the same addresses at the llvm lowering? vector.gather is underspecified for non-1d and non-stride-1 memrefs/tensors, so maybe it would be better to restrict this to 1d types only for the time being... See https://discourse.llvm.org/t/rfc-vector-gather-vector-scatter-must-take-multiple-index-vectors/90728

cc: @krzysz00 

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