[Mlir-commits] [mlir] 50882b4 - [mlir] List more elementwise ops in VectorToGPU MMA conversion

[Lei Zhang]

Author: Lei Zhang
Date: 2022-12-05T22:51:19Z
New Revision: 50882b4daf77b9d93e025f804b0855c94a83f237

URL: https://github.com/llvm/llvm-project/commit/50882b4daf77b9d93e025f804b0855c94a83f237
DIFF: https://github.com/llvm/llvm-project/commit/50882b4daf77b9d93e025f804b0855c94a83f237.diff

LOG: [mlir] List more elementwise ops in VectorToGPU MMA conversion

Reviewed By: ThomasRaoux

Differential Revision: https://reviews.llvm.org/D139244

Added: 
    

Modified: 
    mlir/lib/Conversion/VectorToGPU/VectorToGPU.cpp

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Conversion/VectorToGPU/VectorToGPU.cpp b/mlir/lib/Conversion/VectorToGPU/VectorToGPU.cpp
index 937913445f30c..85492f0e0d295 100644
--- a/mlir/lib/Conversion/VectorToGPU/VectorToGPU.cpp
+++ b/mlir/lib/Conversion/VectorToGPU/VectorToGPU.cpp
@@ -191,12 +191,26 @@ convertElementwiseOpToMMA(Operation *op) {
     return gpu::MMAElementwiseOp::ADDF;
   if (isa<arith::MulFOp>(op))
     return gpu::MMAElementwiseOp::MULF;
+  if (isa<arith::SubFOp>(op))
+    return gpu::MMAElementwiseOp::SUBF;
   if (isa<arith::MaxFOp>(op))
     return gpu::MMAElementwiseOp::MAXF;
   if (isa<arith::MinFOp>(op))
     return gpu::MMAElementwiseOp::MINF;
   if (isa<arith::DivFOp>(op))
     return gpu::MMAElementwiseOp::DIVF;
+  if (isa<arith::AddIOp>(op))
+    return gpu::MMAElementwiseOp::ADDI;
+  if (isa<arith::MulIOp>(op))
+    return gpu::MMAElementwiseOp::MULI;
+  if (isa<arith::SubIOp>(op))
+    return gpu::MMAElementwiseOp::SUBI;
+  if (isa<arith::DivSIOp>(op))
+    return gpu::MMAElementwiseOp::DIVS;
+  if (isa<arith::DivUIOp>(op))
+    return gpu::MMAElementwiseOp::DIVU;
+  if (isa<arith::NegFOp>(op))
+    return gpu::MMAElementwiseOp::NEGATEF;
   return std::nullopt;
 }