Add Int4 support for tosa::ConstOp

[Jerry-Ge]

Summary

• Also added Tosa_Weight and Tosa_WeightTensorXD specifically for weights

Signed-off-by: Jerry Ge <[e-mail]>

[Jerry-Ge]

added a simple lit test for this change

[eric-k256]

LGTM.

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

def Tosa_Conv2DOp : Tosa_Op<"conv2d", [
 
   let description = [{
     Performs a 2D convolution over the given tensor input, using the weight
     tensor.
   }];
 
   let arguments = (ins
     Tosa_Tensor4D:$input,
-    Tosa_Tensor4D:$weight,
+    4DTensorOf<[Tosa_Weight]>:$weight,
     Tosa_Tensor1D:$bias,
 
     Tosa_IntArrayAttr4:$pad,
     Tosa_IntArrayAttr2:$stride,
     Tosa_IntArrayAttr2:$dilation,
     OptionalAttr<Tosa_ConvOpQuantizationAttr>:$quantization_info
   );
 
def Tosa_Conv3DOp : Tosa_Op<"conv3d", [
   let summary = "3D Convolution operator";
 
   let description = [{
     Performs a 3D convolution over the given input tensor.
   }];
 
   let arguments = (ins
     Tosa_Tensor5D:$input,
-    Tosa_Tensor5D:$weight,
+    TensorRankOf<[Tosa_Weight], [5]>:$weight,
     Tosa_Tensor1D:$bias,
 
     Tosa_IntArrayAttr6:$pad,
     Tosa_IntArrayAttr3:$stride,
     Tosa_IntArrayAttr3:$dilation,
     OptionalAttr<Tosa_ConvOpQuantizationAttr>:$quantization_info
   );
 
def Tosa_DepthwiseConv2DOp : Tosa_Op<"depthwise_conv2d", [
 
   let description = [{
     Performs 2D convolutions separately over each channel of the given tensor
     input, using the weight tensor.
   }];
 
   let arguments = (ins
     Tosa_Tensor4D:$input,
-    Tosa_Tensor4D:$weight,
+    4DTensorOf<[Tosa_Weight]>:$weight,
     Tosa_Tensor1D:$bias,
 
     Tosa_IntArrayAttr4:$pad,
     Tosa_IntArrayAttr2:$stride,
     Tosa_IntArrayAttr2:$dilation,
     OptionalAttr<Tosa_ConvOpQuantizationAttr>:$quantization_info
   );
 
def Tosa_FullyConnectedOp : Tosa_Op<"fully_connected", [
   let summary = "Fully Connected operator";
 
   let description = [{
     Performs a fully connected network.
   }];
 
   let arguments = (ins
     Tosa_Tensor2D:$input,
-    Tosa_Tensor2D:$weight,
+    2DTensorOf<[Tosa_Weight]>:$weight,
     Tosa_Tensor1D:$bias,
     OptionalAttr<Tosa_ConvOpQuantizationAttr>:$quantization_info
   );
 
   let results = (outs
     Tosa_Tensor2D:$output
   );
 
def Tosa_TransposeConv2DOp : Tosa_Op<"transpose_conv2d", [
 
   let description = [{
     Performs a 2D transposed convolution over the given tensor input, using the
     weights tensor.
   }];
 
   let arguments = (ins
     Tosa_Tensor4D:$input,
-    Tosa_Tensor4D:$filter,
+    4DTensorOf<[Tosa_Weight]>:$filter,
     Tosa_Tensor1D:$bias,
 
     Tosa_IntArrayAttr4:$out_pad,
     Tosa_IntArrayAttr2:$stride,
     Tosa_IntArrayAttrUpto4:$out_shape,
     OptionalAttr<Tosa_ConvOpQuantizationAttr>:$quantization_info
   );
 
let description = [{
     hold data in any of the supported data formats.
   }];
 
   let arguments = (ins
     ElementsAttr:$value
   );
 
   let results = (outs
-    Tosa_Tensor_Plus_F64:$output
+    TensorOf<[AnyTypeOf<[Tosa_AnyNumber_Plus_F64, Tosa_Int4]>]>:$output
   );
   let hasFolder = 1;
 }
 
 //===----------------------------------------------------------------------===//
 // Operator: identity
 //===----------------------------------------------------------------------===//
 def Tosa_IdentityOp: Tosa_Op<"identity", [Pure,

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

 //===----------------------------------------------------------------------===//
 // Non-Quantized Signed Integer Types.
 // Used to express accumulator results or compare results.
 //===----------------------------------------------------------------------===//
 
 def Tosa_UInt8 : UI<8>;
 def Tosa_UInt16 : UI<16>;
 
+def Tosa_Int4 : I<4>;
 def Tosa_Int8 : I<8>;
 def Tosa_Int16 : I<16>;
 def Tosa_Int32 : I<32>;
 def Tosa_Int48 : I<48>;
 def Tosa_Int64 : I<64>;
 
 def Tosa_SignedInt : AnyTypeOf<[Tosa_Int8,
                                 Tosa_Int16,
def Tosa_Float : AnyTypeOf<[
           F16,
           BF16]>;
 
 //===----------------------------------------------------------------------===//
 // Multi-category types.
 //===----------------------------------------------------------------------===//
 def Tosa_AnyNumber : AnyTypeOf<[Tosa_Int, Tosa_QuantizedInt, Tosa_Float],
                                "number">;
+
 // Add F64 type support just for tosa::CastOp and tosa::ConstOp
 def Tosa_AnyNumber_Plus_F64 : AnyTypeOf<[Tosa_Int, Tosa_QuantizedInt, Tosa_Float, F64],
                                "number_plus_f64">;
 
+// For weight tensors from tosa::Conv2DOp, tosa::Conv3DOp,
+// tosa::DepthwiseConv2DOp, tosa::TransposeConv2DOp, tosa::FullyConnectedOp
+def Tosa_Weight : AnyTypeOf<[Tosa_Int4, Tosa_Int8,
+                             Tosa_QuantizedInt, Tosa_Float]>;
+
 //===----------------------------------------------------------------------===//
 // Tensor types
 //===----------------------------------------------------------------------===//
 
 def Tosa_Int32Tensor : TensorOf<[Tosa_Int32]>;
 def Tosa_Int32Or64Tensor : TensorOf<[Tosa_Int32Or64]>;
 
 // Either ranked or unranked tensor of TOSA supported element types.
 def Tosa_Tensor : TensorOf<[Tosa_AnyNumber]>;
 def Tosa_Tensor_Plus_F64 : TensorOf<[Tosa_AnyNumber_Plus_F64]>;
+
 // Must be ranked but no further constraints
 def Tosa_RankedTensor : RankedTensorOf<[Tosa_AnyNumber]>;
 
 // Any tensor element type allowed in Tosa ops.
 def Tosa_ElementType : Type<Or<[Tosa_Int.predicate, Tosa_QuantizedInt.predicate,
                                 Tosa_Float.predicate]>, "tosa.dtype">;
 
 class Tosa_TensorOfOrNone<list<Type> allowedTypes, string description = ""> :

mlir/test/Dialect/Tosa/invalid.mlir

func.func @test_conv2d(%arg0: tensor<*xi8>, %arg1: tensor<16x3x3x4xi8>, %arg2: tensor<16xi8>) -> tensor<1x27x27x16xi8> {
   %0 = "tosa.conv2d"(%arg0, %arg1, %arg2) {dilation = array<i64: 1, 1>, pad = array<i64: 0, 0, 0, 0>, stride = array<i64: 1, 1>}
            : (tensor<*xi8>, tensor<16x3x3x4xi8>, tensor<16xi8>) -> tensor<1x27x27x16xi8>
   return %0 : tensor<1x27x27x16xi8>
 }
 
 // -----
 
 func.func @test_conv2d(%arg0: tensor<1x29x29x4xi8>, %arg1: tensor<*xi8>, %arg2: tensor<16xi8>) -> tensor<1x27x27x16xi8> {
-  // expected-error@+1 {{expect a ranked tensor for weight, got <block argument> of type 'tensor<*xi8>' at index: 1}}
+  // expected-error@+1 {{'tosa.conv2d' op operand #1 must be 4D tensor of 4-bit signless integer or 8-bit signless integer or Quint8 type or Qint4 type or Qint8 type or Qint16 type or Qint32 type or 32-bit float or 16-bit float or bfloat16 type values, but got 'tensor<*xi8>'}}
   %0 = "tosa.conv2d"(%arg0, %arg1, %arg2) {dilation = array<i64: 1, 1>, pad = array<i64: 0, 0, 0, 0>, stride = array<i64: 1, 1>}
            : (tensor<1x29x29x4xi8>, tensor<*xi8>, tensor<16xi8>) -> tensor<1x27x27x16xi8>
   return %0 : tensor<1x27x27x16xi8>
 }
 
-
 // -----
 
 func.func @test_conv2d(%arg0: tensor<1x29x29x4xi8>, %arg1: tensor<16x3x3x4xi8>, %arg2: tensor<16xi8>) -> tensor<1x27x27x16xi8> {
   // expected-error@+1 {{'tosa.conv2d' op quantizationattr is required for quantized type, and not allowed for float type}}
   %0 = "tosa.conv2d"(%arg0, %arg1, %arg2) {dilation = array<i64: 1, 1>, pad = array<i64: 0, 0, 0, 0>, stride = array<i64: 1, 1>}
            : (tensor<1x29x29x4xi8>, tensor<16x3x3x4xi8>, tensor<16xi8>) -> tensor<1x27x27x16xi8>
   return %0 : tensor<1x27x27x16xi8>
 }

mlir/test/Dialect/Tosa/ops.mlir

 // -----
 // CHECK-LABEL: conv2d
 func.func @test_conv2d(%arg0: tensor<1x4x4x4xf32>, %arg1: tensor<8x1x1x4xf32>, %arg2: tensor<8xf32>) -> tensor<1x4x4x8xf32> {
     %0 = "tosa.conv2d"(%arg0, %arg1, %arg2) {dilation = array<i64: 1, 1>, pad = array<i64: 0, 0, 0, 0>, stride = array<i64: 1, 1>} : (tensor<1x4x4x4xf32>, tensor<8x1x1x4xf32>, tensor<8xf32>) -> tensor<1x4x4x8xf32>
     return %0 : tensor<1x4x4x8xf32>
 }
 
 // -----
+// CHECK-LABEL: conv2d_q8xi4
+func.func @test_conv2d_q8xi4(%arg0: tensor<1x11x11x3xi8>) -> tensor<1x1x1x3xi8> {
+  %0 = "tosa.const"() {value = dense<0> : tensor<3x11x11x3xi4>} : () -> tensor<3x11x11x3xi4>
+  %1 = "tosa.const"() {value = dense<[12, 23, 55]> : tensor<3xi32>} : () -> tensor<3xi32>
+  %2 = "tosa.conv2d"(%arg0, %0, %1) {dilation = array<i64: 1, 1>, pad = array<i64: 0, 0, 0, 0>, quantization_info = #tosa.conv_quant<input_zp = 0, weight_zp = 0>, stride = array<i64: 1, 1>} : (tensor<1x11x11x3xi8>, tensor<3x11x11x3xi4>, tensor<3xi32>) -> tensor<1x1x1x3xi32>
+  %3 = "tosa.rescale"(%2) {double_round = true, input_zp = 0 : i32, multiplier = array<i32: 2026291432, 1079222024, 1693132724>, output_zp = 27 : i32, per_channel = true, scale32 = true, shift = array<i32: 37, 36, 37>} : (tensor<1x1x1x3xi32>) -> tensor<1x1x1x3xi8>
+  return %3 : tensor<1x1x1x3xi8>
+}
+
+// -----
 // CHECK-LABEL: depthwise_conv2d
 func.func @test_depthwise_conv2d(%arg0: tensor<1x4x4x4xf32>, %arg1: tensor<1x1x4x2xf32>, %arg2: tensor<8xf32>) -> tensor<1x4x4x8xf32> {
   %2 = "tosa.depthwise_conv2d"(%arg0, %arg1, %arg2) {dilation = array<i64: 1, 1>, pad = array<i64: 0, 0, 0, 0>, stride = array<i64: 1, 1>} : (tensor<1x4x4x4xf32>, tensor<1x1x4x2xf32>, tensor<8xf32>) -> tensor<1x4x4x8xf32>
   return %2 : tensor<1x4x4x8xf32>
 }
 
 // -----
 // CHECK-LABEL: fft2d