[mlir][tosa] Add accumulator type attribute to TOSA dialect

[tatwaichong]

Summary

Tosa supports fp16 and fp32 accumulator type for fp16 input, but
no way to tell for computational operators whether accumulator
should be fp16 or fp32 from input type. Add this new attribute
to specify the type.

Set to fp32 by default for now. When the time fp16 is supported,
the accumulator type can be selected based on trade-off between
performance and accuracy.

[rsuderman]

Could you include the corresponding change in TosaToLinalg? It should be a fairly trivial change.

[tatwaichong]

Rebase onto top-of-tree to include the fix that isn't relevant to this patch.

[tatwaichong]

Read the acc type attr in TosaToLinalg.

[tatwaichong]

Rebase to include OUTSIDE fix for bazel build error.

[tatwaichong]

Good spot. Thanks for the reminder. I missed TosaToLinalg part. Reading from the acctype attr is complemented in newer patch. Do I still miss something necessary else?

[tatwaichong]

Rebase to include bazel fix for transform dialect (not relevant to this patch)

[tatwaichong]

Look like the recent build of mlir repo isn't stable.
Have to be careful to pick a commit to rebase.
Now rebase this PR to a commit that can pass bazel build test locally.

[rsuderman]

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

91

Rather than doing an enum for the accumulator type you should be able to use a TypeAttr. Then in a separate verifier you can check that the subtype is correct. It will make the TosaToLinalg pass cleaner as you won't need a switch on the enum.

[tatwaichong]

Define the new accumulator type as TypeAttr rather than EnumAttr.

[tatwaichong]

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

91

I realize this is a awesome suggestion. It makes the change become intuitive and simple.

mlir/include/mlir/Dialect/Tosa/IR/CMakeLists.txt

 add_mlir_dialect(TosaOps tosa)
 add_mlir_doc(TosaOps TosaOps Dialects/ -gen-op-doc)
 add_mlir_interface(TosaInterfaces)
 
 set(LLVM_TARGET_DEFINITIONS TosaOps.td)
+mlir_tablegen(TosaOpsEnums.h.inc -gen-enum-decls)
+mlir_tablegen(TosaOpsEnums.cpp.inc -gen-enum-defs)
+add_public_tablegen_target(MLIRTosaOpsEnumsGen)
+
+set(LLVM_TARGET_DEFINITIONS TosaOps.td)
 mlir_tablegen(TosaAttributes.h.inc -gen-attrdef-decls)
 mlir_tablegen(TosaAttributes.cpp.inc -gen-attrdef-defs)
 add_public_tablegen_target(MLIRTosaAttributesIncGen)

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

