[mlir][tosa] Spec v0.23 updates

[sjarus]

Summary

Add pad_const field to tosa.pad.
Add builders to enable optional construction of pad_const in pad op.
Update documentation of tosa.clamp to match spec wording.

Signed-off-by: Suraj Sudhir <[e-mail]>

[sjarus]

This change keeps the new pad_const field an optional construct and adds a builder to handle it.

Existing frontend legalizations do not break. They can be updated later and the parameter optionality can then be removed without breaking anything.

[rsuderman]

Okay, I will review and integrate this on Monday to guarantee there are not any cascading failures. Looks really great though. I was just starting work on something where this is very useful.

[sjarus]

In D113322#3112990, @rsuderman wrote:

Okay, I will review and integrate this on Monday to guarantee there are not any cascading failures. Looks really great though. I was just starting work on something where this is very useful.

Sounds good! We ran bazel test in the TF tree and all the TOSA stuff was ok, but there are other pieces there outside of our control.

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

 // relationship between their input and output, expressed by the
 // UnaryOpQuantizationAttr.
 def Tosa_UnaryOpQuantInfoBuilder : OpBuilder<
   (ins "Type":$outputType, "Value":$input),
   [{
     buildUnaryOpWithQuantInfo($_builder, $_state, outputType, input);
   }]>;
 
-// This builder is called on the TOSA pad operator that needs to create its own
-// OptionalAttr quantization_attr parameter to scale the padding values
+// These builders are called on the TOSA pad operator that needs to create its
+// own OptionalAttr quantization_attr parameter to scale the padding values
 // correctly.
 def Tosa_PadOpQuantInfoBuilder : OpBuilder<
   (ins "Type":$outputType, "Value":$input, "Value":$paddings),
   [{
     buildPadOpWithQuantInfo($_builder, $_state, outputType,
                             input, paddings);
   }]>;
 
