Address review comments

Author: victor-eds

lib/Dialect/Triton/IR/Ops.cpp

@@ -502,12 +502,14 @@ namespace {
 /// reduction dimension:
 /// ```mlir
 /// "tt.reduce"(%0, ...) <{axis = N}> ({...})
-/// : (tensor<S0x...xSN-1x1xSN+1x...>, ...) -> tensor<S0x...xSN-1xSN+1x...>
+///     : (tensor<S0 x ... x SN-1 x 1 x SN+1 x ...>, ...) ->
+///       (tensor<S0 x ... x SN-1 x SN+1 x ...>, ...)
 /// ```
 /// With equivalent reshape operations (one per operand):
 /// ```mlir
 /// tt.reshape %0 allow_reorder
-///     : tensor<S0x...xSN-1x1xSN+1x...> -> tensor<S0x...xSN-1xSN+1x...>
+///     : tensor<S0 x ... x SN-1 x 1 x SN+1 x ...> ->
+///       tensor<S0 x ... x SN-1 x SN+1 x ...>
 /// ```
 struct CanonicalizeReshapeReduceOpPattern final : OpRewritePattern<ReduceOp> {
   using OpRewritePattern<ReduceOp>::OpRewritePattern;

test/Triton/canonicalize.mlir

@@ -57,13 +57,14 @@ tt.func @fn(%arg0: tensor<1xf32, #sliced0>) -> (tensor<32x1xf32, #blocked0>){
 // CHECK-SAME:                  %[[ARG0:.*]]: tensor<2x1x16xf32>,
 // CHECK-SAME:                  %[[ARG1:.*]]: tensor<2x1x16xf16>
 tt.func @reduce(%arg0: tensor<2x1x16xf32>, %arg1: tensor<2x1x16xf16>) -> (tensor<2x16xf32>, tensor<2x16xf16>) {
-  // CHECK: tt.reshape %[[ARG0]] allow_reorder : tensor<2x1x16xf32> -> tensor<2x16xf32>
-  // CHECK: tt.reshape %[[ARG1]] allow_reorder : tensor<2x1x16xf16> -> tensor<2x16xf16>
+  // CHECK: %[[VAL0:.*]] = tt.reshape %[[ARG0]] allow_reorder : tensor<2x1x16xf32> -> tensor<2x16xf32>
+  // CHECK: %[[VAL1:.*]] = tt.reshape %[[ARG1]] allow_reorder : tensor<2x1x16xf16> -> tensor<2x16xf16>
   %0:2 = "tt.reduce"(%arg0, %arg1) <{axis=1 : i32}> ({
   ^bb0(%acc0: f32, %acc1: f16, %curr0: f32, %curr1: f16):
     %1 = arith.addf %acc0, %curr0 : f32
     %2 = arith.mulf %acc1, %curr1 : f16
     tt.reduce.return %1, %2 : f32, f16
   }) : (tensor<2x1x16xf32>, tensor<2x1x16xf16>) -> (tensor<2x16xf32>, tensor<2x16xf16>)
+  // CHECK: tt.return %[[VAL0]], %[[VAL1]] : tensor<2x16xf32>, tensor<2x16xf16>
   tt.return %0#0, %0#1 : tensor<2x16xf32>, tensor<2x16xf16>
 }