Bump llvm to llvm/llvm-project@5d6d982

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

@@ -127,11 +127,6 @@ LogicalResult getConvOpsAccType(PatternRewriter &rewriter,
                                 RankedTensorType weightTy,
                                 RankedTensorType outputTy, TypeAttr &accType);
 
-// Temporary function to get TOSA const shape
-// TODO: Remove this function when getTosaConstShape is available in
-// externals/llvm-project/mlir/include/mlir/Dialect/Tosa/Utils/ConversionUtils.h
-Value getTosaConstShape(PatternRewriter &rewriter, Location loc,
-                        llvm::ArrayRef<int64_t> shape);
 } // namespace tosa
 } // namespace mlir
 

lib/Conversion/TorchToTosa/TosaLegalizeCommon.cpp

@@ -351,7 +351,7 @@ std::optional<Value> convertGatherNdOp(PatternRewriter &rewriter, Operation *op,
 
   // %3 = "tosa.reshape"(%1) {new_shape = [8, 3]} : (tensor<1x4x2x3xi32>) ->
   // tensor<8x3xi32> Flatten the input indices tensor to an [W, ND] matrix.
-  auto indicesMatrixReshapeOp = tosa::CreateOpAndInfer<tosa::ReshapeOp>(
+  Value indicesMatrixReshapeOp = tosa::CreateOpAndInfer<tosa::ReshapeOp>(
       rewriter, op->getLoc(),
       GetTypeFromTensorShape(indicesMatrixShape, indicesType.getElementType()),
       indicesValue, rewriter.getDenseI64ArrayAttr(indicesMatrixShape));
@@ -378,13 +378,18 @@ std::optional<Value> convertGatherNdOp(PatternRewriter &rewriter, Operation *op,
   if (!flattenedCoeffValue)
     return std::nullopt;
 
+  if (mlir::tosa::EqualizeRanks(rewriter, op->getLoc(), indicesMatrixReshapeOp,
+                                flattenedCoeffValue.value())
+          .failed())
+    return std::nullopt;
+
   // Multiply the coefficients by the coordinates
   // %5 = "tosa.mul"(%3, %4) {shift = 0 : i32} : (tensor<8x3xi32>,
   // tensor<3xi32>) -> tensor<8x3xi32>
   auto flattenedIndicesMulOp = tosa::CreateOpAndInfer<tosa::MulOp>(
       rewriter, op->getLoc(),
       GetTypeFromTensorShape(indicesMatrixShape, indicesType.getElementType()),
-      indicesMatrixReshapeOp.getResult(), flattenedCoeffValue.value(), 0);
+      indicesMatrixReshapeOp, flattenedCoeffValue.value(), 0);
 
   // Sum up the products of the coefficients and coordinates
   // %6 = "tosa.reduce_sum"(%5) {axis = 1 : i64} : (tensor<8x3xi32>) ->
@@ -616,7 +621,7 @@ std::optional<Value> convertScatterNdOp(PatternRewriter &rewriter,
   // [[0, 1], [0, 2],  [0, 3]] -> [[0, 1], [0, 2],  [0, 3]]
   // %11 = "tosa.reshape"(%8) {new_shape = array<i64: 3, 2>} : (tensor<3x2xi32>)
   // -> tensor<3x2xi32>
-  auto indicesMatrixReshapeOp = tosa::CreateOpAndInfer<tosa::ReshapeOp>(
+  Value indicesMatrixReshapeOp = tosa::CreateOpAndInfer<tosa::ReshapeOp>(
       rewriter, op->getLoc(),
       GetTypeFromTensorShape(indicesMatrixShape, indicesType.getElementType()),
       indicesValue, rewriter.getDenseI64ArrayAttr(indicesMatrixShape));
@@ -643,14 +648,19 @@ std::optional<Value> convertScatterNdOp(PatternRewriter &rewriter,
   if (!flattenedCoeffValue)
     return std::nullopt;
 
+  if (mlir::tosa::EqualizeRanks(rewriter, op->getLoc(), indicesMatrixReshapeOp,
+                                flattenedCoeffValue.value())
+          .failed())
+    return std::nullopt;
+
   // Multiply the coefficients by the coordinates.
   // [[0, 1], [0, 2],  [0, 3]] X [4, 1] -> [[4*0, 1*1], [4*0, 1*2], [4*0, 1*3]]
   // %13 = "tosa.mul"(%11, %12) {shift = 0 : i32} : (tensor<3x2xi32>,
   // tensor<2xi32>) -> tensor<3x2xi32>
   auto flattenedIndicesMulOp = tosa::CreateOpAndInfer<tosa::MulOp>(
       rewriter, op->getLoc(),
       GetTypeFromTensorShape(indicesMatrixShape, indicesType.getElementType()),
-      indicesMatrixReshapeOp.getResult(), flattenedCoeffValue.value(), 0);
+      indicesMatrixReshapeOp, flattenedCoeffValue.value(), 0);
 
   // Sum up the products of the coefficients and coordinates
   // [[4*0 + 1*1], [4*0 + 1*2], [4*0 + 1*3]] = [[1],[2],[3]]
@@ -734,10 +744,20 @@ std::optional<Value> convertReduceOpCommon(
       RankedTensorType reduce_type =
           RankedTensorType::get(shape_vec, reduce_element_type);
 
-      auto reduce_op = CreateOpAndInfer<T>(rewriter, op->getLoc(), reduce_type,
-                                           val, axis_attr);
+      Value reduce_op;
+      if constexpr (std::is_same<T, tosa::ReduceMinOp>() ||
+                    std::is_same<T, tosa::ReduceMaxOp>()) {
+        // Use default NaN Propagation mode "PROPAGATE" for tosa.reduce_min
+        // and tosa.reduce_max
+        reduce_op = CreateOpAndInfer<T>(
+            rewriter, op->getLoc(), reduce_type, val, axis_attr,
+            /*nan_mode=*/rewriter.getStringAttr("PROPAGATE"));
+      } else {
+        reduce_op = CreateOpAndInfer<T>(rewriter, op->getLoc(), reduce_type,
+                                        val, axis_attr);
+      }
 
-      val = reduce_op.getResult();
+      val = reduce_op;
     }
 
     if (is_quantized) {
@@ -973,6 +993,12 @@ convertReduceMeanOp(PatternRewriter &rewriter, Operation *op,
 
   if (!input_is_qtype) {
     Value div_const = getTosaConstTensorSingleF32(rewriter, op, div_scale);
+
+    if (mlir::tosa::EqualizeRanks(rewriter, op->getLoc(), val.value(),
+                                  div_const)
+            .failed())
+      return std::nullopt;
+
     return CreateOpAndInfer<tosa::MulOp>(rewriter, op->getLoc(), output_type,
                                          val.value(), div_const, 0)
         .getResult();
@@ -1021,6 +1047,11 @@ convertLinalgVectorNormOp(PatternRewriter &rewriter, Operation *op,
     return std::nullopt;
   }
 
+  Value ordValRank0 = ordVal;
+  if (mlir::tosa::EqualizeRanks(rewriter, op->getLoc(), input_value, ordVal)
+          .failed())
+    return std::nullopt;
+
   if (fabs(ordLiteralFloat) < epsilon ||
       fabs(static_cast<double>(ordLiteralInt)) < epsilon) {
     op->emitOpError("unimplemented: L0 norm");
@@ -1049,9 +1080,17 @@ convertLinalgVectorNormOp(PatternRewriter &rewriter, Operation *op,
       rewriter, op, output_type, powVal, axes_elems, keep_dims);
   if (!result)
     return std::nullopt;
-  auto reciprocalVal = CreateOpAndInfer<tosa::ReciprocalOp>(
-                           rewriter, op->getLoc(), ordVal.getType(), ordVal)
-                           .getResult();
+
+  Value reciprocalVal =
+      CreateOpAndInfer<tosa::ReciprocalOp>(rewriter, op->getLoc(),
+                                           ordValRank0.getType(), ordValRank0)
+          .getResult();
+
+  if (mlir::tosa::EqualizeRanks(rewriter, op->getLoc(), result.value(),
+                                reciprocalVal)
+          .failed())
+    return std::nullopt;
+
   return CreateOpAndInfer<tosa::PowOp>(rewriter, op->getLoc(), output_type,
                                        result.value(), reciprocalVal)
       .getResult();

lib/Conversion/TorchToTosa/TosaLegalizeUtils.cpp

@@ -8,7 +8,8 @@
 //===----------------------------------------------------------------------===//
 
 #include "torch-mlir/Conversion/TorchToTosa/TosaLegalizeUtils.h"
-#include "mlir/Dialect/Tosa/IR/TosaOps.h"       // from @llvm-project
+#include "mlir/Dialect/Tosa/IR/TosaOps.h" // from @llvm-project
+#include "mlir/Dialect/Tosa/Utils/ConversionUtils.h"
 #include "mlir/Dialect/Tosa/Utils/QuantUtils.h" // from @llvm-project
 
 namespace mlir {
@@ -301,31 +302,31 @@ std::optional<Value> getConstTensor<float>(PatternRewriter &rewriter,
       (src.isF32() && dest.isInteger(8)) ||
       (src.isF32() && dest.isBF16()) ||
       (src.isF32() && dest.isF16()) ||
-      (src.isF32() && dest.isFloat8E4M3()) ||
-      (src.isF32() && dest.isFloat8E5M2()) ||
+      (src.isF32() && isa<Float8E4M3Type>(dest)) ||
+      (src.isF32() && isa<Float8E5M2Type>(dest)) ||
       // f16 -> *
       (src.isF16() && dest.isInteger(32)) ||
       (src.isF16() && dest.isInteger(16)) ||
       (src.isF16() && dest.isInteger(8)) ||
       (src.isF16() && dest.isBF16()) ||
       (src.isF16() && dest.isF32()) ||
-      (src.isF16() && dest.isFloat8E4M3()) ||
-      (src.isF16() && dest.isFloat8E5M2()) ||
+      (src.isF16() && isa<Float8E4M3Type>(dest)) ||
+      (src.isF16() && isa<Float8E5M2Type>(dest)) ||
       // bf16 -> *
       (src.isBF16() && dest.isInteger(32)) ||
       (src.isBF16() && dest.isInteger(16)) ||
       (src.isBF16() && dest.isInteger(8)) ||
       (src.isBF16() && dest.isF32()) ||
-      (src.isBF16() && dest.isFloat8E4M3()) ||
-      (src.isBF16() && dest.isFloat8E5M2()) ||
+      (src.isBF16() && isa<Float8E4M3Type>(dest)) ||
+      (src.isBF16() && isa<Float8E5M2Type>(dest)) ||
       // fp8e4m3 -> *
-      (src.isFloat8E4M3() && dest.isBF16()) ||
-      (src.isFloat8E4M3() && dest.isF32()) ||
-      (src.isFloat8E4M3() && dest.isF16()) ||
+      (isa<Float8E4M3Type>(src) && dest.isBF16()) ||
+      (isa<Float8E4M3Type>(src) && dest.isF32()) ||
+      (isa<Float8E4M3Type>(src) && dest.isF16()) ||
       // fp8e5m2 -> *
-      (src.isFloat8E5M2() && dest.isBF16()) ||
-      (src.isFloat8E5M2() && dest.isF32()) ||
-      (src.isFloat8E5M2() && dest.isF16())) {
+      (isa<Float8E5M2Type>(src) && dest.isBF16()) ||
+      (isa<Float8E5M2Type>(src) && dest.isF32()) ||
+      (isa<Float8E5M2Type>(src) && dest.isF16())) {
     return success();
   }
   // clang-format on
@@ -393,6 +394,11 @@ LogicalResult tosaCastTensorToType(PatternRewriter &rewriter, Operation *op,
       auto zeroValue =
           tosa::getConstTensor<float>(rewriter, op, 0, {}, srcElemTy).value();
 
+      if (mlir::tosa::EqualizeRanks(rewriter, op->getLoc(), src, zeroValue)
+              .failed())
+        return rewriter.notifyMatchFailure(
+            op, "Failed to equalize ranks among operands and result");
+
       auto boolType = srcType.clone(rewriter.getIntegerType(1));
       auto isNegative = tosa::CreateOpAndInfer<tosa::GreaterOp>(
           rewriter, op->getLoc(), boolType, zeroValue, src);
@@ -488,10 +494,10 @@ LogicalResult getConvOpsAccType(PatternRewriter &rewriter,
   } else if (inputElemTy.isInteger(16) && weightElemTy.isInteger(8) &&
              outputElemTy.isInteger(48)) {
     accType = mlir::TypeAttr::get(rewriter.getIntegerType(48));
-  } else if ((inputElemTy.isFloat8E4M3() && weightElemTy.isFloat8E4M3() &&
-              outputElemTy.isF16()) ||
-             (inputElemTy.isFloat8E5M2() && weightElemTy.isFloat8E5M2() &&
-              outputElemTy.isF16())) {
+  } else if ((isa<Float8E4M3Type>(inputElemTy) &&
+              isa<Float8E4M3Type>(weightElemTy) && outputElemTy.isF16()) ||
+             (isa<Float8E5M2Type>(inputElemTy) &&
+              isa<Float8E5M2Type>(weightElemTy) && outputElemTy.isF16())) {
     accType = mlir::TypeAttr::get(rewriter.getF16Type());
   } else {
     accType = mlir::TypeAttr::get(outputElemTy);
@@ -500,17 +506,5 @@ LogicalResult getConvOpsAccType(PatternRewriter &rewriter,
   return success();
 }
 
-// Temporary function to get TOSA const shape
-// TODO: Remove this function when getTosaConstShape is available in
-// externals/llvm-project/mlir/include/mlir/Dialect/Tosa/Utils/ConversionUtils.h
-Value getTosaConstShape(PatternRewriter &rewriter, Location loc,
-                        llvm::ArrayRef<int64_t> shape) {
-  auto attr = rewriter.getIndexTensorAttr(shape);
-  auto type = mlir::tosa::shapeType::get(rewriter.getContext(), shape.size());
-  mlir::Operation *mlir_op =
-      rewriter.create<tosa::ConstShapeOp>(loc, type, attr);
-  return mlir_op->getResult(0);
-}
-
 } // namespace tosa
 } // namespace mlir

lib/Dialect/Torch/Utils/Utils.cpp

@@ -652,13 +652,13 @@ Type Torch::getDefaultAccType(PatternRewriter &rewriter, Type inputType) {
     return rewriter.getF32Type();
   if (isa<Float64Type>(inputType))
     return rewriter.getF64Type();
-  if (inputType.isFloat8E5M2())
+  if (isa<Float8E5M2Type>(inputType))
     return rewriter.getF32Type();
-  if (inputType.isFloat8E4M3FN())
+  if (isa<Float8E4M3FNType>(inputType))
     return rewriter.getF32Type();
-  if (inputType.isFloat8E5M2FNUZ())
+  if (isa<Float8E5M2FNUZType>(inputType))
     return rewriter.getF32Type();
-  if (inputType.isFloat8E4M3FNUZ())
+  if (isa<Float8E4M3FNUZType>(inputType))
     return rewriter.getF32Type();
   if (inputType.isInteger(8))
     // this is an intentional deviation from CUDA (which accumulates i8 to i64)