additional cleanup

google · Mar 11, 2024 · 1995ae7 · 1995ae7
1 parent c6d0b87
commit 1995ae7
Show file tree

Hide file tree

Showing 4 changed files with 9 additions and 9 deletions.
diff --git a/include/Dialect/TensorExt/IR/TensorExtCanonicalization.td b/include/Dialect/TensorExt/IR/TensorExtCanonicalization.td
@@ -6,7 +6,7 @@ include "mlir/Dialect/Arith/IR/ArithOps.td"
 include "mlir/Dialect/Tensor/IR/TensorOps.td"
 include "mlir/IR/PatternBase.td"
 
-// FIXME: refactor these helpers to a common file with InsertRotate.td
+// TODO(#515): refactor these helpers to a common file with InsertRotate.td
 defvar DefOverflow = ConstantEnumCase<Arith_IntegerOverflowAttr, "none">;
 
 def MakeSingleResultVariadic: NativeCodeCall<"{ $0 }">;

diff --git a/include/Dialect/TensorExt/Transforms/InsertRotate.td b/include/Dialect/TensorExt/Transforms/InsertRotate.td
@@ -6,6 +6,8 @@ include "mlir/Dialect/Arith/IR/ArithOps.td"
 include "mlir/Dialect/Tensor/IR/TensorOps.td"
 include "mlir/IR/PatternBase.td"
 
+// TODO(#512): Support target slot selection when the downstream op is an insert.
+
 // The patterns in this file are intended to align with the automatic-SIMD
 // batching heuristics from the HECO project. See section 4.4 of
 // https://arxiv.org/abs/2202.01649 and the hir2hir passes in
@@ -40,6 +42,7 @@ foreach ArithOp = [Arith_AddIOp, Arith_SubIOp, Arith_MulIOp] in {
 
 // Pre-align the first op's operands to the index that the result is
 // used for in a subsequent op.
+// TODO(#514): handle OuterOp with two different InnerOps on the LHS and RHS
 foreach InnerOp = [Arith_AddIOp, Arith_SubIOp, Arith_MulIOp] in {
   foreach OuterOp = [Arith_AddIOp, Arith_SubIOp, Arith_MulIOp] in {
     // Left associated grouping handles (add (add (rotate t1 i1) (rotate t2 i2)) (rotate t3 i3))
@@ -75,8 +78,6 @@ foreach InnerOp = [Arith_AddIOp, Arith_SubIOp, Arith_MulIOp] in {
         (TensorExt_RotateOp $output, $i3),
       ]
     >;
-
-    // FIXME: do I need a pattern where OuterOp has two different InnerOp on the LHS and RHS?
   }
 }
 

diff --git a/include/Dialect/TensorExt/Transforms/Passes.td b/include/Dialect/TensorExt/Transforms/Passes.td
@@ -29,6 +29,7 @@ def InsertRotate : Pass<"insert-rotate"> {
   let dependentDialects = ["mlir::heir::tensor_ext::TensorExtDialect"];
 }
 
+// TODO(#512): Investigate replacing this pattern with a tensor_ext.combine op
 def CollapseInsertionChains : Pass<"collapse-insertion-chains"> {
   let summary = "Collapse chains of extract/insert ops into rotate ops when possible";
   let description = [{

diff --git a/tests/tensor_ext/insert_rotations.mlir b/tests/tensor_ext/insert_rotations.mlir
@@ -1,6 +1,6 @@
 // RUN: heir-opt --insert-rotate --canonicalize --cse %s | FileCheck %s
 
-func.func @test_insert_rotation_for_add(%arg1: tensor<16xi32>) -> tensor<16xi32> {
+func.func @test_insert_rotation_for_add(%arg1: tensor<16xi32>) -> i32 {
     %c4 = arith.constant 4 : index
     %c11 = arith.constant 11 : index
     %c15 = arith.constant 15 : index
@@ -13,14 +13,11 @@ func.func @test_insert_rotation_for_add(%arg1: tensor<16xi32>) -> tensor<16xi32>
 
     %extracted_1 = tensor.extract %arg1[%c4] : tensor<16xi32>
     %2 = arith.addi %1, %extracted_1 : i32
-
-    %inserted = tensor.insert %2 into %arg1[%c4] : tensor<16xi32>
-    // FIXME: try also just returning the extracted value directly
-    return %inserted : tensor<16xi32>
+    return %2 : i32
 }
 
 // CHECK-LABEL: func @test_insert_rotation_for_add
-// CHECK-SAME: (%[[arg0:.*]]: tensor<16xi32>) -> tensor<16xi32> {
+// CHECK-SAME: (%[[arg0:.*]]: tensor<16xi32>) -> i32 {
 // CHECK-NEXT:    %[[c11:.*]] = arith.constant 11
 // CHECK-NEXT:    %[[c4:.*]] = arith.constant 4
 // CHECK-NEXT:    %[[c7:.*]] = arith.constant 7
@@ -29,3 +26,4 @@ func.func @test_insert_rotation_for_add(%arg1: tensor<16xi32>) -> tensor<16xi32>
 // CHECK-NEXT:    %[[v2:.*]] = arith.addi %[[v0]], %[[v1]] : tensor<16xi32>
 // CHECK-NEXT:    %[[v3:.*]] = arith.addi %[[v2]], %[[arg0]] : tensor<16xi32>
 // CHECK-NEXT:    %[[extracted:.*]] = tensor.extract %[[v3]][%[[c4]]] : tensor<16xi32>
+// CHECK-NEXT:    return