[Pipeline] Add TorchToTcpCustomOpPass to frontend pipeline (#20)

This was missed from the earlier PR.
cruise-automation · Nov 6, 2023 · eaaa7cf · eaaa7cf
1 parent 32c2532
commit eaaa7cf
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 0 deletions.
diff --git a/lib/Pipeline/Pipeline.cpp b/lib/Pipeline/Pipeline.cpp
@@ -14,6 +14,7 @@
 #include "mlir-tcp/Conversion/TcpToArith/TcpToArith.h"
 #include "mlir-tcp/Conversion/TcpToLinalg/TcpToLinalg.h"
 #include "mlir-tcp/Conversion/TorchToTcp/TorchToTcp.h"
+#include "mlir-tcp/Conversion/TorchToTcp/TorchToTcpCustomOp.h"
 
 #include "mlir/Conversion/FuncToLLVM/ConvertFuncToLLVMPass.h"
 #include "mlir/Conversion/MathToLLVM/MathToLLVM.h"
@@ -36,6 +37,7 @@ using namespace mlir;
 
 static void createTorchBackendToTcpBackendPipeline(OpPassManager &pm) {
   pm.addNestedPass<func::FuncOp>(tcp::createConvertTorchToTcpPass());
+  pm.addNestedPass<func::FuncOp>(tcp::createConvertTorchToTcpCustomOpPass());
 
   // Clean up any non-canonical code introduced above.
   pm.addNestedPass<func::FuncOp>(createCanonicalizerPass());

diff --git a/test/Pipeline/torch_to_tcp_pipeline.mlir b/test/Pipeline/torch_to_tcp_pipeline.mlir
@@ -113,3 +113,17 @@ func.func @torch.aten.div.Tensor$mixed_type_int(%arg0: !torch.vtensor<[?, ?],si1
   %0 = torch.aten.div.Tensor %arg0, %arg1 : !torch.vtensor<[?, ?],si16>, !torch.vtensor<[?, ?],si32> -> !torch.vtensor<[?, ?],si32>
   return %0 : !torch.vtensor<[?, ?],si32>
 }
+
+// -----
+
+// CHECK-LABEL: torch.aten.gather_op
+// CHECK-SAME: %[[VAL_0:.*]]: tensor<2x2xi64>,
+// CHECK-SAME: %[[VAL_1:.*]]: tensor<2x2xf32>
+// CHECK: %[[VAL_2:.*]] = tcp.custom_op("torch.aten.gather") %[[VAL_1]], %[[VAL_0]] {axis = 1 : i64} : tensor<2x2xf32>, tensor<2x2xi64> -> tensor<2x2xf32>
+// CHECK: return %[[VAL_2]] : tensor<2x2xf32>
+func.func @torch.aten.gather_op(%arg0: !torch.vtensor<[2,2],si64>, %arg1: !torch.vtensor<[2,2],f32>) -> !torch.vtensor<[2,2],f32> {
+  %false = torch.constant.bool false
+  %int1 = torch.constant.int 1
+  %0 = torch.aten.gather %arg1, %int1, %arg0, %false : !torch.vtensor<[2,2],f32>, !torch.int, !torch.vtensor<[2,2],si64>, !torch.bool -> !torch.vtensor<[2,2],f32>
+  return %0 : !torch.vtensor<[2,2],f32>
+}