[AutoBump] Merge with fixes of 297c2709 (May 20) (43)

.gitignore

@@ -11,6 +11,7 @@ externals/pytorch/
 libtorch*
 
 /build/
+.build-cache/
 /setup_build/
 __pycache__
 *.pyc

include/torch-mlir/Conversion/TorchToStablehlo/StablehloLegalizeUtils.h

@@ -22,6 +22,10 @@ namespace hlo {
 
 using mlir::ConversionPatternRewriter;
 
+// Create chlo::ConstantLikeOp
+template <typename T>
+Value getConstantLike(OpBuilder &rewriter, Location loc, T constant, Value val);
+
 // Create a 32-bit float constant operator from a float
 Value getStablehloConstTensorSingleF32(PatternRewriter &rewriter, Operation *op,
                                        float val);

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

@@ -13577,6 +13577,31 @@ def Torch_AtenSplitSizesOp : Torch_Op<"aten.split.sizes", [
       printDefaultTorchOp(printer, *this, 3, 1);
     }
   }];
+  let hasCanonicalizer = 1;
+}
+
+def Torch_AtenTensorSplitSectionsOp : Torch_Op<"aten.tensor_split.sections", [
+    AllowsTypeRefinement,
+    ReadOnly
+  ]> {
+  let summary = "Generated op for `aten::tensor_split.sections : (Tensor, int, int) -> (Tensor[])`";
+  let arguments = (ins
+    AnyTorchTensorType:$self,
+    Torch_IntType:$sections,
+    Torch_IntType:$dim
+  );
+  let results = (outs
+    AnyTorchListOfTensorType:$result
+  );
+  let hasCustomAssemblyFormat = 1;
+  let extraClassDefinition = [{
+    ParseResult AtenTensorSplitSectionsOp::parse(OpAsmParser &parser, OperationState &result) {
+      return parseDefaultTorchOp(parser, result, 3, 1);
+    }
+    void AtenTensorSplitSectionsOp::print(OpAsmPrinter &printer) {
+      printDefaultTorchOp(printer, *this, 3, 1);
+    }
+  }];
 }
 
 def Torch_AtenUnbindIntOp : Torch_Op<"aten.unbind.int", [

include/torch-mlir/Dialect/Torch/Transforms/Passes.h

@@ -73,6 +73,11 @@ struct TorchLoweringPipelineOptions
 void createTorchScriptModuleToTorchBackendPipeline(
     OpPassManager &pm, const TorchLoweringPipelineOptions &options);
 
+/// Creates a pipeline that lowers the graph IR that is produced by
+/// TorchDynamo export into the form expected by torch-verify-backend-contract.
+void createTorchDynamoExportToTorchBackendPipeline(
+    OpPassManager &pm, const TorchLoweringPipelineOptions &options);
+
 /// Creates a pipeline that lowers a flat list of funcs and global slots
 /// with the torch and aten dialects and mutable arrays and converts it to
 /// the form required by torch-verify-backend-contract.

include/torch-mlir/Dialect/Torch/Utils/Utils.h

@@ -101,6 +101,10 @@ int64_t getNumberOfElements(RankedTensorType inputType);
 SmallVector<int64_t> makeShapeLLVMCompatible(ArrayRef<int64_t> shape);
 SmallVector<int64_t> makeShapeTorchCompatible(ArrayRef<int64_t> shape);
 
+ValueTensorType getTensorTypeFromShapeValues(ArrayRef<Value> shapes,
+                                             Type dtype);
+Value getTensorDimSize(PatternRewriter &rewriter, Value tensor, int64_t dim);
+
 // Helper function to squeeze the input tensor at given dim.
 // Return the squeezed tensor or failure.
 FailureOr<Value> squeezeTensor(PatternRewriter &rewriter, Operation *op,

lib/CAPI/TorchTypes.cpp

@@ -51,7 +51,7 @@ MlirType torchMlirTorchOptionalTypeGet(MlirType containedType) {
 }
 
 MlirType torchMlirTorchOptionalTypeGetContained(MlirType t) {
-  auto type = unwrap(t).cast<Torch::OptionalType>();
+  auto type = cast<Torch::OptionalType>(unwrap(t));
   return wrap(type.getContainedType());
 }
 
@@ -77,12 +77,12 @@ MlirType torchMlirTorchTupleTypeGet(MlirContext context,
 }
 
 size_t torchMlirTorchTupleTypeGetNumTypes(MlirType t) {
-  auto type = unwrap(t).cast<Torch::TupleType>();
+  auto type = cast<Torch::TupleType>(unwrap(t));
   return type.getContainedTypes().size();
 }
 
 MlirType torchMlirTorchTupleTypeGetType(MlirType t, intptr_t pos) {
-  auto type = unwrap(t).cast<Torch::TupleType>();
+  auto type = cast<Torch::TupleType>(unwrap(t));
   return wrap(type.getContainedTypes()[pos]);
 }
 
@@ -108,12 +108,12 @@ MlirType torchMlirTorchUnionTypeGet(MlirContext context,
 }
 
 size_t torchMlirTorchUnionTypeGetNumTypes(MlirType t) {
-  auto type = unwrap(t).cast<Torch::UnionType>();
+  auto type = cast<Torch::UnionType>(unwrap(t));
   return type.getContainedTypes().size();
 }
 
 MlirType torchMlirTorchUnionTypeGetType(MlirType t, intptr_t pos) {
-  auto type = unwrap(t).cast<Torch::UnionType>();
+  auto type = cast<Torch::UnionType>(unwrap(t));
   return wrap(type.getContainedTypes()[pos]);
 }
 
@@ -134,7 +134,7 @@ MlirType torchMlirTorchListTypeGet(MlirType containedType) {
 }
 
 MlirType torchMlirTorchListTypeGetContainedType(MlirType t) {
-  return wrap(unwrap(t).cast<Torch::ListType>().getContainedType());
+  return wrap(cast<Torch::ListType>(unwrap(t)).getContainedType());
 }
 
 MlirTypeID torchMlirTorchListTypeGetTypeID() {
@@ -297,26 +297,26 @@ MlirType torchMlirTorchNonValueTensorTypeGetWithLeastStaticInformation(
 
 MlirType torchMlirTorchNonValueTensorTypeGetFromAttribute(MlirAttribute attr) {
   auto attrTensorType =
-      unwrap(attr).cast<TypedAttr>().getType().cast<RankedTensorType>();
+      cast<RankedTensorType>(cast<TypedAttr>(unwrap(attr)).getType());
   return wrap(Torch::NonValueTensorType::get(attrTensorType.getContext(),
                                              attrTensorType.getShape(),
                                              attrTensorType.getElementType()));
 }
 
 int64_t torchMlirTorchNonValueTensorTypeGetRank(MlirType t) {
-  return unwrap(t).cast<Torch::NonValueTensorType>().getSizes().size();
+  return cast<Torch::NonValueTensorType>(unwrap(t)).getSizes().size();
 }
 
 bool torchMlirTorchNonValueTensorTypeHasSizes(MlirType t) {
-  return unwrap(t).cast<Torch::NonValueTensorType>().hasSizes();
+  return cast<Torch::NonValueTensorType>(unwrap(t)).hasSizes();
 }
 
 bool torchMlirTorchNonValueTensorTypeHasDtype(MlirType t) {
-  return unwrap(t).cast<Torch::NonValueTensorType>().hasDtype();
+  return cast<Torch::NonValueTensorType>(unwrap(t)).hasDtype();
 }
 
 int64_t torchMlirTorchNonValueTensorTypeGetSizes(MlirType t, int64_t *sizes) {
-  auto tensorType = unwrap(t).cast<Torch::NonValueTensorType>();
+  auto tensorType = cast<Torch::NonValueTensorType>(unwrap(t));
   bool hasSizes = tensorType.hasSizes();
   if (!hasSizes)
     return -1;
@@ -329,7 +329,7 @@ int64_t torchMlirTorchNonValueTensorTypeGetSizes(MlirType t, int64_t *sizes) {
 }
 
 MlirType torchMlirTorchNonValueTensorTypeGetDtype(MlirType t) {
-  return wrap(unwrap(t).cast<Torch::NonValueTensorType>().getDtype());
+  return wrap(cast<Torch::NonValueTensorType>(unwrap(t)).getDtype());
 }
 
 MlirTypeID torchMlirTorchNonValueTensorTypeGetTypeID() {
@@ -364,26 +364,26 @@ MlirType torchMlirTorchValueTensorTypeGetWithLeastStaticInformation(
 
 MlirType torchMlirTorchValueTensorTypeGetFromAttribute(MlirAttribute attr) {
   auto attrTensorType =
-      unwrap(attr).cast<TypedAttr>().getType().cast<RankedTensorType>();
+      cast<RankedTensorType>(cast<TypedAttr>(unwrap(attr)).getType());
   return wrap(Torch::ValueTensorType::get(attrTensorType.getContext(),
                                           attrTensorType.getShape(),
                                           attrTensorType.getElementType()));
 }
 
 int64_t torchMlirTorchValueTensorTypeGetRank(MlirType t) {
-  return unwrap(t).cast<Torch::ValueTensorType>().getSizes().size();
+  return cast<Torch::ValueTensorType>(unwrap(t)).getSizes().size();
 }
 
 bool torchMlirTorchValueTensorTypeHasSizes(MlirType t) {
-  return unwrap(t).cast<Torch::ValueTensorType>().hasSizes();
+  return cast<Torch::ValueTensorType>(unwrap(t)).hasSizes();
 }
 
 bool torchMlirTorchValueTensorTypeHasDtype(MlirType t) {
-  return unwrap(t).cast<Torch::ValueTensorType>().hasDtype();
+  return cast<Torch::ValueTensorType>(unwrap(t)).hasDtype();
 }
 
 int64_t torchMlirTorchValueTensorTypeGetSizes(MlirType t, int64_t *sizes) {
-  auto tensorType = unwrap(t).cast<Torch::ValueTensorType>();
+  auto tensorType = cast<Torch::ValueTensorType>(unwrap(t));
   bool hasSizes = tensorType.hasSizes();
   if (!hasSizes)
     return -1;
@@ -396,7 +396,7 @@ int64_t torchMlirTorchValueTensorTypeGetSizes(MlirType t, int64_t *sizes) {
 }
 
 MlirType torchMlirTorchValueTensorTypeGetDtype(MlirType t) {
-  return wrap(unwrap(t).cast<Torch::ValueTensorType>().getDtype());
+  return wrap(cast<Torch::ValueTensorType>(unwrap(t)).getDtype());
 }
 
 MlirTypeID torchMlirTorchValueTensorTypeGetTypeID() {
@@ -487,12 +487,12 @@ MlirType torchMlirTorchDictTypeGetChecked(MlirContext context, MlirType keyType,
 }
 
 MlirType torchMlirTorchDictTypeGetKeyType(MlirType t) {
-  auto type = unwrap(t).cast<Torch::DictType>();
+  auto type = cast<Torch::DictType>(unwrap(t));
   return wrap(type.getKeyType());
 }
 
 MlirType torchMlirTorchDictTypeGetValueType(MlirType t) {
-  auto type = unwrap(t).cast<Torch::DictType>();
+  auto type = cast<Torch::DictType>(unwrap(t));
   return wrap(type.getValueType());
 }
 

lib/Conversion/TorchConversionToMLProgram/TorchConversionToMLProgram.cpp

@@ -13,10 +13,6 @@
 #include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Dialect/MLProgram/IR/MLProgram.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
-#include "mlir/IR/Builders.h"
-#include "mlir/IR/BuiltinAttributes.h"
-#include "mlir/IR/BuiltinTypes.h"
-#include "mlir/IR/MLIRContext.h"
 #include "torch-mlir/Dialect/TorchConversion/IR/TorchConversionOps.h"
 #include "torch-mlir/Dialect/TorchConversion/Transforms/BackendTypeConversion.h"
 

lib/Conversion/TorchOnnxToTorch/DefaultDomainAtoF.cpp

@@ -63,7 +63,7 @@ LogicalResult windowFunctionImpl(OpBinder binder,
   // Create an f32 ValueTensorType with thse same size as size, the
   // operand
   auto shapeOfOperand =
-      size.getType().dyn_cast<Torch::ValueTensorType>().getOptionalSizes();
+      dyn_cast<Torch::ValueTensorType>(size.getType()).getOptionalSizes();
   auto f32ResultType = rewriter.getType<Torch::ValueTensorType>(
       shapeOfOperand, rewriter.getF32Type());
   Value periodicSizeFloat = b.create<Torch::AtenToDtypeOp>(
@@ -897,8 +897,8 @@ void mlir::torch::onnx_c::populateDefaultDomainAtoF(
         }
 
         if (DenseResourceElementsAttr attr =
-                binder.op->getAttr("torch.onnx.value")
-                    .dyn_cast_or_null<DenseResourceElementsAttr>()) {
+                dyn_cast_or_null<DenseResourceElementsAttr>(
+                    binder.op->getAttr("torch.onnx.value"))) {
           // Bytes are stored in little endian order. Big endian support will
           // require swizzling.
           if (!Endian::little) {
@@ -909,25 +909,34 @@ void mlir::torch::onnx_c::populateDefaultDomainAtoF(
 
           auto ty = cast<ShapedType>(attr.getType());
           ElementsAttr denseAttr;
-          auto ptr = attr.getRawHandle().getBlob()->getData();
+          auto ptr = attr.getRawHandle().getBlob();
+          if (!ptr) {
+            denseAttr = DenseResourceElementsAttr::get(
+                ty, "__onnx_constant_not_found_possibly_due_to_being_elided__",
+                AsmResourceBlob());
+            rewriter.replaceOpWithNewOp<Torch::ValueTensorLiteralOp>(
+                binder.op, resultType, denseAttr);
+            return success();
+          }
+          auto data = ptr->getData();
           if (cast<ShapedType>(attr.getType()).getElementType().isInteger(1)) {
             llvm::SmallVector<APInt> newContents;
-            for (auto val : ptr) {
+            for (auto val : data) {
               APInt apval(1, val);
               newContents.push_back(apval);
             }
             denseAttr = DenseElementsAttr::get(ty, newContents);
           } else {
-            denseAttr = DenseElementsAttr::getFromRawBuffer(ty, ptr);
+            denseAttr = DenseElementsAttr::getFromRawBuffer(ty, data);
           }
 
           rewriter.replaceOpWithNewOp<Torch::ValueTensorLiteralOp>(
               binder.op, resultType, denseAttr);
           return success();
         }
 
-        if (ElementsAttr attr = binder.op->getAttr("torch.onnx.value")
-                                    .dyn_cast_or_null<ElementsAttr>()) {
+        if (ElementsAttr attr = dyn_cast_or_null<ElementsAttr>(
+                binder.op->getAttr("torch.onnx.value"))) {
           rewriter.replaceOpWithNewOp<Torch::ValueTensorLiteralOp>(
               binder.op, resultType, attr);
           return success();
@@ -2283,9 +2292,7 @@ void mlir::torch::onnx_c::populateDefaultDomainAtoF(
             binder.tensorResultType(resultType))
           return failure();
         Type listElemType =
-            tensors[0]
-                .getType()
-                .cast<Torch::BaseTensorType>()
+            cast<Torch::BaseTensorType>(tensors[0].getType())
                 .getWithSizesAndDtype(/*optionalSizes=*/std::nullopt,
                                       /*optionalDtype=*/nullptr);
         Type listType = Torch::ListType::get(listElemType);