onnx.MelWeightMatrix TorchOnnxToTorch (llvm#3503)

Just uploading what I have till now

[Gist](https://gist.github.com/PhaneeshB/761f75f5522d9f4a40ef949a328e93fe)
of pytorch impl that I'm following to implement the OnnxToTorch lowering

Additional Details - (also pasted as comment in gist)
[Op
Description](https://github.com/onnx/onnx/blob/main/docs/Operators.md#melweightmatrix)
in Onnx Documentation

[Example](https://github.com/onnx/onnx/blob/main/docs/Operators.md#examples-93)
Used the same example in this file.
the Expected output is shown in the example

[Reference Onnx
Impl](https://github.com/onnx/onnx/blob/4c3ed5e08be75bbe1eeb6818e490b1b6a370183e/onnx/reference/ops/op_mel_weight_matrix.py#L13)
- This is the base for the above code.

Author: PhaneeshB

lib/Conversion/TorchOnnxToTorch/DefaultDomainGtoP.cpp

@@ -591,6 +591,373 @@ void mlir::torch::onnx_c::populateDefaultDomainGtoP(
                       binder.op, resultType, lhs, rhs);
                   return success();
                 });
+
+  patterns.onOp(
+      "MelWeightMatrix", 17,
+      [](OpBinder binder, ConversionPatternRewriter &rewriter) {
+        llvm::SmallVector<Value> operands;
+        Torch::ValueTensorType resultType;
+        int64_t output_dtype_attr;
+        if (binder.tensorOperands(operands, 5) ||
+            binder.tensorResultType(resultType) || operands.size() != 5 ||
+            binder.s64IntegerAttr(output_dtype_attr, "output_datatype", 1)) {
+          return failure();
+        }
+        // operands sequence :
+        // num_mel_bins, dft_length, sample_rate -> int32/64 tensors
+        // lower_edge_hertz, upper_edge_hertz -> f16/32/64
+
+        // Need to backtrack the values of num_mel_bins and dft_length//2+1 from
+        // result shape since the inputs are tensors and we cannot know their
+        // values at compile time. if the result type does not contain static
+        // shapes, then the implementation will be unsupported.
+        if (!resultType.areAllSizesKnown())
+          return rewriter.notifyMatchFailure(
+              binder.op, "Unknown result sizes, not supported.");
+
+        ArrayRef<int64_t> resShape = resultType.getSizes();
+        if (resShape.size() != 2)
+          return rewriter.notifyMatchFailure(
+              binder.op,
+              "Expected result rank to be 2, not supported for other ranks.");
+
+        std::optional<int64_t> torchDTypeInt =
+            onnxDtypeIntToTorchDtypeInt(output_dtype_attr);
+        if (!torchDTypeInt.has_value()) {
+          return rewriter.notifyMatchFailure(
+              binder.op, "conversion to given output dtype unsupported");
+        }
+
+        // Here Onwards all shapes will be computed using these sizes
+        int64_t numSpectrogramBinsInt = resShape[0];
+        int64_t numMelBinsInt = resShape[1];
+        Torch::ValueTensorType inputIntType = binder.toValidTensorType(
+            operands[0].getType()); // Since operands[0 / 1 / 2] will have the
+                                    // same int type.
+        Torch::ValueTensorType inputFloatType = binder.toValidTensorType(
+            operands[3].getType()); // Since operands[3 / 4] will have the same
+                                    // float type
+
+        Value numMelBinsItem =
+            getItemOp<Torch::IntType>(binder, rewriter, operands[0]);
+        Value dftLengthItem =
+            getItemOp<Torch::IntType>(binder, rewriter, operands[1]);
+        Value sampleRateItem =
+            getItemOp<Torch::IntType>(binder, rewriter, operands[2]);
+        Value lowerEdgeHzItem =
+            getItemOp<Torch::FloatType>(binder, rewriter, operands[3]);
+        Value upperEdgeHzItem =
+            getItemOp<Torch::FloatType>(binder, rewriter, operands[4]);
+
+        // Helpers
+        ImplicitLocOpBuilder b(binder.getLoc(), rewriter);
+        auto ctx = binder.op->getContext();
+
+        // Recurring shapes
+        SmallVector<int64_t> unranked({});
+        SmallVector<int64_t> shapeNMB({numMelBinsInt});
+        SmallVector<int64_t> shapeNMBp2({numMelBinsInt + 2});
+        SmallVector<int64_t> shape1xNMB({1, numMelBinsInt});
+        SmallVector<int64_t> shapeNSB({numSpectrogramBinsInt});
+        SmallVector<int64_t> shapeNSBxNMB(
+            {numSpectrogramBinsInt, numMelBinsInt});
+
+        // Recurring DTypes
+        Type inpFpDType = inputFloatType.getDtype();
+        Type inpIntDType = inputIntType.getDtype();
+        Type si32Ty = rewriter.getIntegerType(32, true);
+        Type f32Ty = rewriter.getF32Type();
+        Type i1Ty = rewriter.getI1Type();
+
+        // Value constants
+        Value noneConst = b.create<Torch::ConstantNoneOp>();
+        Value negTwoConst =
+            b.create<Torch::ConstantIntOp>(rewriter.getI64IntegerAttr(-2));
+        Value negOneConst =
+            b.create<Torch::ConstantIntOp>(rewriter.getI64IntegerAttr(-1));
+        Value zeroConst =
+            b.create<Torch::ConstantIntOp>(rewriter.getI64IntegerAttr(0));
+        Value oneConst =
+            b.create<Torch::ConstantIntOp>(rewriter.getI64IntegerAttr(1));
+        Value twoConst =
+            b.create<Torch::ConstantIntOp>(rewriter.getI64IntegerAttr(2));
+        Value float32DTypeConst =
+            b.create<Torch::ConstantIntOp>(rewriter.getI64IntegerAttr(6));
+
+        Torch::ValueTensorType dftLenType =
+            Torch::ValueTensorType::get(ctx, unranked, inpIntDType);
+        Type freqBinsIntType =
+            Torch::ValueTensorType::get(ctx, shapeNMBp2, si32Ty);
+        Type freqBinsFltType =
+            Torch::ValueTensorType::get(ctx, shapeNMBp2, f32Ty);
+
+        Value dftLengthDivTwoFlt =
+            b.create<Torch::AtenDivIntOp>(dftLengthItem, twoConst);
+        Value dftLengthDivTwo =
+            b.create<Torch::AtenIntFloatOp>(dftLengthDivTwoFlt);
+        Value numSpectrogramBins =
+            b.create<Torch::AtenAddIntOp>(dftLengthDivTwo, oneConst);
+        Value numSpectrogramBinsItem = numSpectrogramBins;
+        Value freqBinsInit = b.create<Torch::AtenArangeOp>(
+            freqBinsIntType, numMelBinsItem, /*dtype=*/float32DTypeConst,
+            /*layout=*/noneConst, /*device=*/noneConst,
+            /*pin_memory=*/noneConst);
+
+        // From Ref Impl of Onnx.MelWeightMatrix:
+        // https://github.com/onnx/onnx/blob/main/onnx/reference/ops/op_mel_weight_matrix.py#L25-L32
+        // convert input Freq Hz to Mel
+        Value twoFiveNineFiveConst =
+            b.create<Torch::ConstantFloatOp>(rewriter.getF64FloatAttr(2595));
+        Value sevenHConst =
+            b.create<Torch::ConstantFloatOp>(rewriter.getF64FloatAttr(700));
+        Value tenConst =
+            b.create<Torch::ConstantFloatOp>(rewriter.getF64FloatAttr(10));
+
+        Value lfDiv7Hfloat =
+            b.create<Torch::AtenDivFloatOp>(lowerEdgeHzItem, sevenHConst);
+        Type freqType = Torch::ValueTensorType::get(ctx, unranked, inpFpDType);
+        Value lfDiv7H =
+            b.create<Torch::PrimNumToTensorScalarOp>(freqType, lfDiv7Hfloat);
+        Value lfDiv7HAdd1 = b.create<Torch::AtenAddScalarOp>(
+            freqType, lfDiv7H, oneConst, /*alpha =*/oneConst);
+        Value lfDiv7HAdd1Log10 =
+            b.create<Torch::AtenLog10Op>(freqType, lfDiv7HAdd1);
+        Value lfMel = b.create<Torch::AtenMulScalarOp>(
+            freqType, lfDiv7HAdd1Log10, twoFiveNineFiveConst);
+
+        Value hfDiv7Hfloat =
+            b.create<Torch::AtenDivFloatOp>(upperEdgeHzItem, sevenHConst);
+        Value hfDiv7H =
+            b.create<Torch::PrimNumToTensorScalarOp>(freqType, hfDiv7Hfloat);
+        Value hfDiv7HAdd1 = b.create<Torch::AtenAddScalarOp>(
+            freqType, hfDiv7H, oneConst, /*alpha =*/oneConst);
+        Value hfDiv7HAdd1Log10 =
+            b.create<Torch::AtenLog10Op>(freqType, hfDiv7HAdd1);
+        Value hfMel = b.create<Torch::AtenMulScalarOp>(
+            freqType, hfDiv7HAdd1Log10, twoFiveNineFiveConst);
+
+        Value hfSubLf = b.create<Torch::AtenSubTensorOp>(
+            hfMel.getType(), hfMel, lfMel, /*alpha=*/oneConst);
+        Value melStep = b.create<Torch::AtenDivScalarOp>(
+            hfSubLf.getType(), hfSubLf, numMelBinsItem);
+
+        Value freqBinsMulMelStep = b.create<Torch::AtenMulTensorOp>(
+            freqBinsFltType, freqBinsInit, melStep);
+        Value freqBinsScaled = b.create<Torch::AtenAddTensorOp>(
+            freqBinsFltType, freqBinsMulMelStep, lfMel, /*alpha=*/oneConst);
+
+        // Mel to Hz conv
+
+        Value fbDiv = b.create<Torch::AtenDivScalarOp>(
+            freqBinsFltType, freqBinsScaled, twoFiveNineFiveConst);
+        Value fbClone = b.create<Torch::AtenCloneOp>(
+            freqBinsFltType, freqBinsScaled, /*memory_format=*/noneConst);
+        Value tenTensor = b.create<Torch::AtenFillScalarOp>(freqBinsFltType,
+                                                            fbClone, tenConst);
+        Value fbPow = b.create<Torch::AtenPowTensorTensorOp>(freqBinsFltType,
+                                                             tenTensor, fbDiv);
+        Value fbPowSubOne = b.create<Torch::AtenSubScalarOp>(
+            freqBinsFltType, fbPow, oneConst, /*alpha=*/oneConst);
+        Value freqBinsHz = b.create<Torch::AtenMulScalarOp>(
+            freqBinsFltType, fbPowSubOne, sevenHConst);
+
+        // Normalize freqBinsHz
+        Value dftLenPlusOne = b.create<Torch::AtenAddScalarOp>(
+            dftLenType, operands[1], oneConst, /*alpha=*/oneConst);
+        Value dftLenPlusOneItem =
+            getItemOp<Torch::IntType>(binder, rewriter, dftLenPlusOne);
+        Value fbMulDft = b.create<Torch::AtenMulScalarOp>(
+            freqBinsFltType, freqBinsHz, dftLenPlusOneItem);
+        Value freqBinsNormalized = b.create<Torch::AtenDivScalarOp>(
+            freqBinsFltType, fbMulDft, sampleRateItem);
+
+        // cast to int32
+        Value int32DTypeConst =
+            b.create<Torch::ConstantIntOp>(rewriter.getI64IntegerAttr(3));
+        Value falseConst = b.create<Torch::ConstantBoolOp>(false);
+        Value freqBins = b.create<Torch::AtenToDtypeOp>(
+            freqBinsIntType, freqBinsNormalized, /*dtype=*/int32DTypeConst,
+            /*non_blocking=*/falseConst, /*copy=*/falseConst,
+            /*memory_format=*/noneConst);
+
+        Torch::ValueTensorType sliceResType =
+            Torch::ValueTensorType::get(ctx, shapeNMB, si32Ty);
+        Type unsqueezeResType =
+            sliceResType.getWithSizesAndDtype(shape1xNMB, si32Ty);
+        Value lfTensor = b.create<Torch::AtenSliceTensorOp>(
+            sliceResType, freqBins, /*dim=*/zeroConst, /*start=*/zeroConst,
+            /*end=*/negTwoConst, /*step=*/oneConst);
+        Value lowFreqTensor = b.create<Torch::AtenUnsqueezeOp>(
+            unsqueezeResType, lfTensor, /*dim=*/zeroConst);
+
+        Value cfTensor = b.create<Torch::AtenSliceTensorOp>(
+            sliceResType, freqBins, /*dim=*/zeroConst, /*start=*/oneConst,
+            /*end=*/negOneConst, /*step=*/oneConst);
+        Value centerFreqTensor = b.create<Torch::AtenUnsqueezeOp>(
+            unsqueezeResType, cfTensor, /*dim=*/zeroConst);
+
+        Value hfTensor = b.create<Torch::AtenSliceTensorOp>(
+            sliceResType, freqBins, /*dim=*/zeroConst, /*start=*/zeroConst,
+            /*end=*/noneConst, /*step=*/oneConst);
+        Value highFreqTensor = b.create<Torch::AtenUnsqueezeOp>(
+            unsqueezeResType, hfTensor, /*dim=*/zeroConst);
+
+        Value lowToCenter =
+            b.create<Torch::AtenSubTensorOp>(unsqueezeResType, centerFreqTensor,
+                                             lowFreqTensor, /*alpha=*/oneConst);
+        Value centerToHigh = b.create<Torch::AtenSubTensorOp>(
+            unsqueezeResType, highFreqTensor, centerFreqTensor,
+            /*alpha=*/oneConst);
+
+        Type zeroToNInitType =
+            inputIntType.getWithSizesAndDtype(shapeNSB, f32Ty);
+        Value zeroToNInit = b.create<Torch::AtenArangeOp>(
+            zeroToNInitType, numSpectrogramBinsItem,
+            /*dtype=*/float32DTypeConst,
+            /*layout=*/noneConst, /*device=*/noneConst,
+            /*pin_memory=*/noneConst);
+
+        Type zeroToNBaseType = inputIntType.getWithSizesAndDtype(
+            ArrayRef<int64_t>{numSpectrogramBinsInt, 1}, f32Ty);
+        Value zeroToNBase = b.create<Torch::AtenUnsqueezeOp>(
+            zeroToNBaseType, zeroToNInit, /*dim=*/oneConst);
+        Type zeroToNumElesType =
+            inputIntType.getWithSizesAndDtype(shapeNSBxNMB, f32Ty);
+        Value expandShapeList = b.create<Torch::PrimListConstructOp>(
+            rewriter.getType<Torch::ListType>(
+                rewriter.getType<Torch::IntType>()),
+            SmallVector<Value>{numSpectrogramBinsItem, numMelBinsItem});
+        Value zeroToNumEles = b.create<Torch::AtenExpandOp>(
+            zeroToNumElesType, zeroToNBase, expandShapeList,
+            /*implicit=*/falseConst);
+
+        Type maskType = inputIntType.getWithSizesAndDtype(shape1xNMB, i1Ty);
+        Value maskLowToCenterZero =
+            b.create<Torch::AtenEqScalarOp>(maskType, lowToCenter, zeroConst);
+
+        // L2C computation
+        Value lowToCenterNoZero = b.create<Torch::AtenWhereScalarSelfOp>(
+            unsqueezeResType, maskLowToCenterZero, negOneConst, lowToCenter);
+        Type maskL2CAfterCType =
+            inputIntType.getWithSizesAndDtype(shapeNSBxNMB, i1Ty);
+        Value maskL2CAfterC = b.create<Torch::AtenGtTensorOp>(
+            maskL2CAfterCType, zeroToNumEles, centerFreqTensor);
+        Type maxLFResTy =
+            inputIntType.getWithSizesAndDtype(ArrayRef<int64_t>{1}, si32Ty);
+        Value maxLowerFreq =
+            b.create<Torch::AtenMaxOp>(maxLFResTy, lowFreqTensor);
+        Value maxLowerFreqItem =
+            getItemOp<Torch::IntType>(binder, rewriter, maxLowerFreq);
+        Value zeroToNumElesL2C = b.create<Torch::AtenWhereScalarSelfOp>(
+            zeroToNumElesType, maskLowToCenterZero, maxLowerFreqItem,
+            zeroToNumEles);
+        Value upslopeDiff = b.create<Torch::AtenSubTensorOp>(
+            zeroToNumElesType, zeroToNumElesL2C, lowFreqTensor,
+            /*alpha=*/oneConst);
+        Type l2cNZFltTy = inputIntType.getWithSizesAndDtype(shape1xNMB, f32Ty);
+        Value l2cNZFlt = b.create<Torch::AtenToDtypeOp>(
+            l2cNZFltTy, lowToCenterNoZero, /*dtype=*/float32DTypeConst,
+            /*non_blocking=*/falseConst, /*copy=*/falseConst,
+            /*memory_format=*/noneConst);
+        Value upslopeL2C0 = b.create<Torch::AtenDivTensorOp>(
+            zeroToNumElesType, upslopeDiff, l2cNZFlt);
+        Type maskUpslopeL2C0PosType =
+            inputIntType.getWithSizesAndDtype(shapeNSBxNMB, i1Ty);
+        Value maskUpslopeL2C0Pos = b.create<Torch::AtenGtScalarOp>(
+            maskUpslopeL2C0PosType, upslopeL2C0, zeroConst);
+        Value upslopeL2C0PosRanged = b.create<Torch::AtenWhereScalarOtherOp>(
+            zeroToNumElesType, maskUpslopeL2C0Pos, upslopeL2C0, zeroConst);
+        Value maskIdxL2CAfterCList = b.create<Torch::PrimListConstructOp>(
+            rewriter.getType<Torch::ListType>(maskL2CAfterC.getType()),
+            ValueRange{maskL2CAfterC});
+        Value zeroConstTensor = Torch::createRank0Tensor(
+            rewriter, binder.getLoc(),
+            Torch::ValueTensorType::get(ctx, std::nullopt, f32Ty), zeroConst);
+        Value upslopeL2C1 = b.create<Torch::AtenIndexPutOp>(
+            zeroToNumElesType, upslopeL2C0PosRanged, maskIdxL2CAfterCList,
+            zeroConstTensor, falseConst);
+        Value maskIdxL2CZeroList = b.create<Torch::PrimListConstructOp>(
+            rewriter.getType<Torch::ListType>(maskLowToCenterZero.getType()),
+            ValueRange{maskLowToCenterZero});
+        Type centerFreqTensorL2CZeroType =
+            inputIntType.getWithSizesAndDtype(ArrayRef<int64_t>{-1}, si32Ty);
+        Value centerFreqTensorL2CZero = b.create<Torch::AtenIndexTensorOp>(
+            centerFreqTensorL2CZeroType, centerFreqTensor, maskIdxL2CZeroList);
+        Type maskSqueezeType =
+            inputIntType.getWithSizesAndDtype(shapeNMB, i1Ty);
+        Value maskLowToCenterZeroSqueeze = b.create<Torch::AtenSqueezeOp>(
+            maskSqueezeType, maskLowToCenterZero);
+        Type maskL2CIntTy = inputIntType.getWithSizesAndDtype(shapeNMB, si32Ty);
+        Value maskLowToCenterInt = b.create<Torch::AtenToDtypeOp>(
+            maskL2CIntTy, maskLowToCenterZeroSqueeze, /*dtype=*/int32DTypeConst,
+            /*non_blocking=*/falseConst, /*copy=*/falseConst,
+            /*memory_format=*/noneConst);
+        Value upslopeOneIdxList = b.create<Torch::PrimListConstructOp>(
+            rewriter.getType<Torch::ListType>(
+                centerFreqTensorL2CZero.getType()),
+            ValueRange{centerFreqTensorL2CZero, maskLowToCenterInt});
+        Value oneConstTensor = Torch::createRank0Tensor(
+            rewriter, binder.getLoc(),
+            Torch::ValueTensorType::get(ctx, std::nullopt, f32Ty), oneConst);
+        Value upslopeL2C = b.create<Torch::AtenIndexPutOp>(
+            zeroToNumElesType, upslopeL2C1, upslopeOneIdxList, oneConstTensor,
+            falseConst);
+
+        // H2C computation
+        Value maskCenterToHighZero =
+            b.create<Torch::AtenEqScalarOp>(maskType, centerToHigh, zeroConst);
+        Value maskH2CBeforeC = b.create<Torch::AtenLtTensorOp>(
+            maskL2CAfterCType, zeroToNumEles, centerFreqTensor);
+        Value centerToHighNoZero = b.create<Torch::AtenWhereScalarSelfOp>(
+            unsqueezeResType, maskCenterToHighZero, negOneConst, centerToHigh);
+        Value c2hNZFlt = b.create<Torch::AtenToDtypeOp>(
+            l2cNZFltTy, centerToHighNoZero, /*dtype=*/float32DTypeConst,
+            /*non_blocking=*/falseConst, /*copy=*/falseConst,
+            /*memory_format=*/noneConst);
+        Value zeroToNumElesC2H = b.create<Torch::AtenWhereScalarSelfOp>(
+            zeroToNumElesType, maskCenterToHighZero, zeroConst, zeroToNumEles);
+        Value downslopeDiff = b.create<Torch::AtenSubTensorOp>(
+            zeroToNumElesType, highFreqTensor, zeroToNumElesC2H,
+            /*alpha=*/oneConst);
+        Value downslopeC2H0 = b.create<Torch::AtenDivTensorOp>(
+            zeroToNumElesType, downslopeDiff, c2hNZFlt);
+        Value maskDownslopeC2H0Pos = b.create<Torch::AtenGtScalarOp>(
+            maskUpslopeL2C0PosType, downslopeC2H0, zeroConst);
+        Value downslopeC2H0Pos = b.create<Torch::AtenWhereScalarOtherOp>(
+            zeroToNumElesType, maskDownslopeC2H0Pos, downslopeC2H0, zeroConst);
+        Value idxH2CBeforeCList = b.create<Torch::PrimListConstructOp>(
+            rewriter.getType<Torch::ListType>(maskH2CBeforeC.getType()),
+            ValueRange{maskH2CBeforeC});
+        Value downslopeC2H = b.create<Torch::AtenIndexPutOp>(
+            zeroToNumElesType, downslopeC2H0Pos, idxH2CBeforeCList,
+            zeroConstTensor, falseConst);
+
+        // final result Calculation
+        Value maskH2CNonZero = b.create<Torch::AtenNeScalarOp>(
+            maskL2CAfterCType, downslopeC2H, zeroConst);
+        Value idxH2CNZList = b.create<Torch::PrimListConstructOp>(
+            rewriter.getType<Torch::ListType>(maskH2CNonZero.getType()),
+            ValueRange{maskH2CNonZero});
+        Value upslopeL2CMasked = b.create<Torch::AtenIndexPutOp>(
+            zeroToNumElesType, upslopeL2C, idxH2CNZList, zeroConstTensor,
+            falseConst);
+
+        Value slopesFinal = b.create<Torch::AtenAddTensorOp>(
+            zeroToNumElesType, upslopeL2CMasked, downslopeC2H,
+            /*alpha=*/oneConst);
+
+        Value outputDTypeConst = b.create<Torch::ConstantIntOp>(
+            rewriter.getType<Torch::IntType>(),
+            rewriter.getI64IntegerAttr(torchDTypeInt.value()));
+        Value finalOutput = b.create<Torch::AtenToDtypeOp>(
+            resultType, slopesFinal, /*dtype=*/outputDTypeConst,
+            /*non_blocking=*/falseConst, /*copy=*/falseConst,
+            /*memory_format=*/noneConst);
+
+        rewriter.replaceOp(binder.op, finalOutput);
+        return success();
+      });
+
   patterns.onOp(
       "Multinomial", 7,
       [](OpBinder binder, ConversionPatternRewriter &rewriter) {