Add flow.tensor.bitcast for torch view as complex/real

compiler/plugins/input/Torch/InputConversion/BitCastQuantTensor.cpp → compiler/plugins/input/Torch/InputConversion/BitCastTensor.cpp

@@ -16,7 +16,7 @@
 
 namespace mlir::iree_compiler::TorchInput {
 
-#define GEN_PASS_DEF_BITCASTQUANTTENSORPASS
+#define GEN_PASS_DEF_BITCASTTENSORPASS
 #include "compiler/plugins/input/Torch/InputConversion/Passes.h.inc"
 
 namespace {
@@ -62,8 +62,44 @@ class BitCastViewDtype
   }
 };
 
-class BitCastQuantizedMatmul
-    : public OpRewritePattern<torch::Torch::OperatorOp> {
+template <typename BitcastOp>
+class BitCastViewComplex : public OpRewritePattern<BitcastOp> {
+public:
+  using OpRewritePattern<BitcastOp>::OpRewritePattern;
+  LogicalResult matchAndRewrite(BitcastOp op,
+                                PatternRewriter &rewriter) const override {
+
+    Value in = op.getSelf();
+    auto loc = op.getLoc();
+    auto inType = cast<torch::Torch::ValueTensorType>(in.getType());
+    auto resultType = cast<torch::Torch::ValueTensorType>(op.getType());
+    auto bType = inType.toBuiltinTensor();
+
+    Value builtinCast =
+        rewriter.create<torch::TorchConversion::ToBuiltinTensorOp>(loc, bType,
+                                                                   in);
+    auto rType = resultType.toBuiltinTensor();
+
+    // Cast to the builtin tensor type.
+    llvm::SmallVector<Value> dynDims;
+    for (int i = 0, s = bType.getRank(); i < s; ++i) {
+      if (bType.isDynamicDim(i)) {
+        dynDims.push_back(rewriter.create<tensor::DimOp>(loc, builtinCast, i));
+      }
+    }
+
+    Value flowBitcast = rewriter.create<IREE::Flow::TensorBitCastOp>(
+        loc, rType, builtinCast, dynDims, dynDims);
+
+    auto torchCast =
+        rewriter.create<torch::TorchConversion::FromBuiltinTensorOp>(
+            loc, resultType, flowBitcast);
+    rewriter.replaceOp(op, torchCast);
+    return success();
+  }
+};
+
+class BitCastizedMatmul : public OpRewritePattern<torch::Torch::OperatorOp> {
 public:
   using OpRewritePattern::OpRewritePattern;
   LogicalResult matchAndRewrite(torch::Torch::OperatorOp op,
@@ -148,8 +184,8 @@ class BitCastQuantizedMatmul
 } // namespace
 
 namespace {
-class BitCastQuantTensorPass final
-    : public impl::BitCastQuantTensorPassBase<BitCastQuantTensorPass> {
+class BitCastTensorPass final
+    : public impl::BitCastTensorPassBase<BitCastTensorPass> {
   void getDependentDialects(DialectRegistry &registry) const override {
     registry.insert<IREE::Flow::FlowDialect>();
     registry.insert<torch::Torch::TorchDialect>();
@@ -159,7 +195,9 @@ class BitCastQuantTensorPass final
   void runOnOperation() override {
     MLIRContext *context = &getContext();
     RewritePatternSet patterns(context);
-    patterns.add<BitCastQuantizedMatmul, BitCastViewDtype>(context);
+    patterns.add<BitCastizedMatmul, BitCastViewDtype,
+                 BitCastViewComplex<torch::Torch::AtenViewAsComplexOp>,
+                 BitCastViewComplex<torch::Torch::AtenViewAsRealOp>>(context);
     if (failed(
             applyPatternsAndFoldGreedily(getOperation(), std::move(patterns))))
       signalPassFailure();

compiler/plugins/input/Torch/InputConversion/CMakeLists.txt

@@ -35,7 +35,7 @@ iree_cc_library(
     "Passes.h"
   SRCS
     "BindSymbolicShapes.cpp"
-    "BitCastQuantTensor.cpp"
+    "BitCastTensor.cpp"
     "ConvertTMTensorToLinalgExt.cpp"
     "FuncConversion.cpp"
     "SetStrictSymbolicShapes.cpp"

compiler/plugins/input/Torch/InputConversion/Passes.cpp

@@ -46,7 +46,7 @@ void createTorchToIREEPipeline(
     // now be simplified.
     pm.addNestedPass<func::FuncOp>(createCanonicalizerPass());
   }
-  pm.addNestedPass<func::FuncOp>(createBitCastQuantTensorPass());
+  pm.addNestedPass<func::FuncOp>(createBitCastTensorPass());
   pm.addNestedPass<func::FuncOp>(
       torch::Torch::createReduceOpVariantsPass(llvm::StringRef()));
   pm.addNestedPass<func::FuncOp>(

compiler/plugins/input/Torch/InputConversion/Passes.td

@@ -14,8 +14,8 @@ def BindSymbolicShapesPass :
   let summary = "Process torch dynamic shape bindings into IREE analyzable forms";
 }
 
-def BitCastQuantTensorPass :
-    InterfacePass<"torch-iree-bitcast-quant-tensor", "mlir::FunctionOpInterface"> {
+def BitCastTensorPass :
+    InterfacePass<"torch-iree-bitcast-tensor", "mlir::FunctionOpInterface"> {
   let summary = "Bitcasts i8 packed tensors of sub-byte types to the actual bit width";
 }
 

compiler/plugins/input/Torch/InputConversion/test/CMakeLists.txt

@@ -7,7 +7,7 @@ iree_lit_test_suite(
     "auto_input_conversion.mlir"
     "attention.mlir"
     "bind_symbolic_shapes.mlir"
-    "bitcast_quant_tensor.mlir"
+    "bitcast_tensor.mlir"
     "func_conversion.mlir"
     "func_conversion_invalid.mlir"
     "scan.mlir"

compiler/plugins/input/Torch/InputConversion/test/bitcast_quant_tensor.mlir → compiler/plugins/input/Torch/InputConversion/test/bitcast_tensor.mlir

@@ -1,4 +1,4 @@
-// RUN: iree-opt %s --pass-pipeline="builtin.module(func.func(torch-iree-bitcast-quant-tensor))" -split-input-file -verify-diagnostics | FileCheck %s
+// RUN: iree-opt %s --pass-pipeline="builtin.module(func.func(torch-iree-bitcast-tensor))" -split-input-file -verify-diagnostics | FileCheck %s
 
 // CHECK-LABEL: func @forward
 func.func @forward(%arg0: !torch.vtensor<[1,1,8],f16>) -> !torch.vtensor<[1,1,8],f16> {
@@ -24,3 +24,22 @@ func.func @view_type(%arg0 : !torch.vtensor<[295501824],ui8>) -> !torch.vtensor<
     %0 = torch.aten.view.dtype %arg0, %int4 : !torch.vtensor<[295501824],ui8>, !torch.int -> !torch.vtensor<[147750912],si16>
     return %0 : !torch.vtensor<[147750912],si16>
 }
+
+
+// -----
+
+// CHECK-LABEL: @view_as_complex
+func.func @view_as_complex(%arg0 : !torch.vtensor<[128,2],f32>) -> !torch.vtensor<[128], complex<f32>> {
+    // CHECK: flow.tensor.bitcast %[[IN:.+]] : tensor<128x2xf32> -> tensor<128xcomplex<f32>>
+    %0 = torch.aten.view_as_complex %arg0 : !torch.vtensor<[128,2],f32> -> !torch.vtensor<[128],complex<f32>>
+    return %0 : !torch.vtensor<[128],complex<f32>>
+}
+
+// -----
+
+// CHECK-LABEL: @view_as_real
+func.func @view_as_real(%arg0 : !torch.vtensor<[128], complex<f32>>) -> !torch.vtensor<[128,2],f32> {
+    // CHECK: flow.tensor.bitcast %[[IN:.+]] : tensor<128xcomplex<f32>> -> tensor<128x2xf32>
+    %0 = torch.aten.view_as_real %arg0 : !torch.vtensor<[128],complex<f32>> -> !torch.vtensor<[128,2],f32>
+    return %0 : !torch.vtensor<[128,2],f32>
+}