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

// Copyright 2022 The IREE Authors
//
// Licensed under the Apache License v2.0 with LLVM Exceptions.
// See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception

#ifndef IREE_COMPILER_PLUGINS_INPUT_TORCH_INPUTCONVERSION_PASSES
#define IREE_COMPILER_PLUGINS_INPUT_TORCH_INPUTCONVERSION_PASSES

include "mlir/Pass/PassBase.td"

def BindSymbolicShapesPass :
    InterfacePass<"torch-iree-bind-symbolic-shapes", "mlir::FunctionOpInterface"> {
  let summary = "Process torch dynamic shape bindings into IREE analyzable forms";
}

def BitCastTensorPass :
    InterfacePass<"torch-iree-bitcast-tensor", "mlir::FunctionOpInterface"> {
  let summary = "Bitcasts i8 packed tensors of sub-byte types to the actual bit width";
}

def ConvertTMTensorToLinalgExtPass :
    InterfacePass<"torch-iree-tm-tensor-to-linalg-ext", "mlir::FunctionOpInterface"> {
  let summary = "Convert from TMTensor ops to LinalgExt ops on tensors";
}

def ConvertTorchUnstructuredToLinalgExtPass :
    InterfacePass<"torch-iree-torch-unstructured-to-linalg-ext", "mlir::FunctionOpInterface"> {
  let summary = "Convert unstructured Torch ops to LinalgExt ops";
}

def ConvertTorchQuantizationToLinalgExtPass :
    InterfacePass<"torch-iree-torch-quantization-to-linalg-ext", "mlir::FunctionOpInterface"> {
  let summary = "Convert PT2E quantization operations to LinalgExt";
  let description = [{
    Preserves PT2E per-tensor, per-channel, per-channel-group, and per-token
    quantization as `iree_linalg_ext.quantize_affine` and `dequantize_affine`
    operations. This pass runs before Torch-to-Linalg expands the PT2E
    operations into generic arithmetic so later IREE optimization can still
    identify quantized computations.
  }];
}

def SetStrictSymbolicShapesPass :
    InterfacePass<"torch-iree-set-strict-symbolic-shapes", "mlir::FunctionOpInterface"> {
  let summary = "Adds the attribute indicating strict symbolic shapes in Torch IR";
}

def FuncConversionPass :
    Pass<"torch-iree-func-conversion", "ModuleOp"> {
  let summary = "Finalizes conversion from torch to IREE";
  let description = [{
    Conversion pass for finalizing functions and ABI. Replaces the generic
    torch-func-backend-type-conversion pass.
  }];
  let options = [
    Option<"externalizeTransients", "externalize-transients",
           "bool",
           /*default=*/"false",
          "If enabled, this option will append an external hal buffer to "
          "program inputs. This buffer will be used for storing transient "
          "memory and must be provided by the user.">];
}

def SyncFuncConversionPass :
    Pass<"torch-iree-sync-func-conversion", "ModuleOp"> {
  let summary = "Finalizes conversion from torch to IREE with a sync-only ABI";
  let description = [{
    Synchronous-only alternative to torch-iree-func-conversion for targets
    where the coarse-fences ABI is undesirable, such as programs compiled for
    the inline HAL. Entry points keep their original name and a plain builtin
    tensor ABI: no fences, no HAL imports/exports, and no $async variant.
    Mutable tensor arguments require in-place aliasing through the HAL and are
    not supported.
  }];
}

#endif // IREE_COMPILER_PLUGINS_INPUT_TORCH_INPUTCONVERSION_PASSES