+def Tosa_ExplicitValuePadOpQuantInfoBuilder : OpBuilder<
+  (ins "Type":$outputType, "Value":$input, "Value":$paddings,
+       "Value":$pad_value),
+  [{
+    buildExplicitValuePadOpWithQuantInfo($_builder, $_state, outputType,
+                                         input, paddings, pad_value);
+  }]>;
+
 //===----------------------------------------------------------------------===//
 // TOSA Operator.
 //===----------------------------------------------------------------------===//
 
 class Tosa_Op<string mnemonic, list<OpTrait> traits = []> :
     Op<Tosa_Dialect, mnemonic, !listconcat(traits, [TosaOpInterface])> {
 }
 
 #endif // TOSA_OP_BASE

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

 //===----------------------------------------------------------------------===//
 def Tosa_ClampOp : Tosa_Op<"clamp", [
     DeclareOpInterfaceMethods<InferShapedTypeOpInterface,
                               ["inferReturnTypeComponents"]>,
     NoSideEffect]> {
   let summary = "Computes clamp(features, min, max).";
 
   let description = [{
-    Clamp to an arbitrary minimum and maximum value. Note that the maximum and
-    minimum values are specified as signed quantized values, no scaling happens
-    before or after this operation.
+    Clamp to an arbitrary minimum and maximum value.
+    Maximum and minimum values are specified as values in the range of the 
+    input type.
+    No zero point subtraction is done to the values, thus to clamp to the zero 
+    point value, the zero point itself should be supplied as the minimum value.
   }];
 
   let arguments = (ins
     Tosa_Tensor:$input,
     I64Attr:$min_int,
     I64Attr:$max_int,
     F32Attr:$min_fp,
     F32Attr:$max_fp

 
 //===----------------------------------------------------------------------===//
 // Operator: pad
 //===----------------------------------------------------------------------===//
 def Tosa_PadOp : Tosa_Op<"pad", [
     DeclareOpInterfaceMethods<InferShapedTypeOpInterface,
                               ["inferReturnTypeComponents"]>,
     NoSideEffect]> {
-  let summary = "Pads a tensor with zeros.";
+  let summary = "Pads a tensor with value specified.";
 
   let description = [{
-    Zero-pads a tensor along borders of each dimension.
+    Pads a tensor along borders of each dimension with pad_value. 
   }];
 
   let arguments = (ins
     Tosa_RankedTensor:$input1,
     Tosa_Int32Or64Tensor:$padding,
+    Optional<Tosa_ScalarTensor>:$pad_const,
     OptionalAttr<Tosa_PadOpQuantizationAttr>:$quantization_info
   );
 
   let results = (outs
     Tosa_RankedTensor:$output
   );
 
-  let builders = [Tosa_PadOpQuantInfoBuilder];
+  let builders = [Tosa_PadOpQuantInfoBuilder, 
+                  Tosa_ExplicitValuePadOpQuantInfoBuilder];
 }
 
 //===----------------------------------------------------------------------===//
 // Operator: reshape
 //===----------------------------------------------------------------------===//
 def Tosa_ReshapeOp: Tosa_Op<"reshape", [
     DeclareOpInterfaceMethods<InferShapedTypeOpInterface,
                               ["inferReturnTypeComponents"]>,

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

 
 class Tosa_TensorOfOrNone<list<Type> allowedTypes, string description = ""> :
   AnyTypeOf<[TensorOf<allowedTypes>, NoneType], description>;
 
 //===----------------------------------------------------------------------===//
 // Tensor types with constrained ranks.
 //===----------------------------------------------------------------------===//
 
+// Rank-0 (scalar) tensor
+def Tosa_ScalarTensor : TensorRankOf<[Tosa_AnyNumber], [0]>;
+
 // We include unranked tensors as a supported type for all possible tosa
 // Tensors as unranked does not guarantee invalid. If unranked tensors exist
 // they should be shape propagate used Tosa's shape inference pass and verified
 // to not include any remaining unranked tensors.
 def Tosa_UnrankedTensor : UnrankedTensorOf<[Tosa_AnyNumber]>;
 
 def Tosa_Tensor1D : AnyTypeOf<[Tosa_UnrankedTensor, 1DTensorOf<[Tosa_AnyNumber]>]>;
 def Tosa_Tensor2D : AnyTypeOf<[Tosa_UnrankedTensor, 2DTensorOf<[Tosa_AnyNumber]>]>;

mlir/lib/Dialect/Tosa/IR/TosaOps.cpp

static void buildUnaryOpWithQuantInfo(OpBuilder &builder,
   auto quantAttr = buildUnaryOpQuantizationAttr(builder, input, outputType);
   if (quantAttr)
     result.addAttribute("quantization_info", quantAttr);
   result.types.push_back(outputType);
 }
 
 /// This builder is called on TOSA pad operator that needs to create its own
 /// OptionalAttr quantization_attr parameter to scale the padding values
-/// correctly.
+/// correctly. No pad_const is interpreted as zero-padding.
 static void buildPadOpWithQuantInfo(OpBuilder &builder, OperationState &result,
                                     Type outputType, Value input,
                                     Value paddings) {
   result.addOperands({input, paddings});
   auto quantAttr = buildPadOpQuantizationAttr(builder, input);
   if (quantAttr)
     result.addAttribute("quantization_info", quantAttr);
   result.types.push_back(outputType);
 }
 
+/// This builder is called on TOSA pad operator when an explicit pad_const
+/// value is passed in. It also optionally constructs quantization_attr.
+static void buildExplicitValuePadOpWithQuantInfo(OpBuilder &builder,
+                                                 OperationState &result,
+                                                 Type outputType, Value input,
+                                                 Value paddings,
+                                                 Value pad_const) {
+  result.addOperands({input, paddings, pad_const});
+  auto quantAttr = buildPadOpQuantizationAttr(builder, input);
+  if (quantAttr)
+    result.addAttribute("quantization_info", quantAttr);
+  result.types.push_back(outputType);
+}
+
 //===----------------------------------------------------------------------===//
 // TOSA Operator Return Type Inference.
 //===----------------------------------------------------------------------===//
 
 static void getI64Values(ArrayAttr arrayAttr, SmallVector<int64_t> &values) {
   for (auto it : arrayAttr) {
     values.push_back(it.cast<IntegerAttr>().getValue().getSExtValue());
   }

mlir/test/Dialect/Tosa/ops.mlir

 // -----
 // CHECK-LABEL: pad
 func @test_pad(%arg0: tensor<13x21x3xf32>, %arg1: tensor<3x2xi32>) -> tensor<13x21x3xf32> {
   %0 = "tosa.pad"(%arg0, %arg1) : (tensor<13x21x3xf32>, tensor<3x2xi32>) -> tensor<13x21x3xf32>
   return %0 : tensor<13x21x3xf32>
 }
 
 // -----
+// CHECK-LABEL: pad_explicit_value
+func @test_pad_explicit_value(%arg0: tensor<13x21x3xf32>, %arg1: tensor<3x2xi32>) -> tensor<13x21x3xf32> {
+  %0 = "tosa.const"() {value = dense<3.14> : tensor<f32>} : () -> tensor<f32>
+  %1 = "tosa.pad"(%arg0, %arg1, %0) : (tensor<13x21x3xf32>, tensor<3x2xi32>, tensor<f32>) -> tensor<13x21x3xf32>
+  return %1 : tensor<13x21x3xf32>
+}
+
+// -----
 // CHECK-LABEL: reshape
 func @test_reshape(%arg0: tensor<13x21x3xf32>) -> tensor<1x819xf32> {
   %0 = "tosa.reshape"(%arg0) {new_shape = [1, 819]} : (tensor<13x21x3xf32>) -> tensor<1x819xf32>
   return %0 : tensor<1x819xf32>
 }
 
 // -----
 // CHECK-LABEL: reverse