change to using llvm::SplitString for splitting the string

Author: Matthew Francis-Landau

mlir-tensorrt/tensorrt/lib/TensorRT/Transforms/TransposeReshapeElimination.cpp

@@ -28,6 +28,7 @@
 #include "mlir/IR/Matchers.h"
 #include "mlir/Pass/Pass.h"
 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
+#include "llvm/ADT/StringExtras.h"
 #include "llvm/Support/Debug.h"
 #include <numeric>
 
@@ -567,37 +568,31 @@ class RankChangeToReshape : public OpRewritePattern<OpType> {
 
 namespace {
 struct EinsumEquation {
-  std::string equation;
-  SmallVector<std::string> lhsParts;
-  std::string lhs;
-  std::string rhs;
+  llvm::StringRef equation;
+  SmallVector<llvm::SmallString<128>> lhsParts;
+  llvm::SmallString<128> lhs;
+  llvm::SmallString<128> rhs;
 
   LogicalResult parse(llvm::StringRef einsumEquation) {
-    std::string e{einsumEquation};
-    return parse(e);
-  }
-
-  LogicalResult parse(const std::string &einsumEquation) {
     size_t pos = einsumEquation.find("->");
     if (pos == std::string::npos)
       return failure();
     equation = einsumEquation;
-    lhs = einsumEquation.substr(0, pos);
-    rhs = einsumEquation.substr(pos + 2);
-    std::istringstream lhsStream(lhs);
-    std::string currentPart;
-    while (std::getline(lhsStream, currentPart, ',')) {
-      lhsParts.push_back(currentPart);
-    }
+    lhs = equation.substr(0, pos);
+    rhs = equation.substr(pos + 2);
+    SmallVector<llvm::StringRef> parts;
+    llvm::SplitString(lhs, parts, ",");
+    for (llvm::StringRef part : parts)
+      lhsParts.push_back(part); // cast from StringRef to SmallString
     return success();
   }
 
-  std::string generateEquation() const {
-    std::string ret = lhsParts[0];
+  StringRef generateEquation() const {
+    llvm::SmallString<128> ret = lhsParts[0];
     for (size_t i = 1; i < lhsParts.size(); i++) {
-      ret += "," + lhsParts[i];
+      ret.append({",", lhsParts[i]});
     }
-    ret += "->" + rhs;
+    ret.append({"->", rhs});
     return ret;
   }
 };
@@ -653,7 +648,7 @@ class PushDownTransposeToEinsum : public OpRewritePattern<tensorrt::EinsumOp> {
     if (!hasTransposeInput)
       return failure();
 
-    std::string newEinsumEquation = einsumEquation.generateEquation();
+    StringRef newEinsumEquation = einsumEquation.generateEquation();
 
     rewriter.replaceOpWithNewOp<tensorrt::EinsumOp>(op, op.getType(), newInputs,
                                                     newEinsumEquation);
@@ -696,7 +691,7 @@ class PushUpTransposeToEinsum : public OpRewritePattern<tensorrt::TransposeOp> {
     for (size_t i = 0; i < einsumRhs.size(); i++)
       einsumEquation.rhs += (char)einsumRhs[i];
 
-    std::string newEinsumEquation = einsumEquation.generateEquation();
+    StringRef newEinsumEquation = einsumEquation.generateEquation();
 
     auto newEinsum = rewriter.create<tensorrt::EinsumOp>(
         op.getLoc(), op.getType(), einsum.getInputs(), newEinsumEquation);
@@ -735,7 +730,7 @@ class EinsumPushDownTranspose : public OpRewritePattern<tensorrt::EinsumOp> {
     std::sort(outputAxes.begin(), outputAxes.end(),
               [&](const std::pair<char, int64_t> &a,
                   const std::pair<char, int64_t> &b) {
-                for (std::string &eqLhs : equation.lhsParts) {
+                for (auto &eqLhs : equation.lhsParts) {
                   if (eqLhs.find(a.first) != std::string::npos) {
                     if (eqLhs.find(b.first) != std::string::npos) {
                       return eqLhs.find(a.first) < eqLhs.find(b.first);
@@ -751,7 +746,7 @@ class EinsumPushDownTranspose : public OpRewritePattern<tensorrt::EinsumOp> {
 
     SmallVector<int64_t> newEinsumShape;
     SmallVector<int64_t> outputPerm;
-    std::string newEinsumRhs = "";
+    SmallString<128> newEinsumRhs{""};
     for (auto &[c, i] : outputAxes) {
       newEinsumRhs += c;
       newEinsumShape.push_back(op.getType().getDimSize(i));
@@ -760,7 +755,7 @@ class EinsumPushDownTranspose : public OpRewritePattern<tensorrt::EinsumOp> {
     if (newEinsumRhs == equation.rhs)
       return failure(); // no change
 
-    std::string newEinsumEquation = equation.generateEquation();
+    StringRef newEinsumEquation = equation.generateEquation();
 
     auto newEinsum = rewriter.create<tensorrt::EinsumOp>(
         op.getLoc(), op.getType().clone(newEinsumShape), op.getInputs(),
@@ -852,7 +847,7 @@ class EinsumPushUpTranspose : public OpRewritePattern<tensorrt::EinsumOp> {
     if (!didChange)
       return failure();
 
-    std::string newEquation = equation.generateEquation();
+    StringRef newEquation = equation.generateEquation();
     rewriter.replaceOpWithNewOp<tensorrt::EinsumOp>(op, op.getType(), newInputs,
                                                     newEquation);
     return success();
@@ -928,7 +923,7 @@ class EinsumEliminate1Axis : public OpRewritePattern<tensorrt::EinsumOp> {
         newOutputShape.push_back(outputType.getDimSize(i));
       }
     }
-    std::string newEquation = newEinsumEquation.generateEquation();
+    StringRef newEquation = newEinsumEquation.generateEquation();
 
     if (changeOutput) {
       auto newEinsum = rewriter.create<tensorrt::EinsumOp>(
@@ -1004,7 +999,7 @@ class EinsumMergeDown1Axis : public OpRewritePattern<tensorrt::EinsumOp> {
     if (!madeChange)
       return failure();
 
-    std::string newEquation = equation.generateEquation();
+    StringRef newEquation = equation.generateEquation();
 
     rewriter.replaceOpWithNewOp<tensorrt::EinsumOp>(op, op.getType(), newInputs,
                                                     newEquation);
@@ -1064,7 +1059,7 @@ class EinsumMergeUp1Axis : public OpRewritePattern<tensorrt::ExpandRankOp> {
       }
     }
 
-    std::string newEquation = equation.lhs + "->" + newRhs;
+    SmallString<128> newEquation{equation.lhs, "->", newRhs};
     rewriter.replaceOpWithNewOp<tensorrt::EinsumOp>(
         op, op.getType(), einsum.getInputs(), newEquation);
     return success();
@@ -1166,7 +1161,7 @@ class PushReshapeUpThroughEinsum
     // check that all of the inputs are have the right groupping.  If this
     // doesn't happen then that means that the reshape can not get pushed
     // through
-    for (std::string &eqLhs : equation.lhsParts) {
+    for (auto &eqLhs : equation.lhsParts) {
       for (char c : eqLhs) {
         auto it = charToGroup.find(c);
         if (it == charToGroup.end())
@@ -1193,7 +1188,7 @@ class PushReshapeUpThroughEinsum
     for (size_t i = 0; i < einsum.getInputs().size(); i++) {
       Value input = einsum.getInputs()[i];
       auto inputType = cast<RankedTensorType>(input.getType());
-      std::string newInputEquation = "";
+      SmallString<128> newInputEquation{""};
       SmallVector<int64_t> newInputShape;
       SmallVector<int64_t> newInputTranspose;
       for (int j = 0; j < inputType.getRank(); j++) {
@@ -1229,7 +1224,7 @@ class PushReshapeUpThroughEinsum
       newEquation.lhsParts.push_back(newInputEquation);
     }
 
-    std::string newEquationStr = newEquation.generateEquation();
+    StringRef newEquationStr = newEquation.generateEquation();
 
     if (has1OutputShape) {
       SmallVector<int64_t> newShape;
@@ -1419,13 +1414,13 @@ class PushReshapeDownThroughEinsum
       }
     }
 
-    for (std::string &part : equation.lhsParts) {
+    for (auto &part : equation.lhsParts) {
       for (char c : part) {
         auto group = charToGroup.find(c);
         if (group == charToGroup.end())
           continue;
         for (char c2 : group->second) {
-          if (part.find(c2) == std::string::npos)
+          if (part.find(c2) == StringRef::npos)
             return failure(
                 /* Missing dimensions that need to be reshaped together */);
         }
@@ -1437,7 +1432,7 @@ class PushReshapeDownThroughEinsum
       if (group == charToGroup.end())
         continue;
       for (char c2 : group->second) {
-        if (equation.rhs.find(c2) == std::string::npos)
+        if (equation.rhs.find(c2) == StringRef::npos)
           return failure(
               /* Missing dimensions that need to be reshaped together */);
       }
@@ -1486,7 +1481,7 @@ class PushReshapeDownThroughEinsum
       RankedTensorType inputType = cast<RankedTensorType>(input.getType());
       SmallVector<int64_t> newInputShape;
       SmallVector<int64_t> newInputTranspose;
-      std::string newEinsumStr = "";
+      SmallString<128> newEinsumStr{""};
       for (int j = 0; j < inputType.getRank(); j++) {
         char c = equation.lhsParts[i][j];
         auto it = charToGroup.find(c);
@@ -1504,7 +1499,7 @@ class PushReshapeDownThroughEinsum
           }
           for (char c2 : group->first) {
             size_t pos = equation.lhsParts[i].find(c2);
-            assert(pos != std::string::npos);
+            assert(pos != StringRef::npos);
             newInputTranspose.push_back(pos);
           }
         }
@@ -1555,13 +1550,13 @@ class PushReshapeDownThroughEinsum
         }
         for (char c2 : it->second) {
           size_t pos = equation.rhs.find(c2);
-          assert(pos != std::string::npos);
+          assert(pos != StringRef::npos);
           afterReshapeTranspose.push_back(pos);
         }
       }
     }
 
-    std::string newEinsumEquation = newEquation.generateEquation();
+    StringRef newEinsumEquation = newEquation.generateEquation();
 
     auto newEinsum = rewriter.create<tensorrt::EinsumOp>(
         op.getLoc(), outputType.clone(einsumOutputShape), newInputs,