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

[sjarus]

Summary

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

[rsuderman]

Something appears to be wrong with your patch as its based on a revision not in the LLVM tree. Could you verify your diff history or resubmit based on head?

[sjarus]

In D103854#2806111, @rsuderman wrote:

Something appears to be wrong with your patch as its based on a revision not in the LLVM tree. Could you verify your diff history or resubmit based on head?

Yeah I noticed I did arc diff from a different branch.

Edit: I think I had a stale branch.

[sjarus]

Temporarily support 2D and 3D tensor types as valid.
Legalizations and tests across multiple repositories must be updated to emit 3D type.
Change to supporting just Tosa_Tensor3D once everything is updated.

[rsuderman]

Looks good, I am just checking why the windows build failed (and whether its unrelated).

[rsuderman]

Updated description as external projects are not relevant for LLVM updates.

[sjarus]

In D103854#2806565, @rsuderman wrote:

Looks good, I am just checking why the windows build failed (and whether its unrelated).

Thanks a lot for your assistance!

mlir/include/mlir/Dialect/Tosa/IR/TosaOps.td

def Tosa_MatMulOp : Tosa_Op<"matmul", [NoSideEffect]> {
 
   let description = [{
     Performs a two dimensional matrix multiplication. This allows both inputs to
     be activations, rather than reserving weights as an attribute in the
     FULLY_CONNECTED operator.
   }];
 
   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];
 }
 
 //===----------------------------------------------------------------------===//
 // Operator: max_pool2d
 //===----------------------------------------------------------------------===//

mlir/include/mlir/Dialect/Tosa/IR/TosaTypesBase.td

 def Tosa_Tensor6D : TensorRankOf<[Tosa_AnyNumber], [6]>;
 
 // Ranked tensors up to given rank.
 def Tosa_Tensor1Dto2D : TensorRankOf<[Tosa_AnyNumber], [1,2]>;
 def Tosa_Tensor1Dto4D : TensorRankOf<[Tosa_AnyNumber], [1,2,3,4]>;
 def Tosa_Tensor1Dto5D : TensorRankOf<[Tosa_AnyNumber], [1,2,3,4,5]>;
 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_TensorUpto6D : TensorRankOf<[Tosa_AnyNumber], [0,1,2,3,4,5,6]>;
 
 def Tosa_Int32TensorUpto4D : TensorRankOf<[Tosa_Int32], [0,1,2,3,4]>;
 
 //===----------------------------------------------------------------------===//
 // Generic scalar, vector, or tensor of a particular type.
 //===----------------------------------------------------------------------===//