[mlir][tosa] Add accumulator verifier for FP8 types for AvgPool2D

[Jerry-Ge]

No description provided.

[llvmbot]

@llvm/pr-subscribers-mlir

@llvm/pr-subscribers-mlir-tosa

Author: Jerry-Ge (Jerry-Ge)

Changes

---

Full diff: https://github.com/llvm/llvm-project/pull/134727.diff

1 Files Affected:

• (modified) mlir/lib/Dialect/Tosa/IR/TosaOps.cpp (+5)

diff --git a/mlir/lib/Dialect/Tosa/IR/TosaOps.cpp b/mlir/lib/Dialect/Tosa/IR/TosaOps.cpp
index c8e9ad8bd3346..80bcc7a618080 100644
--- a/mlir/lib/Dialect/Tosa/IR/TosaOps.cpp
+++ b/mlir/lib/Dialect/Tosa/IR/TosaOps.cpp
@@ -600,6 +600,11 @@ LogicalResult tosa::AvgPool2dOp::verify() {
   if (inputETy.isF32() && !accType.isF32())
     return emitOpError("accumulator type for f32 tensor is not f32");
 
+  if ((llvm::isa<Float8E5M2Type>(inputETy) ||
+       llvm::isa<Float8E4M3FNType>(inputETy)) &&
+      !accType.isF16())
+    return emitOpError("accumulator type for f8 tensor is not f16");
+    
   if (inputETy != inputZpETy)
     return emitOpError("expect both input and its zero point are the same "
                        "element type, got ")

[github-actions]

✅ With the latest revision this PR passed the C/C++ code formatter.

[GeorgeARM]

Probably add a negative test as well?

[GeorgeARM]

Should we move the checks in the validation pass? Maybe we can leave the operator verifier a bit more relaxed ie just check that the accumulator is of fp type?
Thoughts?
Cc: @lhutton1 @FranklandJack

[lhutton1]

+1 from me on this - it allows other users of TOSA to express operations that may be important for their workflow e.g. an fp8 avg_pool with an fp32 accumulator

[Jerry-Ge]

sounds good, thanks for the review. i will then close this PR.