[Mlir-commits] [mlir] 0cfc5bc - [mlir][tosa][tosa-to-linalg] Use 0 weights at skippable RFFT2D locations (#212293)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Thu Jul 30 09:19:05 PDT 2026
Author: Ian Tayler Lessa
Date: 2026-07-30T17:18:59+01:00
New Revision: 0cfc5bc3aa10274776306e4a4a179974d3fadb4b
URL: https://github.com/llvm/llvm-project/commit/0cfc5bc3aa10274776306e4a4a179974d3fadb4b
DIFF: https://github.com/llvm/llvm-project/commit/0cfc5bc3aa10274776306e4a4a179974d3fadb4b.diff
LOG: [mlir][tosa][tosa-to-linalg] Use 0 weights at skippable RFFT2D locations (#212293)
The TOSA specifications indicates that when computing the imaginary
output for locations (0,0), (0,W/2), (H/2,0), (H/2, W/2),
implementations may choose to skip the computation entirely and return
0, or compute each accumulation term as `-val_real * 0.0`.
The previous legalizations used `-val_real * sin(a)` but `sin(a)` may
not return exactly 0.0 due to representation error in the input to
`sin`.
This patch updates the legalization to instead check for indices where
`sin(a)` should be replaced by `0.0` and uses a `select` op to match the
expected behaviour according to the Specification when
tosa_extra_multiplies is set to `true`.
Signed-off-by: Ian Tayler Lessa <ian.taylerlessa at arm.com>
Added:
Modified:
mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp
mlir/test/Conversion/TosaToLinalg/tosa-to-linalg.mlir
Removed:
################################################################################
diff --git a/mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp b/mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp
index e9c9e17fe6274..e3f40c57eb312 100644
--- a/mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp
+++ b/mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp
@@ -2827,10 +2827,18 @@ struct RFFT2dConverter final : public OpRewritePattern<RFFT2dOp> {
auto dimW = rewriter.createOrFold<tensor::DimOp>(loc, input, 2);
// Constants and dimension sizes
+ auto zeroFloat = arith::ConstantOp::create(
+ rewriter, loc, rewriter.getZeroAttr(elementType));
auto twoPiAttr = rewriter.getFloatAttr(elementType, 6.283185307179586);
auto twoPi = arith::ConstantOp::create(rewriter, loc, twoPiAttr);
+
+ auto zeroIndex = arith::ConstantIndexOp::create(rewriter, loc, 0);
+ auto twoIndex = arith::ConstantIndexOp::create(rewriter, loc, 2);
+
auto constH = castIndexToFloat(rewriter, loc, elementType, dimH);
auto constW = castIndexToFloat(rewriter, loc, elementType, dimW);
+ auto halfH = index::DivUOp::create(rewriter, loc, dimH, twoIndex);
+ auto halfW = index::DivUOp::create(rewriter, loc, dimW, twoIndex);
auto buildBody = [&](OpBuilder &builder, Location loc, ValueRange args) {
Value valReal = args[0];
@@ -2860,14 +2868,37 @@ struct RFFT2dConverter final : public OpRewritePattern<RFFT2dOp> {
auto sumXY = arith::AddFOp::create(builder, loc, yComponent, xComponent);
auto angle = arith::MulFOp::create(builder, loc, twoPi, sumXY);
+ // We will check the indices to see if this is a position that should use
+ // a 0.0 weight for the imaginary value computation following the TOSA
+ // specification with `tosa_extra_multiplies=true`.
+ //
+ // These are the relevant locations: (0,0), (0,W/2), (H/2,0), (H/2, W/2).
+ auto iyIs0 = arith::CmpIOp::create(builder, loc, arith::CmpIPredicate::eq,
+ iyRem, zeroIndex);
+ auto iyIsHalfH = arith::CmpIOp::create(
+ builder, loc, arith::CmpIPredicate::eq, iyRem, halfH);
+ auto ixIs0 = arith::CmpIOp::create(builder, loc, arith::CmpIPredicate::eq,
+ ixRem, zeroIndex);
+ auto ixIsHalfW = arith::CmpIOp::create(
+ builder, loc, arith::CmpIPredicate::eq, ixRem, halfW);
+
+ auto iyIsSinSkippable =
+ arith::OrIOp::create(builder, loc, iyIs0, iyIsHalfH);
+ auto ixIsSinSkippable =
+ arith::OrIOp::create(builder, loc, ixIs0, ixIsHalfW);
+ auto shouldSkipSin = arith::AndIOp::create(builder, loc, iyIsSinSkippable,
+ ixIsSinSkippable);
+
// realComponent = valReal * cos(angle)
- // imagComponent = valReal * sin(angle)
+ // imagComponent = valReal * (shouldSkipSin ? 0.0 : sin(angle))
auto cosAngle = math::CosOp::create(builder, loc, angle);
auto sinAngle = math::SinOp::create(builder, loc, angle);
+ auto imagWeight = arith::SelectOp::create(builder, loc, shouldSkipSin,
+ zeroFloat, sinAngle);
auto realComponent =
arith::MulFOp::create(builder, loc, valReal, cosAngle);
auto imagComponent =
- arith::MulFOp::create(builder, loc, valReal, sinAngle);
+ arith::MulFOp::create(builder, loc, valReal, imagWeight);
// outReal = sumReal + realComponent
// outImag = sumImag - imagComponent
diff --git a/mlir/test/Conversion/TosaToLinalg/tosa-to-linalg.mlir b/mlir/test/Conversion/TosaToLinalg/tosa-to-linalg.mlir
index e6bd800a0cf0a..a803ee7d99153 100644
--- a/mlir/test/Conversion/TosaToLinalg/tosa-to-linalg.mlir
+++ b/mlir/test/Conversion/TosaToLinalg/tosa-to-linalg.mlir
@@ -2048,58 +2048,70 @@ func.func @table8_dyn_table(%arg0: tensor<6xi8>, %arg1: tensor<?xi8>) -> () {
// -----
// NOTE: Assertions have been autogenerated by utils/generate-test-checks.py
-// CHECK: #[[$ATTR_0:.+]] = affine_map<(d0, d1, d2, d3, d4) -> (d0, d3, d4)>
-// CHECK: #[[$ATTR_1:.+]] = affine_map<(d0, d1, d2, d3, d4) -> (d0, d1, d2)>
-
+// CHECK: #[[$ATTR_85:.+]] = affine_map<(d0, d1, d2, d3, d4) -> (d0, d3, d4)>
+// CHECK: #[[$ATTR_86:.+]] = affine_map<(d0, d1, d2, d3, d4) -> (d0, d1, d2)>
// CHECK-LABEL: func.func @test_static_rfft2d(
-// CHECK-SAME: %[[VAL_0:.*]]: tensor<5x4x8xf32>) -> (tensor<5x4x5xf32>, tensor<5x4x5xf32>) {
-// CHECK: %[[VAL_1:.*]] = arith.constant 1 : index
-// CHECK: %[[VAL_2:.*]] = arith.constant 2 : index
-// CHECK: %[[VAL_3:.*]] = arith.constant 8 : index
-// CHECK: %[[VAL_4:.*]] = arith.constant 4 : index
-// CHECK: %[[VAL_5:.*]] = arith.constant 5 : index
-// CHECK: %[[VAL_6:.*]] = tensor.empty() : tensor<5x4x5xf32>
-// CHECK: %[[VAL_7:.*]] = arith.constant 0.000000e+00 : f32
-// CHECK: %[[VAL_8:.*]] = linalg.fill ins(%[[VAL_7]] : f32) outs(%[[VAL_6]] : tensor<5x4x5xf32>) -> tensor<5x4x5xf32>
-// CHECK: %[[VAL_9:.*]] = tensor.empty() : tensor<5x4x5xf32>
-// CHECK: %[[VAL_10:.*]] = arith.constant 0.000000e+00 : f32
-// CHECK: %[[VAL_11:.*]] = linalg.fill ins(%[[VAL_10]] : f32) outs(%[[VAL_9]] : tensor<5x4x5xf32>) -> tensor<5x4x5xf32>
-// CHECK: %[[VAL_12:.*]] = arith.constant 1 : index
-// CHECK: %[[VAL_13:.*]] = arith.constant 4 : index
-// CHECK: %[[VAL_14:.*]] = arith.constant 2 : index
-// CHECK: %[[VAL_15:.*]] = arith.constant 8 : index
-// CHECK: %[[VAL_16:.*]] = arith.constant 6.28318548 : f32
-// CHECK: %[[VAL_17:.*]] = arith.index_castui %[[VAL_13]] : index to i32
-// CHECK: %[[VAL_18:.*]] = arith.uitofp %[[VAL_17]] : i32 to f32
-// CHECK: %[[VAL_19:.*]] = arith.index_castui %[[VAL_15]] : index to i32
-// CHECK: %[[VAL_20:.*]] = arith.uitofp %[[VAL_19]] : i32 to f32
-// CHECK: %[[VAL_21:.*]]:2 = linalg.generic {indexing_maps = [#[[$ATTR_0]], #[[$ATTR_1]], #[[$ATTR_1]]], iterator_types = ["parallel", "parallel", "parallel", "reduction", "reduction"]} ins(%[[VAL_0]] : tensor<5x4x8xf32>) outs(%[[VAL_8]], %[[VAL_11]] : tensor<5x4x5xf32>, tensor<5x4x5xf32>) {
-// CHECK: ^bb0(%[[VAL_22:.*]]: f32, %[[VAL_23:.*]]: f32, %[[VAL_24:.*]]: f32):
-// CHECK: %[[VAL_25:.*]] = linalg.index 1 : index
-// CHECK: %[[VAL_26:.*]] = linalg.index 2 : index
-// CHECK: %[[VAL_27:.*]] = linalg.index 3 : index
-// CHECK: %[[VAL_28:.*]] = linalg.index 4 : index
-// CHECK: %[[VAL_29:.*]] = index.mul %[[VAL_27]], %[[VAL_25]]
-// CHECK: %[[VAL_30:.*]] = index.mul %[[VAL_28]], %[[VAL_26]]
-// CHECK: %[[VAL_31:.*]] = index.remu %[[VAL_29]], %[[VAL_13]]
-// CHECK: %[[VAL_32:.*]] = index.remu %[[VAL_30]], %[[VAL_15]]
-// CHECK: %[[VAL_33:.*]] = arith.index_castui %[[VAL_31]] : index to i32
-// CHECK: %[[VAL_34:.*]] = arith.uitofp %[[VAL_33]] : i32 to f32
-// CHECK: %[[VAL_35:.*]] = arith.index_castui %[[VAL_32]] : index to i32
-// CHECK: %[[VAL_36:.*]] = arith.uitofp %[[VAL_35]] : i32 to f32
-// CHECK: %[[VAL_37:.*]] = arith.divf %[[VAL_34]], %[[VAL_18]] : f32
-// CHECK: %[[VAL_38:.*]] = arith.divf %[[VAL_36]], %[[VAL_20]] : f32
-// CHECK: %[[VAL_39:.*]] = arith.addf %[[VAL_37]], %[[VAL_38]] : f32
-// CHECK: %[[VAL_40:.*]] = arith.mulf %[[VAL_16]], %[[VAL_39]] : f32
-// CHECK: %[[VAL_41:.*]] = math.cos %[[VAL_40]] : f32
-// CHECK: %[[VAL_42:.*]] = math.sin %[[VAL_40]] : f32
-// CHECK: %[[VAL_43:.*]] = arith.mulf %[[VAL_22]], %[[VAL_41]] : f32
-// CHECK: %[[VAL_44:.*]] = arith.mulf %[[VAL_22]], %[[VAL_42]] : f32
-// CHECK: %[[VAL_45:.*]] = arith.addf %[[VAL_23]], %[[VAL_43]] : f32
-// CHECK: %[[VAL_46:.*]] = arith.subf %[[VAL_24]], %[[VAL_44]] : f32
-// CHECK: linalg.yield %[[VAL_45]], %[[VAL_46]] : f32, f32
+// CHECK-SAME: %[[ARG0:.*]]: tensor<5x4x8xf32>) -> (tensor<5x4x5xf32>, tensor<5x4x5xf32>) {
+// CHECK: %[[CONSTANT_0:.*]] = arith.constant 1 : index
+// CHECK: %[[CONSTANT_1:.*]] = arith.constant 2 : index
+// CHECK: %[[CONSTANT_2:.*]] = arith.constant 8 : index
+// CHECK: %[[CONSTANT_3:.*]] = arith.constant 4 : index
+// CHECK: %[[CONSTANT_4:.*]] = arith.constant 5 : index
+// CHECK: %[[EMPTY_0:.*]] = tensor.empty() : tensor<5x4x5xf32>
+// CHECK: %[[CONSTANT_5:.*]] = arith.constant 0.000000e+00 : f32
+// CHECK: %[[FILL_0:.*]] = linalg.fill ins(%[[CONSTANT_5]] : f32) outs(%[[EMPTY_0]] : tensor<5x4x5xf32>) -> tensor<5x4x5xf32>
+// CHECK: %[[EMPTY_1:.*]] = tensor.empty() : tensor<5x4x5xf32>
+// CHECK: %[[CONSTANT_6:.*]] = arith.constant 0.000000e+00 : f32
+// CHECK: %[[FILL_1:.*]] = linalg.fill ins(%[[CONSTANT_6]] : f32) outs(%[[EMPTY_1]] : tensor<5x4x5xf32>) -> tensor<5x4x5xf32>
+// CHECK: %[[CONSTANT_7:.*]] = arith.constant 1 : index
+// CHECK: %[[CONSTANT_8:.*]] = arith.constant 4 : index
+// CHECK: %[[CONSTANT_9:.*]] = arith.constant 2 : index
+// CHECK: %[[CONSTANT_10:.*]] = arith.constant 8 : index
+// CHECK: %[[CONSTANT_11:.*]] = arith.constant 0.000000e+00 : f32
+// CHECK: %[[CONSTANT_12:.*]] = arith.constant 6.28318548 : f32
+// CHECK: %[[CONSTANT_13:.*]] = arith.constant 0 : index
+// CHECK: %[[CONSTANT_14:.*]] = arith.constant 2 : index
+// CHECK: %[[INDEX_CASTUI_0:.*]] = arith.index_castui %[[CONSTANT_8]] : index to i32
+// CHECK: %[[UITOFP_0:.*]] = arith.uitofp %[[INDEX_CASTUI_0]] : i32 to f32
+// CHECK: %[[INDEX_CASTUI_1:.*]] = arith.index_castui %[[CONSTANT_10]] : index to i32
+// CHECK: %[[UITOFP_1:.*]] = arith.uitofp %[[INDEX_CASTUI_1]] : i32 to f32
+// CHECK: %[[DIVU_0:.*]] = index.divu %[[CONSTANT_8]], %[[CONSTANT_14]]
+// CHECK: %[[DIVU_1:.*]] = index.divu %[[CONSTANT_10]], %[[CONSTANT_14]]
+// CHECK: %[[GENERIC_0:.*]]:2 = linalg.generic {indexing_maps = [#[[$ATTR_85]], #[[$ATTR_86]], #[[$ATTR_86]]], iterator_types = ["parallel", "parallel", "parallel", "reduction", "reduction"]} ins(%[[ARG0]] : tensor<5x4x8xf32>) outs(%[[FILL_0]], %[[FILL_1]] : tensor<5x4x5xf32>, tensor<5x4x5xf32>) {
+// CHECK: ^bb0(%[[VAL_0:.*]]: f32, %[[VAL_1:.*]]: f32, %[[VAL_2:.*]]: f32):
+// CHECK: %[[INDEX_0:.*]] = linalg.index 1 : index
+// CHECK: %[[INDEX_1:.*]] = linalg.index 2 : index
+// CHECK: %[[INDEX_2:.*]] = linalg.index 3 : index
+// CHECK: %[[INDEX_3:.*]] = linalg.index 4 : index
+// CHECK: %[[MUL_0:.*]] = index.mul %[[INDEX_2]], %[[INDEX_0]]
+// CHECK: %[[MUL_1:.*]] = index.mul %[[INDEX_3]], %[[INDEX_1]]
+// CHECK: %[[REMU_0:.*]] = index.remu %[[MUL_0]], %[[CONSTANT_8]]
+// CHECK: %[[REMU_1:.*]] = index.remu %[[MUL_1]], %[[CONSTANT_10]]
+// CHECK: %[[INDEX_CASTUI_2:.*]] = arith.index_castui %[[REMU_0]] : index to i32
+// CHECK: %[[UITOFP_2:.*]] = arith.uitofp %[[INDEX_CASTUI_2]] : i32 to f32
+// CHECK: %[[INDEX_CASTUI_3:.*]] = arith.index_castui %[[REMU_1]] : index to i32
+// CHECK: %[[UITOFP_3:.*]] = arith.uitofp %[[INDEX_CASTUI_3]] : i32 to f32
+// CHECK: %[[DIVF_0:.*]] = arith.divf %[[UITOFP_2]], %[[UITOFP_0]] : f32
+// CHECK: %[[DIVF_1:.*]] = arith.divf %[[UITOFP_3]], %[[UITOFP_1]] : f32
+// CHECK: %[[ADDF_0:.*]] = arith.addf %[[DIVF_0]], %[[DIVF_1]] : f32
+// CHECK: %[[MULF_0:.*]] = arith.mulf %[[CONSTANT_12]], %[[ADDF_0]] : f32
+// CHECK: %[[CMPI_0:.*]] = arith.cmpi eq, %[[REMU_0]], %[[CONSTANT_13]] : index
+// CHECK: %[[CMPI_1:.*]] = arith.cmpi eq, %[[REMU_0]], %[[DIVU_0]] : index
+// CHECK: %[[CMPI_2:.*]] = arith.cmpi eq, %[[REMU_1]], %[[CONSTANT_13]] : index
+// CHECK: %[[CMPI_3:.*]] = arith.cmpi eq, %[[REMU_1]], %[[DIVU_1]] : index
+// CHECK: %[[ORI_0:.*]] = arith.ori %[[CMPI_0]], %[[CMPI_1]] : i1
+// CHECK: %[[ORI_1:.*]] = arith.ori %[[CMPI_2]], %[[CMPI_3]] : i1
+// CHECK: %[[ANDI_0:.*]] = arith.andi %[[ORI_0]], %[[ORI_1]] : i1
+// CHECK: %[[COS_0:.*]] = math.cos %[[MULF_0]] : f32
+// CHECK: %[[SIN_0:.*]] = math.sin %[[MULF_0]] : f32
+// CHECK: %[[SELECT_0:.*]] = arith.select %[[ANDI_0]], %[[CONSTANT_11]], %[[SIN_0]] : f32
+// CHECK: %[[MULF_1:.*]] = arith.mulf %[[VAL_0]], %[[COS_0]] : f32
+// CHECK: %[[MULF_2:.*]] = arith.mulf %[[VAL_0]], %[[SELECT_0]] : f32
+// CHECK: %[[ADDF_1:.*]] = arith.addf %[[VAL_1]], %[[MULF_1]] : f32
+// CHECK: %[[SUBF_0:.*]] = arith.subf %[[VAL_2]], %[[MULF_2]] : f32
+// CHECK: linalg.yield %[[ADDF_1]], %[[SUBF_0]] : f32, f32
// CHECK: } -> (tensor<5x4x5xf32>, tensor<5x4x5xf32>)
-// CHECK: return %[[VAL_47:.*]]#0, %[[VAL_47]]#1 : tensor<5x4x5xf32>, tensor<5x4x5xf32>
+// CHECK: return %[[VAL_3:.*]]#0, %[[VAL_3]]#1 : tensor<5x4x5xf32>, tensor<5x4x5xf32>
// CHECK: }
func.func @test_static_rfft2d(%arg0: tensor<5x4x8xf32>) -> (tensor<5x4x5xf32>, tensor<5x4x5xf32>) {
%output_real, %output_imag = "tosa.rfft2d"(%arg0) {} : (tensor<5x4x8xf32>) -> (tensor<5x4x5xf32>, tensor<5x4x5xf32>)
@@ -2108,63 +2120,75 @@ func.func @test_static_rfft2d(%arg0: tensor<5x4x8xf32>) -> (tensor<5x4x5xf32>, t
// -----
// NOTE: Assertions have been autogenerated by utils/generate-test-checks.py
-// CHECK: #[[$ATTR_0:.+]] = affine_map<(d0, d1, d2, d3, d4) -> (d0, d3, d4)>
-// CHECK: #[[$ATTR_1:.+]] = affine_map<(d0, d1, d2, d3, d4) -> (d0, d1, d2)>
-
+// CHECK: #[[$ATTR_87:.+]] = affine_map<(d0, d1, d2, d3, d4) -> (d0, d3, d4)>
+// CHECK: #[[$ATTR_88:.+]] = affine_map<(d0, d1, d2, d3, d4) -> (d0, d1, d2)>
// CHECK-LABEL: func.func @test_dynamic_rfft2d(
-// CHECK-SAME: %[[VAL_0:.*]]: tensor<?x?x?xf32>) -> (tensor<?x?x?xf32>, tensor<?x?x?xf32>) {
-// CHECK: %[[VAL_1:.*]] = arith.constant 0 : index
-// CHECK: %[[VAL_2:.*]] = tensor.dim %[[VAL_0]], %[[VAL_1]] : tensor<?x?x?xf32>
-// CHECK: %[[VAL_3:.*]] = arith.constant 1 : index
-// CHECK: %[[VAL_4:.*]] = tensor.dim %[[VAL_0]], %[[VAL_3]] : tensor<?x?x?xf32>
-// CHECK: %[[VAL_5:.*]] = arith.constant 2 : index
-// CHECK: %[[VAL_6:.*]] = tensor.dim %[[VAL_0]], %[[VAL_5]] : tensor<?x?x?xf32>
-// CHECK: %[[VAL_7:.*]] = arith.constant 1 : index
-// CHECK: %[[VAL_8:.*]] = arith.constant 2 : index
-// CHECK: %[[VAL_9:.*]] = arith.divui %[[VAL_6]], %[[VAL_8]] : index
-// CHECK: %[[VAL_10:.*]] = arith.addi %[[VAL_9]], %[[VAL_7]] : index
-// CHECK: %[[VAL_11:.*]] = tensor.empty(%[[VAL_2]], %[[VAL_4]], %[[VAL_10]]) : tensor<?x?x?xf32>
-// CHECK: %[[VAL_12:.*]] = arith.constant 0.000000e+00 : f32
-// CHECK: %[[VAL_13:.*]] = linalg.fill ins(%[[VAL_12]] : f32) outs(%[[VAL_11]] : tensor<?x?x?xf32>) -> tensor<?x?x?xf32>
-// CHECK: %[[VAL_14:.*]] = tensor.empty(%[[VAL_2]], %[[VAL_4]], %[[VAL_10]]) : tensor<?x?x?xf32>
-// CHECK: %[[VAL_15:.*]] = arith.constant 0.000000e+00 : f32
-// CHECK: %[[VAL_16:.*]] = linalg.fill ins(%[[VAL_15]] : f32) outs(%[[VAL_14]] : tensor<?x?x?xf32>) -> tensor<?x?x?xf32>
-// CHECK: %[[VAL_17:.*]] = arith.constant 1 : index
-// CHECK: %[[VAL_18:.*]] = tensor.dim %[[VAL_0]], %[[VAL_17]] : tensor<?x?x?xf32>
-// CHECK: %[[VAL_19:.*]] = arith.constant 2 : index
-// CHECK: %[[VAL_20:.*]] = tensor.dim %[[VAL_0]], %[[VAL_19]] : tensor<?x?x?xf32>
-// CHECK: %[[VAL_21:.*]] = arith.constant 6.28318548 : f32
-// CHECK: %[[VAL_22:.*]] = arith.index_castui %[[VAL_18]] : index to i32
-// CHECK: %[[VAL_23:.*]] = arith.uitofp %[[VAL_22]] : i32 to f32
-// CHECK: %[[VAL_24:.*]] = arith.index_castui %[[VAL_20]] : index to i32
-// CHECK: %[[VAL_25:.*]] = arith.uitofp %[[VAL_24]] : i32 to f32
-// CHECK: %[[VAL_26:.*]]:2 = linalg.generic {indexing_maps = [#[[$ATTR_0]], #[[$ATTR_1]], #[[$ATTR_1]]], iterator_types = ["parallel", "parallel", "parallel", "reduction", "reduction"]} ins(%[[VAL_0]] : tensor<?x?x?xf32>) outs(%[[VAL_13]], %[[VAL_16]] : tensor<?x?x?xf32>, tensor<?x?x?xf32>) {
-// CHECK: ^bb0(%[[VAL_27:.*]]: f32, %[[VAL_28:.*]]: f32, %[[VAL_29:.*]]: f32):
-// CHECK: %[[VAL_30:.*]] = linalg.index 1 : index
-// CHECK: %[[VAL_31:.*]] = linalg.index 2 : index
-// CHECK: %[[VAL_32:.*]] = linalg.index 3 : index
-// CHECK: %[[VAL_33:.*]] = linalg.index 4 : index
-// CHECK: %[[VAL_34:.*]] = index.mul %[[VAL_32]], %[[VAL_30]]
-// CHECK: %[[VAL_35:.*]] = index.mul %[[VAL_33]], %[[VAL_31]]
-// CHECK: %[[VAL_36:.*]] = index.remu %[[VAL_34]], %[[VAL_18]]
-// CHECK: %[[VAL_37:.*]] = index.remu %[[VAL_35]], %[[VAL_20]]
-// CHECK: %[[VAL_38:.*]] = arith.index_castui %[[VAL_36]] : index to i32
-// CHECK: %[[VAL_39:.*]] = arith.uitofp %[[VAL_38]] : i32 to f32
-// CHECK: %[[VAL_40:.*]] = arith.index_castui %[[VAL_37]] : index to i32
-// CHECK: %[[VAL_41:.*]] = arith.uitofp %[[VAL_40]] : i32 to f32
-// CHECK: %[[VAL_42:.*]] = arith.divf %[[VAL_39]], %[[VAL_23]] : f32
-// CHECK: %[[VAL_43:.*]] = arith.divf %[[VAL_41]], %[[VAL_25]] : f32
-// CHECK: %[[VAL_44:.*]] = arith.addf %[[VAL_42]], %[[VAL_43]] : f32
-// CHECK: %[[VAL_45:.*]] = arith.mulf %[[VAL_21]], %[[VAL_44]] : f32
-// CHECK: %[[VAL_46:.*]] = math.cos %[[VAL_45]] : f32
-// CHECK: %[[VAL_47:.*]] = math.sin %[[VAL_45]] : f32
-// CHECK: %[[VAL_48:.*]] = arith.mulf %[[VAL_27]], %[[VAL_46]] : f32
-// CHECK: %[[VAL_49:.*]] = arith.mulf %[[VAL_27]], %[[VAL_47]] : f32
-// CHECK: %[[VAL_50:.*]] = arith.addf %[[VAL_28]], %[[VAL_48]] : f32
-// CHECK: %[[VAL_51:.*]] = arith.subf %[[VAL_29]], %[[VAL_49]] : f32
-// CHECK: linalg.yield %[[VAL_50]], %[[VAL_51]] : f32, f32
+// CHECK-SAME: %[[ARG0:.*]]: tensor<?x?x?xf32>) -> (tensor<?x?x?xf32>, tensor<?x?x?xf32>) {
+// CHECK: %[[CONSTANT_0:.*]] = arith.constant 0 : index
+// CHECK: %[[DIM_0:.*]] = tensor.dim %[[ARG0]], %[[CONSTANT_0]] : tensor<?x?x?xf32>
+// CHECK: %[[CONSTANT_1:.*]] = arith.constant 1 : index
+// CHECK: %[[DIM_1:.*]] = tensor.dim %[[ARG0]], %[[CONSTANT_1]] : tensor<?x?x?xf32>
+// CHECK: %[[CONSTANT_2:.*]] = arith.constant 2 : index
+// CHECK: %[[DIM_2:.*]] = tensor.dim %[[ARG0]], %[[CONSTANT_2]] : tensor<?x?x?xf32>
+// CHECK: %[[CONSTANT_3:.*]] = arith.constant 1 : index
+// CHECK: %[[CONSTANT_4:.*]] = arith.constant 2 : index
+// CHECK: %[[DIVUI_0:.*]] = arith.divui %[[DIM_2]], %[[CONSTANT_4]] : index
+// CHECK: %[[ADDI_0:.*]] = arith.addi %[[DIVUI_0]], %[[CONSTANT_3]] : index
+// CHECK: %[[EMPTY_0:.*]] = tensor.empty(%[[DIM_0]], %[[DIM_1]], %[[ADDI_0]]) : tensor<?x?x?xf32>
+// CHECK: %[[CONSTANT_5:.*]] = arith.constant 0.000000e+00 : f32
+// CHECK: %[[FILL_0:.*]] = linalg.fill ins(%[[CONSTANT_5]] : f32) outs(%[[EMPTY_0]] : tensor<?x?x?xf32>) -> tensor<?x?x?xf32>
+// CHECK: %[[EMPTY_1:.*]] = tensor.empty(%[[DIM_0]], %[[DIM_1]], %[[ADDI_0]]) : tensor<?x?x?xf32>
+// CHECK: %[[CONSTANT_6:.*]] = arith.constant 0.000000e+00 : f32
+// CHECK: %[[FILL_1:.*]] = linalg.fill ins(%[[CONSTANT_6]] : f32) outs(%[[EMPTY_1]] : tensor<?x?x?xf32>) -> tensor<?x?x?xf32>
+// CHECK: %[[CONSTANT_7:.*]] = arith.constant 1 : index
+// CHECK: %[[DIM_3:.*]] = tensor.dim %[[ARG0]], %[[CONSTANT_7]] : tensor<?x?x?xf32>
+// CHECK: %[[CONSTANT_8:.*]] = arith.constant 2 : index
+// CHECK: %[[DIM_4:.*]] = tensor.dim %[[ARG0]], %[[CONSTANT_8]] : tensor<?x?x?xf32>
+// CHECK: %[[CONSTANT_9:.*]] = arith.constant 0.000000e+00 : f32
+// CHECK: %[[CONSTANT_10:.*]] = arith.constant 6.28318548 : f32
+// CHECK: %[[CONSTANT_11:.*]] = arith.constant 0 : index
+// CHECK: %[[CONSTANT_12:.*]] = arith.constant 2 : index
+// CHECK: %[[INDEX_CASTUI_0:.*]] = arith.index_castui %[[DIM_3]] : index to i32
+// CHECK: %[[UITOFP_0:.*]] = arith.uitofp %[[INDEX_CASTUI_0]] : i32 to f32
+// CHECK: %[[INDEX_CASTUI_1:.*]] = arith.index_castui %[[DIM_4]] : index to i32
+// CHECK: %[[UITOFP_1:.*]] = arith.uitofp %[[INDEX_CASTUI_1]] : i32 to f32
+// CHECK: %[[DIVU_0:.*]] = index.divu %[[DIM_3]], %[[CONSTANT_12]]
+// CHECK: %[[DIVU_1:.*]] = index.divu %[[DIM_4]], %[[CONSTANT_12]]
+// CHECK: %[[GENERIC_0:.*]]:2 = linalg.generic {indexing_maps = [#[[$ATTR_87]], #[[$ATTR_88]], #[[$ATTR_88]]], iterator_types = ["parallel", "parallel", "parallel", "reduction", "reduction"]} ins(%[[ARG0]] : tensor<?x?x?xf32>) outs(%[[FILL_0]], %[[FILL_1]] : tensor<?x?x?xf32>, tensor<?x?x?xf32>) {
+// CHECK: ^bb0(%[[VAL_0:.*]]: f32, %[[VAL_1:.*]]: f32, %[[VAL_2:.*]]: f32):
+// CHECK: %[[INDEX_0:.*]] = linalg.index 1 : index
+// CHECK: %[[INDEX_1:.*]] = linalg.index 2 : index
+// CHECK: %[[INDEX_2:.*]] = linalg.index 3 : index
+// CHECK: %[[INDEX_3:.*]] = linalg.index 4 : index
+// CHECK: %[[MUL_0:.*]] = index.mul %[[INDEX_2]], %[[INDEX_0]]
+// CHECK: %[[MUL_1:.*]] = index.mul %[[INDEX_3]], %[[INDEX_1]]
+// CHECK: %[[REMU_0:.*]] = index.remu %[[MUL_0]], %[[DIM_3]]
+// CHECK: %[[REMU_1:.*]] = index.remu %[[MUL_1]], %[[DIM_4]]
+// CHECK: %[[INDEX_CASTUI_2:.*]] = arith.index_castui %[[REMU_0]] : index to i32
+// CHECK: %[[UITOFP_2:.*]] = arith.uitofp %[[INDEX_CASTUI_2]] : i32 to f32
+// CHECK: %[[INDEX_CASTUI_3:.*]] = arith.index_castui %[[REMU_1]] : index to i32
+// CHECK: %[[UITOFP_3:.*]] = arith.uitofp %[[INDEX_CASTUI_3]] : i32 to f32
+// CHECK: %[[DIVF_0:.*]] = arith.divf %[[UITOFP_2]], %[[UITOFP_0]] : f32
+// CHECK: %[[DIVF_1:.*]] = arith.divf %[[UITOFP_3]], %[[UITOFP_1]] : f32
+// CHECK: %[[ADDF_0:.*]] = arith.addf %[[DIVF_0]], %[[DIVF_1]] : f32
+// CHECK: %[[MULF_0:.*]] = arith.mulf %[[CONSTANT_10]], %[[ADDF_0]] : f32
+// CHECK: %[[CMPI_0:.*]] = arith.cmpi eq, %[[REMU_0]], %[[CONSTANT_11]] : index
+// CHECK: %[[CMPI_1:.*]] = arith.cmpi eq, %[[REMU_0]], %[[DIVU_0]] : index
+// CHECK: %[[CMPI_2:.*]] = arith.cmpi eq, %[[REMU_1]], %[[CONSTANT_11]] : index
+// CHECK: %[[CMPI_3:.*]] = arith.cmpi eq, %[[REMU_1]], %[[DIVU_1]] : index
+// CHECK: %[[ORI_0:.*]] = arith.ori %[[CMPI_0]], %[[CMPI_1]] : i1
+// CHECK: %[[ORI_1:.*]] = arith.ori %[[CMPI_2]], %[[CMPI_3]] : i1
+// CHECK: %[[ANDI_0:.*]] = arith.andi %[[ORI_0]], %[[ORI_1]] : i1
+// CHECK: %[[COS_0:.*]] = math.cos %[[MULF_0]] : f32
+// CHECK: %[[SIN_0:.*]] = math.sin %[[MULF_0]] : f32
+// CHECK: %[[SELECT_0:.*]] = arith.select %[[ANDI_0]], %[[CONSTANT_9]], %[[SIN_0]] : f32
+// CHECK: %[[MULF_1:.*]] = arith.mulf %[[VAL_0]], %[[COS_0]] : f32
+// CHECK: %[[MULF_2:.*]] = arith.mulf %[[VAL_0]], %[[SELECT_0]] : f32
+// CHECK: %[[ADDF_1:.*]] = arith.addf %[[VAL_1]], %[[MULF_1]] : f32
+// CHECK: %[[SUBF_0:.*]] = arith.subf %[[VAL_2]], %[[MULF_2]] : f32
+// CHECK: linalg.yield %[[ADDF_1]], %[[SUBF_0]] : f32, f32
// CHECK: } -> (tensor<?x?x?xf32>, tensor<?x?x?xf32>)
-// CHECK: return %[[VAL_52:.*]]#0, %[[VAL_52]]#1 : tensor<?x?x?xf32>, tensor<?x?x?xf32>
+// CHECK: return %[[VAL_3:.*]]#0, %[[VAL_3]]#1 : tensor<?x?x?xf32>, tensor<?x?x?xf32>
// CHECK: }
func.func @test_dynamic_rfft2d(%arg0: tensor<?x?x?xf32>) -> (tensor<?x?x?xf32>, tensor<?x?x?xf32>) {
%output_real, %output_imag = "tosa.rfft2d"(%arg0) {} : (tensor<?x?x?xf32>) -> (tensor<?x?x?xf32>, tensor<?x?x?xf32>)
More information about the Mlir-commits
mailing list