[mlir][TOSA]Add optional attributes to TOSA custom op

[eric-k256]

Summary

The implementation_attrs attr allows passing of backend specific
attributes to TOSA custom ops.

Also adds a config option to avoid namespace collisions on the
identifier.

[silvas]

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

1803

Can we make this more flexible than just a string? I really don't want to have to stringify a bunch of metadata. A typical example from Torch-MLIR is

def Torch_AtenPadOp : Torch_Op<"aten.pad", [
    AllowsTypeRefinement,
    HasValueSemantics,
    ReadOnly
  ]> {
  let summary = "Generated op for `aten::pad : (Tensor, int[], str, float?) -> (Tensor)`";
  let arguments = (ins
    AnyTorchTensorType:$self,
    AnyTorchListOfTorchIntType:$pad,
    Torch_StringType:$mode,
    AnyTorchOptionalFloatType:$value
  );

So we would need a list of integers, a string mode, and an optional float. My ideal representation for this would be something like an ArrayAttr of (dense array attr $pad, StringAttr $mode, FloatAttr | "none") -- I'm not sure about the representation for "none".

The alternative would be to JSONify everything :/

[eric-k256]

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

1803

I thought a bit about options for the attributes when making the change. I can certainly see that string may not be the best way. My other thought was to do a DenseI8ArrayAttr, which makes things slightly easier than strings. It still forces the frontend to encode the attributes in a way easily digested by the backend. That encoding can define the representation of none. From there, going to your idea of an ArrayAttr isn't too big a change. Are there places where an attribute doesn't pack neatly into an existing attribute type? Your example of none comes to mind. Could we use OpaqueAttr here?

Whatever the solution is, the contents of the attributes will still be opaque to TOSA itself. In your example, the backend consuming a tosa.cusom of 'aten.pad' needs to know that attribute 0 is the pad list, 1 is mode, etc. Whether that's done through an ArrayAttr or DenseI8ArrayAttr, that doesn't change.

[silvas]

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

1803

Does using MLIR native attributes affect flatbuffer serialization of the TOSA MLIR? I think that using JSON might actually be better in that case, even if it is more awkward.

[eric-k256]

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

1803

It creates some issues, but they are solvable. We would need to put an encoding scheme in the serialization to encode any MLIR attribute. Anyone consuming the serialized form would need to be able to undo the encoding.

In the prototype serialization code I wrote, I'm already writing the string out as a byte array, as MLIR strings can have embedded null bytes, but flatbuffers doesn't. It's easier to write/read as bytes rather than spending the effort to stringify.

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

def Tosa_CustomOp : Tosa_Op<"custom"> {
 
   let summary = "Custom operator wrapper for Tosa";
 
   let description = [{
     Hardware implementing TOSA may choose to add additional custom operators
     that are not expressed in the existing TOSA operations. These operators are
     not expected to be portable across TOSA implementations. The input and
     output signatures must be expressed in the corresponding TOSA node.
+
+    `identifier` is a string that tells the backend which custom operator is being
+    called.
+
+    `config` is a string identifier which can help avoid name collisions on the
+    identifier field.
+
+    `implementation_attrs` is a string which is a backend and identifier specific
+    set of attributes to the custom operator.
+
+    `inputs` is the set of tensor inputs to the custom operator.
+
+    `outputs is the list of tensors returned by the operator. The number of operators
+    is backend specific.
   }];
 
   let arguments = (ins
     StrAttr:$identifier,
+    StrAttr:$config,
+    StrAttr:$implementation_attrs,
     Variadic<Tosa_Tensor>:$inputs
   );
 
   let results = (outs
     Variadic<Tosa_Tensor>:$outputs
   );
 }
 

mlir/test/Dialect/Tosa/ops.mlir

^bb0(%arg2: tensor<i32>, %arg3: tensor<i32>, %arg4: tensor<10xi32>):
     %3 = "tosa.add"(%arg3, %2) : (tensor<i32>, tensor<i32>) -> tensor<i32>
     %4 = "tosa.reshape"(%2) {new_shape = [1]} : (tensor<i32>) -> tensor<1xi32>
     %5 = "tosa.add"(%arg4, %4) : (tensor<10xi32>, tensor<1xi32>) -> tensor<10xi32>
     %6 = "tosa.add"(%arg2, %2) : (tensor<i32>, tensor<i32>) -> tensor<i32>
     "tosa.yield"(%6, %3, %5) : (tensor<i32>, tensor<i32>, tensor<10xi32>) -> ()
   }) : (tensor<i32>, tensor<i32>, tensor<10xi32>) -> (tensor<i32>, tensor<i32>, tensor<10xi32>)
   return
 }
+
+// -----
+// CHECK-LABEL: custom
+func.func @test_custom(%arg0: tensor<10xi32>) -> tensor<10xi32> {
+  %0 = "tosa.custom"(%arg0) {identifier="custom_test", config="tosa_mlir_test", implementation_attrs=""} : (tensor<10xi32>) -> (tensor<10xi32>)
+  return %0 : tensor<10xi32>
+}