torch-mlir change for dense resource implementation (llvm#2513)

Co-authored-by: Avinash Sharma <avinash@nod-labs.com>
ashwinpn · Nov 3, 2023 · 88adf38 · 88adf38
1 parent 1b9fb1b
commit 88adf38
Showing 1 changed file with 14 additions and 9 deletions.
diff --git a/lib/Conversion/TorchToArith/TorchToArith.cpp b/lib/Conversion/TorchToArith/TorchToArith.cpp
@@ -171,14 +171,18 @@ class ConvertTorchTensorLiteralOp
                   ConversionPatternRewriter &rewriter) const override {
     MLIRContext *context = op->getContext();
     if (auto elements = op.getValueAttr().dyn_cast<DenseIntElementsAttr>()) {
-      Type elemTy = op.getValueAttr().getElementType();
-      unsigned bitWidth = elemTy.getIntOrFloatBitWidth();
-      Type builtinTensorElemTy = IntegerType::get(context, bitWidth);
-      rewriter.replaceOpWithNewOp<arith::ConstantOp>(
-          op, elements.mapValues(builtinTensorElemTy, [&](const APInt &v) {
-            return APInt(bitWidth, v.getSExtValue());
-          }));
-      return success();
+      if (auto type = elements.getType().dyn_cast<RankedTensorType>()) {
+        Type elemTy = op.getValueAttr().getElementType();
+        unsigned bitWidth = elemTy.getIntOrFloatBitWidth();
+        Type builtinTensorElemTy = IntegerType::get(context, bitWidth);
+        auto shapedType =
+              RankedTensorType::get(type.getShape(), builtinTensorElemTy);
+        auto rawData = elements.getRawData();
+        DenseElementsAttr newAttr = DenseElementsAttr::getFromRawBuffer(
+          shapedType, rawData);
+        rewriter.replaceOpWithNewOp<arith::ConstantOp>(op, newAttr);
+        return success();
+      }
     }
     if (auto elements = op.getValueAttr().dyn_cast<DenseResourceElementsAttr>()) {
       if (auto type = elements.getType().dyn_cast<RankedTensorType>()) {
@@ -190,7 +194,8 @@ class ConvertTorchTensorLiteralOp
           AsmResourceBlob *blob = elements.getRawHandle().getBlob();
           assert(blob && "Expecting dense resource with a valid blob");
           rewriter.replaceOpWithNewOp<arith::ConstantOp>(
-              op, DenseElementsAttr::get(shapedType, blob->getData()));
+              op, DenseResourceElementsAttr::get(shapedType,
+                                                 elements.getRawHandle()));
           return success();
         }
       }