Merge pull request #2 from llvm/main

Get latest changes

Author: ivangarcia44

.github/workflows/pre-commit-all.yml

@@ -9,7 +9,7 @@ jobs:
     runs-on: ubuntu-22.04
     steps:
     - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2
-    - uses: actions/setup-python@0b93645e9fea7318ecaed2b359559ac225c90a2b # v5.3.0
+    - uses: actions/setup-python@42375524e23c412d93fb67b49958b491fce71c38 # v5.4.0
     - uses: pre-commit/action@2c7b3805fd2a0fd8c1884dcaebf91fc102a13ecd # v3.0.1
       with:
         extra_args: --color=always --all-files

.github/workflows/pre-commit.yml

@@ -11,7 +11,7 @@ jobs:
       with:
         # requites to grab the history of the PR
         fetch-depth: 0
-    - uses: actions/setup-python@0b93645e9fea7318ecaed2b359559ac225c90a2b # v5.3.0
+    - uses: actions/setup-python@42375524e23c412d93fb67b49958b491fce71c38 # v5.4.0
     - uses: pre-commit/action@2c7b3805fd2a0fd8c1884dcaebf91fc102a13ecd # v3.0.1
       with:
         extra_args: --color=always --from-ref ${{ github.event.pull_request.base.sha }} --to-ref ${{ github.event.pull_request.head.sha }}

CMakeLists.txt

@@ -35,13 +35,17 @@ option(TORCH_MLIR_ENABLE_WERROR_FLAG "Enable `-Werror` flag on supported directo
 option(TORCH_MLIR_USE_INSTALLED_PYTORCH "If depending on PyTorch use it as installed in the current Python environment" ON)
 
 option(TORCH_MLIR_ENABLE_REFBACKEND "Enable reference backend" ON)
+
 if(TORCH_MLIR_ENABLE_REFBACKEND)
   add_definitions(-DTORCH_MLIR_ENABLE_REFBACKEND)
 endif()
 
+set(TORCH_MLIR_TABLEGEN_FLAGS "")
+
 option(TORCH_MLIR_ENABLE_STABLEHLO "Add stablehlo dialect" ON)
 if(TORCH_MLIR_ENABLE_STABLEHLO)
   add_definitions(-DTORCH_MLIR_ENABLE_STABLEHLO)
+  list(APPEND TORCH_MLIR_TABLEGEN_FLAGS "-DTORCH_MLIR_ENABLE_STABLEHLO")
 endif()
 # It is possible that both stablehlo and torch_mlir projects are used in some compiler project.
 # In this case, we don't want to use stablehlo that is downloaded by torch_mlir (in external/stablehlo)
@@ -50,6 +54,12 @@ endif()
 # stablehlo targets AND includes available (for example with `add_subdirectory` and `include_directories`).
 option(TORCH_MLIR_USE_EXTERNAL_STABLEHLO "Use stablehlo from top level project" OFF)
 
+option(TORCH_MLIR_ENABLE_TOSA "Add TOSA support" ON)
+if(TORCH_MLIR_ENABLE_TOSA)
+  add_definitions(-DTORCH_MLIR_ENABLE_TOSA)
+  list(APPEND TORCH_MLIR_TABLEGEN_FLAGS "-DTORCH_MLIR_ENABLE_TOSA")
+endif()
+
 option(TORCH_MLIR_OUT_OF_TREE_BUILD "Specifies an out of tree build" OFF)
 
 # PyTorch native extension gate. If OFF, then no features which depend on

build_tools/autogen_ltc_backend.py

@@ -30,6 +30,12 @@
 TORCHGEN_DIR = Path(torchgen.__path__[0]).resolve()
 TORCH_MLIR_DIR = Path(__file__).resolve().parent.parent
 
+# Safely load fast C Yaml loader if it is are available
+try:
+    from yaml import CSafeLoader as Loader
+except ImportError:
+    from yaml import SafeLoader as Loader  # type:ignore[assignment, misc]
+
 
 def reindent(text, prefix=""):
     return indent(dedent(text), prefix)
@@ -175,7 +181,7 @@ def generate_native_functions(self):
         )
         ts_native_yaml = None
         if ts_native_yaml_path.exists():
