Merge branch 'main' into tosa-clamp-f16-bf16

Author: justin-ngo-arm

externals/llvm-project

@@ -62,7 +62,9 @@ def TMTensor_ScanOp : TMTensor_Op<"scan",
 
   let builders = [
     OpBuilder<(ins "ValueRange":$inputs, "ValueRange":$outputs,
-      CArg<"int64_t", "0">:$dimension, CArg<"bool", "true">:$inclusive)>
+      CArg<"int64_t", "0">:$dimension, CArg<"bool", "true">:$inclusive), [{
+        build($_builder, $_state, TypeRange(outputs), inputs, outputs, dimension, inclusive);
+      }]>
   ];
 
   let results = (outs Variadic<AnyRankedTensor>:$results);
@@ -267,7 +269,9 @@ def TMTensor_AttentionOp : TMTensor_Op<"attention",
   );
 
   let builders = [
-    OpBuilder<(ins "ValueRange":$inputs, "ValueRange":$outputs)>
+    OpBuilder<(ins "ValueRange":$inputs, "ValueRange":$outputs), [{
+        build($_builder, $_state, TypeRange(outputs), inputs, outputs);
+    }]>
   ];
 
   let results = (outs Variadic<AnyRankedTensor>:$result);

externals/stablehlo

@@ -8668,6 +8668,30 @@ def Torch_AtenPixelShuffleOp : Torch_Op<"aten.pixel_shuffle", [
   }];
 }
 
+def Torch_AtenChannelShuffleOp : Torch_Op<"aten.channel_shuffle", [
+    AllowsTypeRefinement,
+    HasValueSemantics,
+    ReadOnly
+  ]> {
+  let summary = "Generated op for `aten::channel_shuffle : (Tensor, int) -> (Tensor)`";
+  let arguments = (ins
+    AnyTorchTensorType:$self,
+    Torch_IntType:$groups
+  );
+  let results = (outs
+    AnyTorchOptionalTensorType:$result
+  );
+  let hasCustomAssemblyFormat = 1;
+  let extraClassDefinition = [{
+    ParseResult AtenChannelShuffleOp::parse(OpAsmParser &parser, OperationState &result) {
+      return parseDefaultTorchOp(parser, result, 2, 1);
+    }
+    void AtenChannelShuffleOp::print(OpAsmPrinter &printer) {
+      printDefaultTorchOp(printer, *this, 2, 1);
+    }
+  }];
+}
+
 def Torch_AtenPermuteOp : Torch_Op<"aten.permute", [
     AllowsTypeRefinement,
     ReadOnly
@@ -14350,7 +14374,7 @@ def Torch_Aten_AssertTensorMetadataOp : Torch_Op<"aten._assert_tensor_metadata",
       printDefaultTorchOp(printer, *this, 6, 0);
     }
   }];
-  let hasFolder = 1;
+  let hasCanonicalizer = 1;
 }
 
 def Torch_AtenDiagonalOp : Torch_Op<"aten.diagonal", [

include/torch-mlir-dialects/Dialect/TMTensor/IR/TMTensorOps.td

@@ -1606,6 +1606,74 @@ void mlir::torch::onnx_c::populateDefaultDomainGtoP(
             /* cudnn enabled */ boolFalse);
         return success();
       });
+  patterns.onOp(
+      "MeanVarianceNormalization", 13,
+      [](OpBinder binder, ConversionPatternRewriter &rewriter) {
+        Torch::ValueTensorType resultType;
+        Value input;
+        SmallVector<int64_t> axes;
+
+        if (binder.tensorOperand(input) ||
+            binder.s64IntegerArrayAttr(axes, "axes",
+                                       llvm::SmallVector<int64_t>({0, 2, 3})) ||
+            binder.tensorResultType(resultType)) {
+          return failure();
+        }
+        if (!resultType.hasSizes() || !resultType.hasDtype()) {
+          return failure();
+        }
+        auto inputTy = cast<Torch::ValueTensorType>(input.getType());
+        if (!inputTy || !inputTy.hasSizes()) {
+          return failure();
+        }
+        int64_t inputRank = inputTy.getSizes().size();
+
+        Location loc = binder.getLoc();
+        Value keepDim = rewriter.create<Torch::ConstantBoolOp>(loc, true);
+        Value unBiased = rewriter.create<Torch::ConstantBoolOp>(loc, false);
+        Value none = rewriter.create<Torch::ConstantNoneOp>(loc);
+
+        ArrayRef<int64_t> output_shape = resultType.getSizes();
+        SmallVector<int64_t> reduced_shape(output_shape);
+
+        for (int64_t i : axes) {
+          int64_t dim = Torch::toPositiveDim(i, inputRank);
+          if (!Torch::isValidDim(dim, inputRank)) {
+            return failure();
+          }
+          reduced_shape[dim] = 1;
+        }
+        Torch::ValueTensorType reducedOutTy = Torch::ValueTensorType::get(
+            resultType.getContext(), reduced_shape, resultType.getDtype());
+        SmallVector<Value> cstAxes;
+        for (int64_t i : axes) {
+          cstAxes.push_back(rewriter.create<Torch::ConstantIntOp>(
+              loc, rewriter.getI64IntegerAttr(i)));
+        }
+        Value axes_list = rewriter.create<Torch::PrimListConstructOp>(
+            loc,
+            Torch::ListType::get(Torch::IntType::get(binder.op->getContext())),
+            cstAxes);
+        Value mean = rewriter.create<Torch::AtenMeanDimOp>(
+            loc, reducedOutTy, input, axes_list, keepDim, none);
+        Value variance = rewriter.create<Torch::AtenVarDimOp>(
+            loc, reducedOutTy, input, axes_list, unBiased, keepDim);
+        Value cstOne = rewriter.create<Torch::ConstantIntOp>(
+            loc, rewriter.getI64IntegerAttr(1));
+        Value cstEps = rewriter.create<Torch::ConstantFloatOp>(
+            loc, rewriter.getF64FloatAttr(1e-9));
+        variance = rewriter.create<Torch::AtenAddScalarOp>(
+            loc, reducedOutTy, variance, cstEps, cstOne);
+        Value sqrtVar =
+            rewriter.create<Torch::AtenSqrtOp>(loc, reducedOutTy, variance);
+        Value inputMinusMean = rewriter.create<Torch::AtenSubTensorOp>(
+            loc, resultType, input, mean, cstOne);
+        Value meanVarNorm = rewriter.create<Torch::AtenDivTensorOp>(
+            loc, resultType, inputMinusMean, sqrtVar);
+
+        rewriter.replaceOp(binder.op, meanVarNorm);
+        return success();
+      });
   patterns.onOp(
       "Max", 1, [](OpBinder binder, ConversionPatternRewriter &rewriter) {
         Torch::ValueTensorType resultType;