[Mlir-commits] [mlir] 05cadc6 - [mlir][tosa] Temporarily support 2D and 3D tensor types in matmul

[Rob Suderman]

Author: Suraj Sudhir
Date: 2021-06-08T16:03:59-07:00
New Revision: 05cadc6f71555319882ccabf631d2e6410e3fea4

URL: https://github.com/llvm/llvm-project/commit/05cadc6f71555319882ccabf631d2e6410e3fea4
DIFF: https://github.com/llvm/llvm-project/commit/05cadc6f71555319882ccabf631d2e6410e3fea4.diff

LOG: [mlir][tosa] Temporarily support 2D and 3D tensor types in matmul

Temporarily support 2D and 3D while the TOSA Matmul op is updated to support batched operations.

Reviewed By: rsuderman

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

Added: 
    

Modified: 
    mlir/include/mlir/Dialect/Tosa/IR/TosaOps.td
    mlir/include/mlir/Dialect/Tosa/IR/TosaTypesBase.td

Removed: 
    


################################################################################
diff  --git a/mlir/include/mlir/Dialect/Tosa/IR/TosaOps.td b/mlir/include/mlir/Dialect/Tosa/IR/TosaOps.td
index 06c2ed480a152..2a5e2baa44e43 100644
--- a/mlir/include/mlir/Dialect/Tosa/IR/TosaOps.td
+++ b/mlir/include/mlir/Dialect/Tosa/IR/TosaOps.td
@@ -208,13 +208,13 @@ def Tosa_MatMulOp : Tosa_Op<"matmul", [NoSideEffect]> {
   }];
 
   let arguments = (ins
-    Tosa_Tensor2D:$a,
-    Tosa_Tensor2D:$b,
+    Tosa_Tensor2Dto3D:$a,
+    Tosa_Tensor2Dto3D:$b,
     OptionalAttr<Tosa_MatMulOpQuantizationAttr>:$quantization_info
   );
 
   let results = (outs
-    Tosa_Tensor2D:$c
+    Tosa_Tensor2Dto3D:$c
   );
 
   let builders = [Tosa_MatMulOpQuantInfoBuilder];

diff  --git a/mlir/include/mlir/Dialect/Tosa/IR/TosaTypesBase.td b/mlir/include/mlir/Dialect/Tosa/IR/TosaTypesBase.td
index 08324a15a07b5..5969d98408a9e 100644
--- a/mlir/include/mlir/Dialect/Tosa/IR/TosaTypesBase.td
+++ b/mlir/include/mlir/Dialect/Tosa/IR/TosaTypesBase.td
@@ -124,6 +124,8 @@ def Tosa_Tensor5D : TensorRankOf<[Tosa_AnyNumber], [5]>;
 def Tosa_Tensor1Dto4D : TensorRankOf<[Tosa_AnyNumber], [1,2,3,4]>;
 def Tosa_Tensor1Dto6D : TensorRankOf<[Tosa_AnyNumber], [1,2,3,4,5,6]>;
 
+def Tosa_Tensor2Dto3D : TensorRankOf<[Tosa_AnyNumber], [2,3]>;
+
 def Tosa_TensorUpto4D : TensorRankOf<[Tosa_AnyNumber], [0,1,2,3,4]>;
 
 def Tosa_Int32TensorUpto4D : TensorRankOf<[Tosa_Int32], [0,1,2,3,4]>;