[mlir][tosa] Add RFFT2d operation

[lhutton1]

Summary

Adds the RFFT2d TOSA operation and supporting
shape inference function.

Signed-off-by: Luke Hutton <[e-mail]>
Change-Id: I7e49c47cdd846cdc1b187545ef76d5cda2d5d9ad

[lhutton1]

Updates after a rebase:

• NoSideEffects -> Pure
• ShapedType::kDynamicSize -> ShapedType::kDynamic

[jpienaar]

Nice, thanks. Would you need help in landing?

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

391

Prefer std::optional (I think these are all being flipped)

399

So you'd have an input without rank but a result with? Is rankness verified?

405

And if this were symbolic, we'd have this form too?

E.g., symbolically this would be

(x,y,z) -> (x,y, z/2+1)

?

[lhutton1]

Thanks for the review jpienaar, yes I believe I'll need help landing this. I'm hoping to push a fix for the CI failure and address the comments shortly

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

399

Thanks yes this doesn't make sense, I'll change this to return failure when there is no rank

405

Yes that's correct. Apologies if I missed something (I'm very new to MLIR), my thinking was that symbolic shapes are not supported here?

[jpienaar]

Let me know when ready and I'll take a look (I didn't check failure)

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

405

It isn't here no - I'd like to change that "at some point" :) Feel free to just add a comment to this extent to ease when we come back around here at some point.

[lhutton1]

Addressing review comments regarding type inference.

[lhutton1]

This is ready for another review

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

405

Thanks, sounds exciting! Added a note

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

 //===-- TosaOps.h - TOSA dialect operation definitions ----------*- C++ -*-===//
 //
 // 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 declares the TOSA Dialect in MLIR.
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_DIALECT_TOSA_IR_TOSAOPS_H
 #define MLIR_DIALECT_TOSA_IR_TOSAOPS_H
 
 #include "mlir/Dialect/Traits.h"
+#include "mlir/IR/OpImplementation.h"
 #include "mlir/Interfaces/InferTypeOpInterface.h"
 #include "mlir/Interfaces/LoopLikeInterface.h"
 #include "mlir/Interfaces/SideEffectInterfaces.h"
 
 //===----------------------------------------------------------------------===//
 // TOSA dialect and structs includes.
 //===----------------------------------------------------------------------===//
 

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

def Tosa_MaxPool2dOp : Tosa_Op<"max_pool2d", [
   let results = (outs
     Tosa_Tensor4D:$output
   );
 
   let hasCanonicalizer = 1;
 }
 
 //===----------------------------------------------------------------------===//
+// Operator: rfft2d
+//===----------------------------------------------------------------------===//
+def Tosa_RFFT2dOp : Tosa_Op<"rfft2d", [
+    DeclareOpInterfaceMethods<InferShapedTypeOpInterface,
+                              ["inferReturnTypeComponents"]>,
+    Pure]> {
+  let summary = "Performs RFFT2D operation on the input.";
+
+  let description = [{
+    Performs a batched 2D real-valued Fast Fourier Transform over the input where
+    the input tensor consists of real values producing complex valued output. The
+    complex output values will be split into the output_real and output_imag
+    tensor arguments. RFFT2D takes advantage of Hermitian symmetry to only
+    calculate the first half of the final output axis. Imaginary values with
+    locations (0,0), (0,W/2), (H/2,0) and (H/2,W/2) are zero.
+  }];
+
+  let arguments = (ins
+    Tosa_Tensor3D:$input
+  );
+
+  let results = (outs
+    Tosa_Tensor3D:$output_real,
+    Tosa_Tensor3D:$output_imag
+  );
+}
+
+//===----------------------------------------------------------------------===//
 // Operator: transpose_conv2d
 //===----------------------------------------------------------------------===//
 def Tosa_TransposeConv2DOp : Tosa_Op<"transpose_conv2d", [
     DeclareOpInterfaceMethods<InferShapedTypeOpInterface,
                               ["inferReturnTypeComponents"]>,
     Pure]> {
   let summary = "Transpose 2D Convolution operator.";
 

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

if (i == axisVal)
       continue;
     outShape.push_back(inputShape.getDimSize(i));
   }
 
   inferredReturnShapes.push_back(ShapedTypeComponents(outShape));
   return success();
 }
 
