Allow operands and results of tosa.custom to be None. This enables lowering from dialects that support None operands or results

Author: jorickert

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

@@ -2057,11 +2057,11 @@ def Tosa_CustomOp : Tosa_Op<"custom"> {
     StrAttr:$operator_name,
     StrAttr:$domain_name,
     StrAttr:$implementation_attrs,
-    Variadic<Tosa_Tensor>:$input_list
+    Variadic<Tosa_TensorOrNone>:$input_list
   );
 
   let results = (outs
-    Variadic<Tosa_Tensor>:$output_list
+    Variadic<Tosa_TensorOrNone>:$output_list
   );
 
   let assemblyFormat = "operands attr-dict `:` functional-type(operands, results)";

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

@@ -135,6 +135,8 @@ def Tosa_ElementType : Type<Or<[Tosa_Int.predicate, Tosa_QuantizedInt.predicate,
 class Tosa_TensorOfOrNone<list<Type> allowedTypes, string description = ""> :
   AnyTypeOf<[TosaTensorOf<allowedTypes>, NoneType], description>;
 
+def Tosa_TensorOrNone : Tosa_TensorOfOrNone<[Tosa_AnyNumber]>;
+
 //===----------------------------------------------------------------------===//
 // Tensor types with constrained ranks.
 //===----------------------------------------------------------------------===//

mlir/test/Dialect/Tosa/ops.mlir

@@ -694,6 +694,13 @@ func.func @test_custom(%arg0: tensor<10xi32>) -> tensor<10xi32> {
   return %0 : tensor<10xi32>
 }
 
+// -----
+// CHECK-LABEL: test_custom_none
+func.func @test_custom_none(%arg0: tensor<10xi32>, %arg1: none) -> tensor<10xi32> {
+  %0 = tosa.custom %arg0, %arg1 {operator_name="custom_test", domain_name="tosa.mlir_test", implementation_attrs="" } : (tensor<10xi32>, none) -> (tensor<10xi32>)
+  return %0 : tensor<10xi32>
+}
+
 // -----
 // CHECK-LABEL: const_shape
 func.func @test_const_shape() -> !tosa.shape<4> {