Differences in handling of accumulate and result type for Average Pool and MatMul/Conv

TinaAMD · July 4, 2023, 1:33pm

According to the tosa spec, the input/output type for average pool is always the same, and the accumulator type can be different (especially also larger than the result type, e.g. fp16 with fp32 accumulate).

However, operators such as conv2/3d and matmul do not have an accumulate attribute, and their output type seems to be matching the accumulator type (implicitly).

Is there a reason why these operators have a different handling of the output types and the accumulate attribute?

I came across this as tosa.matmul is currently not properly lowered from torch to tosa, the result type will be the same as the input type:

github.com

llvm/torch-mlir/blob/157e5e529adae25dfbe69faa2fdcc18bcc67fe05/lib/Conversion/TorchToTosa/TorchToTosa.cpp#L1443

    
      
          
          
// The reshape/transpose should ensure the tosa.matmul always has same
          // batch size for either matrix. If if shapes are dynamic, they'll be
          // appropriately handled.
          assert(matmulLhsShape[0] == matmulRhsShape[0] &&
                 "tosa.matmul needs same batchsize on LHS and RHS");
          
          
SmallVector<int64_t> matmulOutputShape(
              {matmulLhsShape[0], matmulLhsShape[1], matmulRhsShape[2]});
          Type outputElemTy;
          if (lhsElemTy.isa<mlir::FloatType>()) {
            outputElemTy = lhsElemTy;
          } else { // qint8 emits i32 matmul output
            outputElemTy = rewriter.getIntegerType(32);
          }
          
          
auto mmOutputTy = RankedTensorType::get(
              makeShapeLLVMCompatible(matmulOutputShape), outputElemTy);
          auto mmOpResult =
              rewriter
                  .create<tosa::MatMulOp>(

p.s.: I wanted to add links to the tosa spec on mlplatform, but I wasn’t allowed to.

EricKunze · July 5, 2023, 3:48pm

Hi Tina,

In general, I believe that the separation of the accumulation and rescale as is done in the conv2/3d operators leads to simpler operation for the compiler. The reason that average pool is different and has the scale back to the original type is due to average pool dividing the sum by only the number of valid elements in the kernel region not including any padding. So on the corners of tensors the number of elements to divide by varies based on kernel size, padding, and strides. This would be hard to specify cleanly if the operators were separate. Once the division is done we expected that the best use of the result would be in the size of the inputs, as the average value is representable in the same range as the inputs.

I think after your first posting or two you should be able to post links. I believe it is an anti-spam mechanism. (For anyone looking, average pool is here: TOSA 0.70.0 specification, and conv2d is just below it)

Eric

TinaAMD · July 12, 2023, 9:07am

Thank you for the clarification!
We will fix the lowering from Torch → TOSA for matmul in torch-mlir.