[Mlir-commits] [mlir] [rfc][mlir][gpu] Add operations to extract/insert/rotate within subgroup (PR #139048)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Thu May 8 06:19:45 PDT 2025
================
@@ -111,6 +111,68 @@ struct WmmaConstantOpToSPIRVLowering final
}
};
+/// Converts GPU MMA ExtractOp to CompositeExtract SPIR-V KHR/NV cooperative
+/// matrix ops.
+struct WmmaExtractOpToSPIRVLowering final
+ : OpConversionPattern<gpu::SubgroupMmaExtractOp> {
+ using OpConversionPattern::OpConversionPattern;
+
+ LogicalResult
+ matchAndRewrite(gpu::SubgroupMmaExtractOp op, OpAdaptor adaptor,
+ ConversionPatternRewriter &rewriter) const override {
+ Value matrix = adaptor.getMatrix();
+ auto coopType =
+ getTypeConverter()->convertType<spirv::CooperativeMatrixType>(
+ matrix.getType());
+ if (!coopType)
+ return rewriter.notifyMatchFailure(op, "type conversion failed");
+
+ SmallVector<int32_t> intValues;
+ for (Value val : op.getIndices()) {
+ if (auto constOp = val.getDefiningOp<arith::ConstantIndexOp>()) {
+ intValues.push_back(static_cast<int32_t>(constOp.value()));
+ } else {
+ return rewriter.notifyMatchFailure(op, "Indices must be constants");
----------------
SoftwareBoi wrote:
```suggestion
return rewriter.notifyMatchFailure(op, "indices must be constants");
```
https://github.com/llvm/llvm-project/pull/139048
More information about the Mlir-commits
mailing list