def Tosa_MatMulOpQuantInfoBuilder : OpBuilder<
   }]>;
 
 // Both the tosa.avg_pool2d and unary ops use the same
 // UnaruOpQuantizationAttr but the avg_pool operator has its own builder as it
 // has additional parameters not part of the unary ops.
 def Tosa_AvgPool2dOpQuantInfoBuilder : OpBuilder<
   (ins "::mlir::Type":$outputType, "::mlir::Value":$input,
        "::mlir::DenseI64ArrayAttr":$kernel, "::mlir::DenseI64ArrayAttr":$stride,
-       "::mlir::DenseI64ArrayAttr":$pad),
+       "::mlir::DenseI64ArrayAttr":$pad, "AccTypeAttr":$acc_type),
   [{
     buildAvgPool2dOpWithQuantInfo($_builder, $_state, outputType,
-                                  input, kernel, stride, pad);
+                                  input, kernel, stride, pad, acc_type);
   }]>;
 
 // This builder is called on single-parameter unary operators that have a scale
 // relationship between their input and output, expressed by the
 // UnaryOpQuantizationAttr.
 def Tosa_UnaryOpQuantInfoBuilder : OpBuilder<
   (ins "Type":$outputType, "Value":$input),
   [{

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

 #include "mlir/Interfaces/SideEffectInterfaces.h"
 #include "mlir/Interfaces/VectorInterfaces.h"
 
 //===----------------------------------------------------------------------===//
 // TOSA dialect and structs includes.
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/Tosa/IR/TosaOpsDialect.h.inc"
+#include "mlir/Dialect/Tosa/IR/TosaOpsEnums.h.inc"
 
 namespace mlir {
 class PatternRewriter;
 
 namespace tosa {
 
 #include "mlir/Dialect/Tosa/IR/TosaInterfaces.h.inc"
 

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

def Tosa_ArgMaxOp : Tosa_Op<"argmax", [
   );
 
   let results = (outs
     Tosa_TensorUpto4D: $output
   );
 }
 
 //===----------------------------------------------------------------------===//
+// Accumulator types.
+//===----------------------------------------------------------------------===//
+
+def Tosa_AccType : I32EnumAttr<"AccType", "Specify the type of the accumulator",
+                           [Tosa_Enum_I32, Tosa_Enum_F16, Tosa_Enum_F32]> {
+  let genSpecializedAttr = 0;
+  let cppNamespace = "::mlir::tosa";
+}
+
+def Tosa_AccTypeAttr : EnumAttr<Tosa_Dialect, Tosa_AccType, "accumulator_type"> {
+  let assemblyFormat = "`<` $value `>`";
+}
+
+//===----------------------------------------------------------------------===//
 // Operator: avg_pool2d
 //===----------------------------------------------------------------------===//
 def Tosa_AvgPool2dOp : Tosa_Op<"avg_pool2d", [
     DeclareOpInterfaceMethods<InferShapedTypeOpInterface,
                               ["inferReturnTypeComponents"]>,
     Pure]> {
   let summary = "Performs max pooling on the input.";
 
   let description = [{
     This performs an average pooling over the given input tensor. A sliding
     window of size given by <kernel size> is passed over the input tensor, with
     the mean value being placed in the output tensor.
   }];
 
   let arguments = (ins
     Tosa_Tensor4D:$input,
 
     Tosa_IntArrayAttr2:$kernel,
     Tosa_IntArrayAttr2:$stride,
     Tosa_IntArrayAttr4:$pad,
+    Tosa_AccTypeAttr:$acc_type,
     OptionalAttr<Tosa_UnaryOpQuantizationAttr>:$quantization_info
   );
 
   let results = (outs
     Tosa_Tensor4D:$output
   );
 
   let builders = [Tosa_AvgPool2dOpQuantInfoBuilder];

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

 //===-- TosaTypesBase.td - TOSA type definitions -----------*- tablegen -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // This file defines the type definitions for the TOSA dialect.
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef TOSA_TYPES_BASE
 #define TOSA_TYPES_BASE
 
 include "mlir/IR/OpBase.td"
+include "mlir/IR/EnumAttr.td"
 
 //===----------------------------------------------------------------------===//
 // Tosa Type Definitions.
 //===----------------------------------------------------------------------===//
 
 // The base class of a quantized type.
 // Param tuple is: [bitwidth, zeropt, smantissa, sexp, low_end, high_end].
 // Where low and high ends are 0,255 when unsigned, -128,127 when signed, for

 //===----------------------------------------------------------------------===//
 def Tosa_AnyNumber : AnyTypeOf<[Tosa_Int, Tosa_QuantizedInt, Tosa_Float],
                                "number">;
 // Add F64 type support for just for tosa::CastOp
 def Tosa_AnyNumber_Cast : AnyTypeOf<[Tosa_Int, Tosa_QuantizedInt, Tosa_Float, F64],
                                "number_cast">;
 
 //===----------------------------------------------------------------------===//
+// Enumeration.
+//===----------------------------------------------------------------------===//
+
+def Tosa_Enum_I32 : I32EnumAttrCase<"I32", 0, "i32">;
+def Tosa_Enum_F16 : I32EnumAttrCase<"F16", 1, "f16">;
+def Tosa_Enum_F32 : I32EnumAttrCase<"F32", 2, "f32">;
+
+//===----------------------------------------------------------------------===//
 // 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]>;

mlir/lib/Conversion/TosaToLinalg/TosaToLinalgNamed.cpp

LogicalResult matchAndRewrite(tosa::MaxPool2dOp op,
     return success();
   }
 };
 
 class AvgPool2dConverter : public OpRewritePattern<tosa::AvgPool2dOp> {
 public:
   using OpRewritePattern<tosa::AvgPool2dOp>::OpRewritePattern;
 
+  Type getAccumulatorType(PatternRewriter &rewriter,
+                          tosa::AvgPool2dOp op) const {
+    switch (op.getAccType()) {
+    case tosa::AccType::I32:
+      return rewriter.getI32Type();
+    case tosa::AccType::F16:
+      return rewriter.getF16Type();
+    case tosa::AccType::F32:
+      return rewriter.getF32Type();
+    default:
+      (void)rewriter.notifyMatchFailure(
+          op, "tosa.avg_pool ops does not support this accumulator type");
+      return rewriter.getNoneType();
+    }
+  }
+
   LogicalResult matchAndRewrite(tosa::AvgPool2dOp op,
                                 PatternRewriter &rewriter) const final {
     Location loc = op.getLoc();
     Value input = op.getInput();
     ShapedType inputTy = input.getType().cast<ShapedType>();
     Type inElementTy = inputTy.getElementType();
 
     ShapedType resultTy = op.getType().template cast<ShapedType>();
     Type resultETy = op.getType().cast<ShapedType>().getElementType();
 
-    Type accETy =
-        inElementTy.isa<IntegerType>() ? rewriter.getI32Type() : inElementTy;
+    Type accETy = getAccumulatorType(rewriter, op);
     ShapedType accTy = resultTy.clone(accETy);
 
     auto dynamicDimsOr =
         checkHasDynamicBatchDims(rewriter, op, {input, op.getOutput()});
     if (!dynamicDimsOr.has_value())
       return failure();
     SmallVector<Value> dynamicDims = *dynamicDimsOr;
 

mlir/lib/Dialect/Tosa/CMakeLists.txt

 add_mlir_dialect_library(MLIRTosaDialect
   IR/TosaOps.cpp
   IR/TosaCanonicalizations.cpp
   Utils/ConversionUtils.cpp
   Utils/QuantUtils.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/Tosa
 
   DEPENDS
   MLIRTosaAttributesIncGen
   MLIRTosaOpsIncGen
+  MLIRTosaOpsEnumsGen
   MLIRTosaInterfacesIncGen
 
   LINK_LIBS PUBLIC
   MLIRIR
   MLIRDialect
   MLIRCallInterfaces
   MLIRControlFlowInterfaces
   MLIRQuantDialect
   MLIRQuantUtils
   MLIRSideEffectInterfaces
   MLIRTensorDialect
   MLIRVectorInterfaces
   MLIRViewLikeInterface
   )
 
 add_subdirectory(Transforms)

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

static void buildMatMulOpWithQuantInfo(OpBuilder &builder,
   } else {
     result.addTypes(outputType);
   }
 }
 
 /// Both the tosa.avg_pool2d and unary ops use the same UnaruOpQuantizationAttr
 /// but avg_pool operator has its own builder as it has additional parameters
 /// not part of the unary ops.
-static void buildAvgPool2dOpWithQuantInfo(
-    OpBuilder &builder, OperationState &result, Type outputType, Value input,
-    DenseArrayAttr kernel, DenseArrayAttr stride, DenseArrayAttr pad) {
+static void
+buildAvgPool2dOpWithQuantInfo(OpBuilder &builder, OperationState &result,
+                              Type outputType, Value input,
+                              DenseArrayAttr kernel, DenseArrayAttr stride,
+                              DenseArrayAttr pad, AccTypeAttr acc_type) {
   result.addOperands(input);
   result.addAttribute("kernel", kernel);
   result.addAttribute("stride", stride);
   result.addAttribute("pad", pad);
+  result.addAttribute("acc_type", acc_type);
   auto quantAttr = buildUnaryOpQuantizationAttr(builder, input, outputType);
   if (quantAttr)
     result.addAttribute("quantization_info", quantAttr);
   result.types.push_back(outputType);
 }
 
 /// This builder is called on single-parameter unary operators that have scale
 /// relationship between their input and output, expressed by the

 //===----------------------------------------------------------------------===//
 // TOSA Attribute Definitions.
 //===----------------------------------------------------------------------===//
 
 #define GET_ATTRDEF_CLASSES
 #include "mlir/Dialect/Tosa/IR/TosaAttributes.cpp.inc"
 
 //===----------------------------------------------------------------------===//
+// TOSA Enum Definitions.
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Dialect/Tosa/IR/TosaOpsEnums.cpp.inc"
+
+//===----------------------------------------------------------------------===//
 // TOSA Operator Definitions.
 //===----------------------------------------------------------------------===//
 
 #define GET_OP_CLASSES
 #include "mlir/Dialect/Tosa/IR/TosaOps.cpp.inc"

mlir/test/Conversion/TosaToLinalg/tosa-to-linalg-named.mlir

func.func @avg_pool_f32(%arg0: tensor<1x6x34x62xf32>) -> (tensor<1x5x33x62xf32>) {
   // CHECK:   %[[KWIDTH:.+]] = arith.select %[[CMP]], %[[ONE]], %[[END_OFFSET]]
 
   // Divide the summed value by the number of values summed.
   // CHECK:   %[[COUNT:.+]] = arith.muli %[[KHEIGHT]], %[[KWIDTH]]
   // CHECK:   %[[CAST:.+]] = arith.index_cast %[[COUNT]]
   // CHECK:   %[[FLT:.+]] = arith.sitofp %[[CAST]]
   // CHECK:   %[[DIV:.+]] = arith.divf %[[IN]], %[[FLT]]
   // CHECK:   linalg.yield %[[DIV]]
-  %0 = "tosa.avg_pool2d"(%arg0) {pad = array<i64: 1, 1, 1, 1>, kernel = array<i64: 4, 4>, stride = array<i64: 1, 1>} : (tensor<1x6x34x62xf32>)  -> (tensor<1x5x33x62xf32>)
+  %0 = "tosa.avg_pool2d"(%arg0) {acc_type = #tosa.accumulator_type<f32>, pad = array<i64: 1, 1, 1, 1>, kernel = array<i64: 4, 4>, stride = array<i64: 1, 1>} : (tensor<1x6x34x62xf32>)  -> (tensor<1x5x33x62xf32>)
   return %0 : tensor<1x5x33x62xf32>
 }
 
 // -----
 
 // CHECK-LABEL: @avg_pool_i8
 func.func @avg_pool_i8(%arg0: tensor<1x6x34x62xi8>) -> (tensor<1x5x33x62xi8>) {
   // CHECK: %[[GENERIC:.+]] = linalg.generic
func.func @avg_pool_i8(%arg0: tensor<1x6x34x62xi8>) -> (tensor<1x5x33x62xi8>) {
   // CHECK: %[[CMIN:.+]] = arith.constant -128
   // CHECK: %[[CMAX:.+]] = arith.constant 127
   // CHECK: %[[CMP:.+]] = arith.cmpi slt, %[[SCALED]], %[[CMIN]]
   // CHECK: %[[SEL:.+]] = arith.select %[[CMP]], %[[CMIN]], %[[SCALED]]
   // CHECK: %[[CMP:.+]] = arith.cmpi slt, %[[CMAX]], %[[SCALED]]
   // CHECK: %[[CLAMP:.+]] = arith.select %[[CMP]], %[[CMAX]], %[[SEL]]
   // CHECK: %[[TRUNC:.+]] = arith.trunci %[[CLAMP]]
   // CHECK: linalg.yield %[[TRUNC]]
-  %0 = "tosa.avg_pool2d"(%arg0) {pad = array<i64: 1, 1, 1, 1>, kernel = array<i64: 4, 4>, stride = array<i64: 1, 1>} : (tensor<1x6x34x62xi8>)  -> (tensor<1x5x33x62xi8>)
+  %0 = "tosa.avg_pool2d"(%arg0) {acc_type = #tosa.accumulator_type<i32>, pad = array<i64: 1, 1, 1, 1>, kernel = array<i64: 4, 4>, stride = array<i64: 1, 1>} : (tensor<1x6x34x62xi8>)  -> (tensor<1x5x33x62xi8>)
   return %0 : tensor<1x5x33x62xi8>
 }
 
 // -----
 
 // CHECK-LABEL: @avg_pool_dyn
 func.func @avg_pool_dyn(%arg0: tensor<?x6x34x62xf32>) -> (tensor<?x5x33x62xf32>) {
   // CHECK: %[[C0:.+]] = arith.constant 0 : index
   // CHECK: %[[BATCH:.+]] = tensor.dim %arg0, %[[C0]]
   // CHECK: %[[F0:.+]] = arith.constant 0.000000e+00 : f32
   // CHECK: %[[PADDED:.+]] = tensor.pad %arg0 low[0, 1, 1, 0] high[0, 1, 1, 0]
   // CHECK:   tensor.yield %[[F0]]
   // CHECK: %[[F0:.+]] = arith.constant 0.000000e+00 : f32
   // CHECK: %[[EMPTY:.+]] = tensor.empty(%[[BATCH]]) : tensor<?x5x33x62xf32>
   // CHECK: %[[FILL:.+]] = linalg.fill ins(%[[F0]] : f32) outs(%[[EMPTY]] : tensor<?x5x33x62xf32>)
   // CHECK: %[[KERNEL:.+]] = tensor.empty() : tensor<4x4xf32>
   // CHECK: %[[POOL:.+]] = linalg.pooling_nhwc_sum
   // CHECK-SAME: dilations = dense<1> : vector<2xi64>, strides = dense<1> : vector<2xi64>
   // CHECK-SAME: ins(%[[PADDED]], %[[KERNEL]] : tensor<?x8x36x62xf32>, tensor<4x4xf32>)
   // CHECK-SAME: outs(%[[FILL]] : tensor<?x5x33x62xf32>) -> tensor<?x5x33x62xf32>
   // CHECK: %[[EMPTY:.+]] = tensor.empty(%[[BATCH]]) : tensor<?x5x33x62xf32>
   // CHECK: %[[GENERIC:.+]] = linalg.generic
-  %0 = "tosa.avg_pool2d"(%arg0) {pad = array<i64: 1, 1, 1, 1>, kernel = array<i64: 4, 4>, stride = array<i64: 1, 1>} : (tensor<?x6x34x62xf32>)  -> (tensor<?x5x33x62xf32>)
+  %0 = "tosa.avg_pool2d"(%arg0) {acc_type = #tosa.accumulator_type<f32>, pad = array<i64: 1, 1, 1, 1>, kernel = array<i64: 4, 4>, stride = array<i64: 1, 1>} : (tensor<?x6x34x62xf32>)  -> (tensor<?x5x33x62xf32>)
   return %0 : tensor<?x5x33x62xf32>
 }
 
 // -----
 
 // CHECK: #[[$MAP1:.+]] = affine_map<(d0, d1, d2, d3) -> (d3)>
 // CHECK: #[[$MAP2:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>
 

mlir/test/Dialect/Tosa/ops.mlir

 // RUN: mlir-opt %s | mlir-opt | FileCheck %s
 // RUN: mlir-opt %s --mlir-print-op-generic | mlir-opt | FileCheck %s
 
 
 // -----
 // CHECK-LABEL: argmax
 func.func @test_argmax(%arg0: tensor<14x19xf32>) -> tensor<14xi32> {
   %0 = "tosa.argmax"(%arg0) {axis = 1 : i64} : (tensor<14x19xf32>) -> tensor<14xi32>
   return %0 : tensor<14xi32>
 }
 
 // -----
 // CHECK-LABEL: avg_pool2d_f32
 func.func @test_avg_pool2d_f32(%arg0: tensor<1x7x7x9xf32>) -> tensor<1x7x7x9xf32> {
-    %0 = "tosa.avg_pool2d"(%arg0) {kernel = array<i64: 2, 2>, pad = array<i64: 0, 1, 0, 1>, stride = array<i64: 1, 1>} : (tensor<1x7x7x9xf32>) -> tensor<1x7x7x9xf32>
+    %0 = "tosa.avg_pool2d"(%arg0) {acc_type = #tosa.accumulator_type<f32>, kernel = array<i64: 2, 2>, pad = array<i64: 0, 1, 0, 1>, stride = array<i64: 1, 1>} : (tensor<1x7x7x9xf32>) -> tensor<1x7x7x9xf32>
     return %0 : tensor<1x7x7x9xf32>
 }
 
 // -----
 // CHECK-LABEL: avg_pool2d_i8
 func.func @test_avg_pool2d_i8(%arg0: tensor<1x7x7x9xi8>) -> tensor<1x7x7x9xi8> {
-    %0 = "tosa.avg_pool2d"(%arg0) {kernel = array<i64: 2, 2>, pad = array<i64: 0, 1, 0, 1>, stride = array<i64: 1, 1>} : (tensor<1x7x7x9xi8>) -> tensor<1x7x7x9xi8>
+    %0 = "tosa.avg_pool2d"(%arg0) {acc_type = #tosa.accumulator_type<i32>, kernel = array<i64: 2, 2>, pad = array<i64: 0, 1, 0, 1>, stride = array<i64: 1, 1>} : (tensor<1x7x7x9xi8>) -> tensor<1x7x7x9xi8>
     return %0 : tensor<1x7x7x9xi8>
 }
 
 // -----
 // CHECK-LABEL: avg_pool2d_i16
 func.func @test_avg_pool2d_i16(%arg0: tensor<1x7x7x9xi16>) -> tensor<1x7x7x9xi16> {
-    %0 = "tosa.avg_pool2d"(%arg0) {kernel = array<i64: 2, 2>, pad = array<i64: 0, 1, 0, 1>, stride = array<i64: 1, 1>} : (tensor<1x7x7x9xi16>) -> tensor<1x7x7x9xi16>
+    %0 = "tosa.avg_pool2d"(%arg0) {acc_type = #tosa.accumulator_type<i32>, kernel = array<i64: 2, 2>, pad = array<i64: 0, 1, 0, 1>, stride = array<i64: 1, 1>} : (tensor<1x7x7x9xi16>) -> tensor<1x7x7x9xi16>
     return %0 : tensor<1x7x7x9xi16>
 }
 
 // -----
 // CHECK-LABEL: avg_pool2d_q8
 func.func @test_avg_pool2d_q8(%arg0: tensor<1x7x7x9x!quant.uniform<i8:f32, 0.01>>) -> tensor<1x7x7x9x!quant.uniform<i8:f32, 0.01>> {
-    %0 = "tosa.avg_pool2d"(%arg0) {kernel = array<i64: 2, 2>, pad = array<i64: 0, 1, 0, 1>, stride = array<i64: 1, 1>} : (tensor<1x7x7x9x!quant.uniform<i8:f32, 0.01>>) -> tensor<1x7x7x9x!quant.uniform<i8:f32, 0.01>>
+    %0 = "tosa.avg_pool2d"(%arg0) {acc_type = #tosa.accumulator_type<i32>, kernel = array<i64: 2, 2>, pad = array<i64: 0, 1, 0, 1>, stride = array<i64: 1, 1>} : (tensor<1x7x7x9x!quant.uniform<i8:f32, 0.01>>) -> tensor<1x7x7x9x!quant.uniform<i8:f32, 0.01>>
     return %0 : tensor<1x7x7x9x!quant.uniform<i8:f32, 0.01>>
 }
 
 // -----
 // 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>

mlir/test/Dialect/Tosa/tosa-infer-shapes.mlir

func.func @scatter_minimum_static(%arg0 : tensor<?x4x?xi32>, %arg1 : tensor<3x?xi32>, %arg2 : tensor<?x?x5xi32>) {
   return
 }
 
 // -----
 
 // CHECK-LABEL: @test_pool_static
 func.func @test_pool_static(%arg0: tensor<3x5x6x7xf32>) {
   // CHECK: -> tensor<3x2x4x7xf32>
-  %0 = "tosa.avg_pool2d"(%arg0) {kernel = array<i64: 4, 3>, pad = array<i64: 0, 0, 0, 0>, stride = array<i64: 1, 1>} : (tensor<3x5x6x7xf32>) -> tensor<?x?x?x?xf32>
+  %0 = "tosa.avg_pool2d"(%arg0) {acc_type = #tosa.accumulator_type<f32>, kernel = array<i64: 4, 3>, pad = array<i64: 0, 0, 0, 0>, stride = array<i64: 1, 1>} : (tensor<3x5x6x7xf32>) -> tensor<?x?x?x?xf32>
 
   // CHECK: -> tensor<3x2x4x7xf32>
   %1 = "tosa.max_pool2d"(%arg0) {kernel = array<i64: 4, 3>, pad = array<i64: 0, 0, 0, 0>, stride = array<i64: 1, 1>} : (tensor<3x5x6x7xf32>) -> tensor<?x?x?x?xf32>
   return
 }
 
 // -----
 
func.func @conv2d_dynamic_input(%input: tensor<?x?x?x?xf32>, %weights: tensor<5x3x6x3xf32>, %bias: tensor<5xf32>) -> () {
   return
 }
 
 // -----
 
 // CHECK-LABEL: @test_pool_dynamic_input
 func.func @test_pool_dynamic_input(%arg0: tensor<?x?x?x?xf32>) {
   // CHECK: -> tensor<?x?x?x?xf32>
-  %0 = "tosa.avg_pool2d"(%arg0) {kernel = array<i64: 4, 3>, pad = array<i64: 0, 0, 0, 0>, stride = array<i64: 1, 1>} : (tensor<?x?x?x?xf32>) -> tensor<?x?x?x?xf32>
+  %0 = "tosa.avg_pool2d"(%arg0) {acc_type = #tosa.accumulator_type<f32>, kernel = array<i64: 4, 3>, pad = array<i64: 0, 0, 0, 0>, stride = array<i64: 1, 1>} : (tensor<?x?x?x?xf32>) -> tensor<?x?x?x?xf32>
 
   // CHECK: -> tensor<?x?x?x?xf32>
   %1 = "tosa.max_pool2d"(%arg0) {kernel = array<i64: 4, 3>, pad = array<i64: 0, 0, 0, 0>, stride = array<i64: 1, 1>} : (tensor<?x?x?x?xf32>) -> tensor<?x?x?x?xf32>
   return
 }
 
 // -----
 
 // CHECK-LABEL: @test_pool_padded
 func.func @test_pool_padded(%arg0: tensor<3x5x6x7xf32>) {
   // CHECK: -> tensor<3x5x11x7xf32>
-  %0 = "tosa.avg_pool2d"(%arg0) {kernel = array<i64: 4, 3>, pad = array<i64: 1, 2, 3, 4>, stride = array<i64: 1, 1>} : (tensor<3x5x6x7xf32>) -> tensor<?x?x?x?xf32>
+  %0 = "tosa.avg_pool2d"(%arg0) {acc_type = #tosa.accumulator_type<f32>, kernel = array<i64: 4, 3>, pad = array<i64: 1, 2, 3, 4>, stride = array<i64: 1, 1>} : (tensor<3x5x6x7xf32>) -> tensor<?x?x?x?xf32>
 
   // CHECK: -> tensor<3x5x11x7xf32>
   %1 = "tosa.max_pool2d"(%arg0) {kernel = array<i64: 4, 3>, pad = array<i64: 1, 2, 3, 4>, stride = array<i64: 1, 1>} : (tensor<3x5x6x7xf32>) -> tensor<?x?x?x?xf32>
   return
 }
 
 // -----
 
func.func @conv2d_dynamic_bias(%input: tensor<2x8x9x3xf32>, %weights: tensor<5x3x6x3xf32>, %bias: tensor<?xf32>) -> () {
   return
 }
 
 // -----
 
 // CHECK-LABEL: @test_pool_stride
 func.func @test_pool_stride(%arg0: tensor<3x11x12x7xf32>) {
   // CHECK: -> tensor<3x4x4x7xf32>
-  %0 = "tosa.avg_pool2d"(%arg0) {kernel = array<i64: 4, 3>, pad = array<i64: 0, 0, 0, 0>, stride = array<i64: 2, 3>} : (tensor<3x11x12x7xf32>) -> tensor<?x?x?x?xf32>
+  %0 = "tosa.avg_pool2d"(%arg0) {acc_type = #tosa.accumulator_type<f32>, kernel = array<i64: 4, 3>, pad = array<i64: 0, 0, 0, 0>, stride = array<i64: 2, 3>} : (tensor<3x11x12x7xf32>) -> tensor<?x?x?x?xf32>
 
   // CHECK: -> tensor<3x4x4x7xf32>
   %1 = "tosa.max_pool2d"(%arg0) {kernel = array<i64: 4, 3>, pad = array<i64: 0, 0, 0, 0>, stride = array<i64: 2, 3>} : (tensor<3x11x12x7xf32>) -> tensor<?x?x?x?xf32>
   return
 }
 
 // -----
 

utils/bazel/llvm-project-overlay/mlir/BUILD.bazel

tbl_outs = [
             ["-gen-dialect-decls"],
             "include/mlir/Dialect/Tosa/IR/TosaOpsDialect.h.inc",
         ),
         (
             ["-gen-dialect-defs"],
             "include/mlir/Dialect/Tosa/IR/TosaOpsDialect.cpp.inc",
         ),
         (
+            ["-gen-enum-decls"],
+            "include/mlir/Dialect/Tosa/IR/TosaOpsEnums.h.inc",
+        ),
+        (
+            ["-gen-enum-defs"],
+            "include/mlir/Dialect/Tosa/IR/TosaOpsEnums.cpp.inc",
+        ),
+        (
             ["-gen-attrdef-decls"],
             "include/mlir/Dialect/Tosa/IR/TosaAttributes.h.inc",
         ),
         (
             ["-gen-attrdef-defs"],
             "include/mlir/Dialect/Tosa/IR/TosaAttributes.cpp.inc",
         ),
         (