[mlir][arith] Add result pretty printing for constant vscale values (#83565)

In scalable code it is very common to have constant multiples of vscale,
e.g. `4 * vscale`. This updates `arith.muli` to pretty print the result
name in cases like this, so `4 * vscale` would be `%c4_vscale`.

This makes reading IR dumps of scalable code a little nicer.

Author: MacDue

mlir/include/mlir/Dialect/Arith/IR/ArithOps.td

@@ -343,7 +343,9 @@ def Arith_SubIOp : Arith_IntBinaryOpWithOverflowFlags<"subi"> {
 // MulIOp
 //===----------------------------------------------------------------------===//
 
-def Arith_MulIOp : Arith_IntBinaryOpWithOverflowFlags<"muli", [Commutative]> {
+def Arith_MulIOp : Arith_IntBinaryOpWithOverflowFlags<"muli",
+  [Commutative, DeclareOpInterfaceMethods<OpAsmOpInterface, ["getAsmResultNames"]>]
+> {
   let summary = [{
     Integer multiplication operation.
   }];

mlir/lib/Dialect/Arith/IR/ArithOps.cpp

@@ -423,6 +423,33 @@ OpFoldResult arith::MulIOp::fold(FoldAdaptor adaptor) {
       [](const APInt &a, const APInt &b) { return a * b; });
 }
 
+void arith::MulIOp::getAsmResultNames(
+    function_ref<void(Value, StringRef)> setNameFn) {
+  if (!isa<IndexType>(getType()))
+    return;
+
+  // Match vector.vscale by name to avoid depending on the vector dialect (which
+  // is a circular dependency).
+  auto isVscale = [](Operation *op) {
+    return op && op->getName().getStringRef() == "vector.vscale";
+  };
+
+  IntegerAttr baseValue;
+  auto isVscaleExpr = [&](Value a, Value b) {
+    return matchPattern(a, m_Constant(&baseValue)) &&
+           isVscale(b.getDefiningOp());
+  };
+
+  if (!isVscaleExpr(getLhs(), getRhs()) && !isVscaleExpr(getRhs(), getLhs()))
+    return;
+
+  // Name `base * vscale` or `vscale * base` as `c<base_value>_vscale`.
+  SmallString<32> specialNameBuffer;
+  llvm::raw_svector_ostream specialName(specialNameBuffer);
+  specialName << 'c' << baseValue.getInt() << "_vscale";
+  setNameFn(getResult(), specialName.str());
+}
+
 void arith::MulIOp::getCanonicalizationPatterns(RewritePatternSet &patterns,
                                                 MLIRContext *context) {
   patterns.add<MulIMulIConstant>(context);

mlir/test/Dialect/Arith/vscale_constants.mlir

@@ -0,0 +1,14 @@
+// RUN: mlir-opt %s | FileCheck %s
+
+// Note: This test is checking value names (so deliberately is not using a regex match).
+
+func.func @test_vscale_constant_names() {
+  %vscale = vector.vscale
+  %c8 = arith.constant 8 : index
+  // CHECK: %c8_vscale = arith.muli
+  %0 = arith.muli %vscale, %c8 : index
+  %c10 = arith.constant 10 : index
+  // CHECK: %c10_vscale = arith.muli
+  %1 = arith.muli %c10, %vscale : index
+  return
+}