+LogicalResult tosa::RFFT2dOp::inferReturnTypeComponents(
+    MLIRContext *context, ::std::optional<Location> location,
+    ValueShapeRange operands, DictionaryAttr attributes, RegionRange regions,
+    SmallVectorImpl<ShapedTypeComponents> &inferredReturnShapes) {
+  ShapeAdaptor inputShape = operands.getShape(0);
+
+  if (!inputShape.hasRank())
+    return failure();
+
+  llvm::SmallVector<int64_t> outputShape;
+  outputShape.resize(3, ShapedType::kDynamic);
+  outputShape[0] = inputShape.getDimSize(0);
+  outputShape[1] = inputShape.getDimSize(1);
+  int64_t inWidth = inputShape.getDimSize(2);
+
+  // Note that we can support this calculation symbolically
+  // in the future e.g. [x, y, z] -> [x, y, z / 2 - 1]
+  if (inWidth != ShapedType::kDynamic)
+    outputShape[2] = inWidth / 2 + 1;
+
+  inferredReturnShapes.push_back(ShapedTypeComponents(outputShape));
+  inferredReturnShapes.push_back(ShapedTypeComponents(outputShape));
+  return success();
+}
+
 LogicalResult tosa::ConcatOp::inferReturnTypeComponents(
     MLIRContext *context, ::std::optional<Location> location,
     ValueShapeRange operands, DictionaryAttr attributes, RegionRange regions,
     SmallVectorImpl<ShapedTypeComponents> &inferredReturnShapes) {
   // Infer all dimension sizes by reducing based on inputs.
   int32_t axis =
       attributes.get("axis").cast<IntegerAttr>().getValue().getSExtValue();
   llvm::SmallVector<int64_t> outputShape;

mlir/test/Dialect/Tosa/ops.mlir

 // -----
 // CHECK-LABEL: max_pool2d
 func.func @test_max_pool2d(%arg0: tensor<1x32x32x8xf32>) -> tensor<1x32x32x8xf32> {
   %0 = "tosa.max_pool2d"(%arg0) {kernel = array<i64: 1, 1>, pad = array<i64: 0, 0, 0, 0>, stride = array<i64: 1, 1>} : (tensor<1x32x32x8xf32>) -> tensor<1x32x32x8xf32>
   return %0 : tensor<1x32x32x8xf32>
 }
 
 // -----
+// CHECK-LABEL: rfft2d
+func.func @test_rfft2d(%arg0: tensor<13x8x16xf32>) -> (tensor<13x8x9xf32>, tensor<13x8x9xf32>) {
+  %0, %1 = "tosa.rfft2d"(%arg0) {} : (tensor<13x8x16xf32>) -> (tensor<13x8x9xf32>, tensor<13x8x9xf32>)
+  return %0, %1 : tensor<13x8x9xf32>, tensor<13x8x9xf32>
+}
+
+// -----
 // CHECK-LABEL: transpose_conv2d
 func.func @test_transpose_conv2d(%arg0: tensor<1x32x32x8xf32>, %arg1: tensor<16x1x1x8xf32>, %arg2: tensor<16xf32>) -> tensor<1x32x32x16xf32> {
   %0 = "tosa.transpose_conv2d"(%arg0, %arg1, %arg2) {out_pad = array<i64: 0, 0, 0, 0>, out_shape = array<i64: 1, 32, 32, 16>, stride = array<i64: 1, 1>} : (tensor<1x32x32x8xf32>, tensor<16x1x1x8xf32>, tensor<16xf32>) -> tensor<1x32x32x16xf32>
   return %0 : tensor<1x32x32x16xf32>
 }
 
 // -----
 // CHECK-LABEL: clamp

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

^bb0(%arg2: tensor<*xi32>, %arg3: tensor<*xi32>):
     // CHECK-SAME: tensor<i32>
     // CHECK-SAME: tensor<?xi32>
     "tosa.yield"(%3, %4) : (tensor<*xi32>, tensor<*xi32>) -> ()
 
   // CHECK:      (tensor<i32>, tensor<1xi32>) -> (tensor<i32>, tensor<?xi32>)
   }) : (tensor<i32>, tensor<1xi32>) -> (tensor<*xi32>, tensor<*xi32>)
   return
 }
+
+// -----
+
+// CHECK-LABEL: @test_static_rfft2d
+func.func @test_static_rfft2d(%arg0: tensor<5x2x8xf32>) -> () {
+  // CHECK: -> (tensor<5x2x5xf32>, tensor<5x2x5xf32>)
+  %output_real, %output_imag = "tosa.rfft2d"(%arg0) {} : (tensor<5x2x8xf32>) -> (tensor<?x?x?xf32>, tensor<?x?x?xf32>)
+  return
+}
+
+// -----
+
+// CHECK-LABEL: @test_dynamic_batch_rfft2d
+func.func @test_dynamic_batch_rfft2d(%arg0 : tensor<?x2x4xf32>) -> () {
+  // CHECK: -> (tensor<?x2x3xf32>, tensor<?x2x3xf32>)
+  %output_real, %output_imag = "tosa.rfft2d"(%arg0) {} : (tensor<?x2x4xf32>) -> (tensor<?x?x?xf32>, tensor<?x?x?xf32>)
+  return
+}
+
+// -----
+
+// CHECK-LABEL: @test_dynamic_width_rfft2d
+func.func @test_dynamic_width_rfft2d(%arg0 : tensor<5x2x?xf32>) -> () {
+  // CHECK: -> (tensor<5x2x?xf32>, tensor<5x2x?xf32>)
+  %output_real, %output_imag = "tosa.rfft2d"(%arg0) {} : (tensor<5x2x?xf32>) -> (tensor<?x?x?xf32>, tensor<?x?x?xf32>)
+  return
+}