-            ts_native_yaml = yaml.load(ts_native_yaml_path.read_text(), yaml.CLoader)
+            ts_native_yaml = yaml.load(ts_native_yaml_path.read_text(), Loader)
         else:
             logging.warning(
                 f"Could not find `ts_native_functions.yaml` at {ts_native_yaml_path}"
@@ -208,7 +214,7 @@ def get_opnames(ops):
         )
 
         with self.config_path.open() as f:
-            config = yaml.load(f, yaml.CLoader)
+            config = yaml.load(f, Loader)
 
         # List of unsupported ops in LTC autogen because of some error
         blacklist = set(config.get("blacklist", []))

externals/llvm-project

@@ -260,7 +260,7 @@ torchMlirTorchNonValueTensorTypeGetWithLeastStaticInformation(
 MLIR_CAPI_EXPORTED MlirType
 torchMlirTorchNonValueTensorTypeGetFromAttribute(MlirAttribute attr);
 
-/// Gets the the rank (number of dimensions) of a !torch.tensor
+/// Gets the rank (number of dimensions) of a !torch.tensor
 MLIR_CAPI_EXPORTED int64_t torchMlirTorchNonValueTensorTypeGetRank(MlirType t);
 
 /// Return true if this type has a list of sizes.
@@ -269,12 +269,12 @@ MLIR_CAPI_EXPORTED bool torchMlirTorchNonValueTensorTypeHasSizes(MlirType t);
 /// Return true if this type has a dtype.
 MLIR_CAPI_EXPORTED bool torchMlirTorchNonValueTensorTypeHasDtype(MlirType t);
 
-/// Gets the the sizes of the dimensions of a !torch.tensor; note -1 size
+/// Gets the sizes of the dimensions of a !torch.tensor; note -1 size
 /// indicates an unrefined/unknown size dimension.
 MLIR_CAPI_EXPORTED int64_t
 torchMlirTorchNonValueTensorTypeGetSizes(MlirType t, int64_t *sizes);
 
-/// Gets the the dtype (data type) of a !torch.tensor.
+/// Gets the dtype (data type) of a !torch.tensor.
 MLIR_CAPI_EXPORTED MlirType
 torchMlirTorchNonValueTensorTypeGetDtype(MlirType t);
 
@@ -307,7 +307,7 @@ torchMlirTorchValueTensorTypeGetWithLeastStaticInformation(MlirContext context);
 MLIR_CAPI_EXPORTED MlirType
 torchMlirTorchValueTensorTypeGetFromAttribute(MlirAttribute attr);
 
-/// Gets the the rank (number of dimensions) of a !torch.vtensor
+/// Gets the rank (number of dimensions) of a !torch.vtensor
 MLIR_CAPI_EXPORTED int64_t torchMlirTorchValueTensorTypeGetRank(MlirType t);
 
 /// Return true if this type has a list of sizes.
@@ -316,12 +316,12 @@ MLIR_CAPI_EXPORTED bool torchMlirTorchValueTensorTypeHasSizes(MlirType t);
 /// Return true if this type has a dtype.
 MLIR_CAPI_EXPORTED bool torchMlirTorchValueTensorTypeHasDtype(MlirType t);
 
-/// Gets the the sizes of the dimensions of a !torch.vtensor; note -1 size
+/// Gets the sizes of the dimensions of a !torch.vtensor; note -1 size
 /// indicates an unrefined/unknown size dimension.
 MLIR_CAPI_EXPORTED int64_t
 torchMlirTorchValueTensorTypeGetSizes(MlirType t, int64_t *sizes);
 
-/// Gets the the dtype (data type) of a !torch.vtensor.
+/// Gets the dtype (data type) of a !torch.vtensor.
 MLIR_CAPI_EXPORTED MlirType torchMlirTorchValueTensorTypeGetDtype(MlirType t);
 
 /// Gets the !torch.vtensor typeid.

include/torch-mlir-c/TorchTypes.h

@@ -1,11 +1,11 @@
 add_subdirectory(TorchOnnxToTorch)
 
 set(LLVM_TARGET_DEFINITIONS Passes.td)
-if(TORCH_MLIR_ENABLE_STABLEHLO)
-  mlir_tablegen(Passes.h.inc -gen-pass-decls -DTORCH_MLIR_ENABLE_STABLEHLO)
-else()
-  mlir_tablegen(Passes.h.inc -gen-pass-decls)
-endif()
+
+
+
+mlir_tablegen(Passes.h.inc -gen-pass-decls ${TORCH_MLIR_TABLEGEN_FLAGS})
+
 add_public_tablegen_target(TorchMLIRConversionPassIncGen)
 
 add_mlir_doc(Passes TorchMLIRConversionPasses ./ -gen-pass-doc)

include/torch-mlir/Conversion/CMakeLists.txt

@@ -114,6 +114,7 @@ def ConvertTorchToTensor : Pass<"convert-torch-to-tensor", "func::FuncOp"> {
   let constructor = "mlir::torch::createConvertTorchToTensorPass()";
 }
 
+#ifdef TORCH_MLIR_ENABLE_TOSA
 def ConvertTorchToTosa : Pass<"convert-torch-to-tosa", "func::FuncOp"> {
   let summary = "Convert Torch ops to TOSA ops";
   let description = [{
@@ -122,6 +123,7 @@ def ConvertTorchToTosa : Pass<"convert-torch-to-tosa", "func::FuncOp"> {
   }];
   let constructor = "mlir::torch::createConvertTorchToTosaPass()";
 }
+#endif
 
 def ConvertTorchToTMTensor : Pass<"convert-torch-to-tmtensor", "func::FuncOp"> {
   let summary = "Convert recognized Torch ops to TMTensor/Linalg ops";

include/torch-mlir/Conversion/Passes.td

@@ -24,6 +24,7 @@ createOneDimTfIndices(PatternRewriter &rewriter, Operation *op,
                       SmallVector<int64_t> indiceOneDimShape, int32_t dim,
                       ArrayRef<int64_t> indexShape);
 
+// Default function to create TOSA op with shift value
 mlir::tosa::MulOp createMulOpAndCast(PatternRewriter &rewriter, Operation *op,
                                      TensorType outType, Value lhs, Value rhs,
                                      int32_t shift);
@@ -32,8 +33,8 @@ mlir::tosa::MulOp createMulOpAndCast(PatternRewriter &rewriter, Operation *op,
 template <typename TosaOpT>
 TosaOpT createBinaryOpAndCast(PatternRewriter &rewriter, Operation *op,
                               TensorType outType, Value lhs, Value rhs) {
-  lhs = promoteType(rewriter, lhs, outType);
-  rhs = promoteType(rewriter, rhs, outType);
+  lhs = tosa::tosaCastTensorToType(rewriter, lhs, outType).value();
+  rhs = tosa::tosaCastTensorToType(rewriter, rhs, outType).value();
   return CreateOpAndInfer<TosaOpT>(rewriter, op->getLoc(), outType, lhs, rhs);
 }
 

include/torch-mlir/Conversion/TorchToTosa/TosaLegalizeCommon.h

@@ -10,13 +10,14 @@
 #ifndef TORCHMLIR_CONVERSION_TORCHTOTOSA_TOSALEGALIZEUTILS_H
 #define TORCHMLIR_CONVERSION_TORCHTOTOSA_TOSALEGALIZEUTILS_H
 
-#include "mlir/Dialect/Quant/IR/QuantTypes.h"     // from @llvm-project
-#include "mlir/Dialect/Tosa/Utils/ShapeUtils.h"   // from @llvm-project
-#include "mlir/IR/BuiltinAttributes.h"            // from @llvm-project
-#include "mlir/IR/BuiltinTypes.h"                 // from @llvm-project
-#include "mlir/IR/PatternMatch.h"                 // from @llvm-project
-#include "mlir/Interfaces/InferTypeOpInterface.h" // from @llvm-project
-#include "mlir/Support/LLVM.h"                    // from @llvm-project
+#include "mlir/Dialect/Quant/IR/QuantTypes.h"        // from @llvm-project
+#include "mlir/Dialect/Tosa/Utils/ConversionUtils.h" // from @llvm-project
+#include "mlir/Dialect/Tosa/Utils/ShapeUtils.h"      // from @llvm-project
+#include "mlir/IR/BuiltinAttributes.h"               // from @llvm-project
+#include "mlir/IR/BuiltinTypes.h"                    // from @llvm-project
+#include "mlir/IR/PatternMatch.h"                    // from @llvm-project
+#include "mlir/Interfaces/InferTypeOpInterface.h"    // from @llvm-project
+#include "mlir/Support/LLVM.h"                       // from @llvm-project
 
 namespace mlir {
 namespace tosa {
@@ -45,6 +46,10 @@ bool isScale32(mlir::quant::UniformQuantizedType output_element_type);
 Value getTosaConstTensorSingleF32(PatternRewriter &rewriter, Operation *op,
                                   float val);
 
+// Create an int8_t const tosa.mul shift tensor from an int
+Value getTosaMulShiftConstTensor(PatternRewriter &rewriter, Operation *op,
+                                 int32_t shift);
+
 // Create a zero constant tensor of the desired type and shape.
 std::optional<Value> getZerosLikeTensor(PatternRewriter &rewriter,
                                         Operation *op, Type type);
@@ -58,55 +63,24 @@ std::optional<Value> getConstTensor(PatternRewriter &rewriter, Operation *op,
                                     ArrayRef<T> vec, ArrayRef<int64_t> shape,
                                     std::optional<Type> dtype = {});
 
-LogicalResult tosaCastTensorToType(PatternRewriter &rewriter, Operation *op,
-                                   Value src, Type destType, Value &result);
-
-Value promoteType(PatternRewriter &rewriter, Value input, TensorType outType);
+// Default function to create tosa.cast op. This should be called instead of
+// directly calling rewriter.create<tosa::CastOp>.
+std::optional<Value> tosaCastTensorToType(PatternRewriter &rewriter, Value src,
+                                          TensorType destType);
 
 // Creates a TOSA operation and performs shape inference on the individual
 // op. This allows shape inference during the framework to TOSA lowering.
+template <typename TosaOp, typename... Args>
+TosaOp CreateOpAndInfer(ImplicitLocOpBuilder &builder, Type result_ty,
+                        Args &&...args) {
+  return CreateOpAndInferShape<TosaOp>(builder, result_ty, args...);
+}
+
 template <typename TosaOp, typename... Args>
 TosaOp CreateOpAndInfer(PatternRewriter &rewriter, Location loc, Type result_ty,
                         Args &&...args) {
-  auto op = rewriter.create<TosaOp>(loc, result_ty, args...);
-
-  InferShapedTypeOpInterface shapeInterface =
-      dyn_cast<InferShapedTypeOpInterface>(op.getOperation());
-  if (!shapeInterface)
-    return op;
-
-  SmallVector<ShapedTypeComponents> returnedShapes;
-  if (shapeInterface
-          .inferReturnTypeComponents(op.getContext(), op.getLoc(),
-                                     op->getOperands(), op->getAttrDictionary(),
-                                     op->getPropertiesStorage(),
-                                     op->getRegions(), returnedShapes)
-          .failed())
-    return op;
-
-  // We need to use the element type of the existing result type to generate
-  // the new result shaped type. This is because rescale can include a cast to
-  // different bit-width types and does not have a TypeAttr to define the
-  // target type.
-  auto result = op->getResult(0);
-  auto predictedShape = returnedShapes[0];
-  auto currentKnowledge = ValueKnowledge::getKnowledgeFromType(result_ty);
-
-  // Compute the knowledge based on the inferred type.
-  auto inferredKnowledge = ValueKnowledge::getPessimisticValueState();
-  inferredKnowledge.dtype = cast<ShapedType>(result_ty).getElementType();
-  inferredKnowledge.hasRank = predictedShape.hasRank();
-  if (predictedShape.hasRank()) {
-    for (auto dim : predictedShape.getDims()) {
-      inferredKnowledge.sizes.push_back(dim);
-    }
-  }
-
-  // Compute the new type based on the joined version.
-  auto newKnowledge = ValueKnowledge::join(currentKnowledge, inferredKnowledge);
-  auto new_ty = newKnowledge.getType();
-  result.setType(new_ty);
-  return op;
+  ImplicitLocOpBuilder builder(loc, rewriter);
+  return CreateOpAndInfer<TosaOp>(builder, result_ty, args...);
 }
 
 template <typename TosaOp, typename... Args>

include/torch-mlir/Conversion/TorchToTosa/TosaLegalizeUtils.h

@@ -5236,6 +5236,7 @@ def Torch_AtenPowTensorScalarOp : Torch_Op<"aten.pow.Tensor_Scalar", [
       printDefaultTorchOp(printer, *this, 2, 1);
     }
   }];
+  let hasCanonicalizer = 1;
 }
 
 def Torch_AtenPowTensorTensorOp : Torch_Op<"aten.pow.Tensor_Tensor", [
@@ -5260,6 +5261,7 @@ def Torch_AtenPowTensorTensorOp : Torch_Op<"aten.pow.Tensor_Tensor", [
       printDefaultTorchOp(printer, *this, 2, 1);
     }
   }];
+  let hasCanonicalizer = 1;
 }
 
 def Torch_AtenPowScalarOp : Torch_Op<"aten.pow.Scalar", [
@@ -13704,7 +13706,7 @@ def Torch_AtenStftOp : Torch_Op<"aten.stft", [
     HasValueSemantics,
     ReadOnly
   ]> {
-  let summary = "Generated op for `aten::stft : (Tensor, int, int?, int?, Tensor?, bool, bool?, bool?) -> (Tensor)`";
+  let summary = "Generated op for `aten::stft : (Tensor, int, int?, int?, Tensor?, bool, bool?, bool?, bool?) -> (Tensor)`";
   let arguments = (ins
     AnyTorchTensorType:$self,
     Torch_IntType:$n_fft,
@@ -13713,18 +13715,19 @@ def Torch_AtenStftOp : Torch_Op<"aten.stft", [
     AnyTorchOptionalTensorType:$window,
     Torch_BoolType:$normalized,
     AnyTorchOptionalBoolType:$onesided,
-    AnyTorchOptionalBoolType:$return_complex
+    AnyTorchOptionalBoolType:$return_complex,
+    AnyTorchOptionalBoolType:$align_to_window
   );
   let results = (outs
     AnyTorchOptionalTensorType:$result
   );
   let hasCustomAssemblyFormat = 1;
   let extraClassDefinition = [{
     ParseResult AtenStftOp::parse(OpAsmParser &parser, OperationState &result) {
-      return parseDefaultTorchOp(parser, result, 8, 1);
+      return parseDefaultTorchOp(parser, result, 9, 1);
     }
     void AtenStftOp::print(OpAsmPrinter &printer) {
-      printDefaultTorchOp(printer, *this, 8, 1);
+      printDefaultTorchOp(printer, *this, 9, 1);
     }
   }];
 }
@@ -15175,6 +15178,31 @@ def Torch_AtenSortOp : Torch_Op<"aten.sort", [
   let hasFolder = 1;
 }
 
+def Torch_AtenArgsortOp : Torch_Op<"aten.argsort", [
+    AllowsTypeRefinement,
+    HasValueSemantics,
+    ReadOnly
+  ]> {
+  let summary = "Generated op for `aten::argsort : (Tensor, int, bool) -> (Tensor)`";
+  let arguments = (ins
+    AnyTorchTensorType:$self,
+    Torch_IntType:$dim,
+    Torch_BoolType:$descending
+  );
+  let results = (outs
+    AnyTorchOptionalTensorType:$result
+  );
+  let hasCustomAssemblyFormat = 1;
+  let extraClassDefinition = [{
+    ParseResult AtenArgsortOp::parse(OpAsmParser &parser, OperationState &result) {
+      return parseDefaultTorchOp(parser, result, 3, 1);
+    }
+    void AtenArgsortOp::print(OpAsmPrinter &printer) {
+      printDefaultTorchOp(printer, *this, 3, 1);
+    }
+  }];
+}
+
 def Torch_AtenSplitTensorOp : Torch_Op<"aten.split.Tensor", [
     AllowsTypeRefinement,
     ReadOnly

include/torch-mlir/Dialect/Torch/IR/GeneratedTorchOps.td

@@ -1,9 +1,7 @@
 set(LLVM_TARGET_DEFINITIONS Passes.td)
-if(TORCH_MLIR_ENABLE_STABLEHLO)
-  mlir_tablegen(Passes.h.inc -gen-pass-decls -DTORCH_MLIR_ENABLE_STABLEHLO)
-else()
-  mlir_tablegen(Passes.h.inc -gen-pass-decls)
-endif()
+
+mlir_tablegen(Passes.h.inc -gen-pass-decls ${TORCH_MLIR_TABLEGEN_FLAGS})
+
 add_public_tablegen_target(TorchMLIRTorchConversionPassIncGen)
 
 add_mlir_doc(Passes TorchMLIRTorchConversionTransforms ./ -gen-pass-doc)