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

[Benjamin Maxwell]

https://github.com/MacDue created https://github.com/llvm/llvm-project/pull/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.

>From f9a30b0ae7993391ce4ded3f85bfa48e9ff7c978 Mon Sep 17 00:00:00 2001
From: Benjamin Maxwell <benjamin.maxwell at arm.com>
Date: Fri, 1 Mar 2024 12:22:21 +0000
Subject: [PATCH] [mlir][arith] Add result pretty printing for constant vscale
 values

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.
---
 .../include/mlir/Dialect/Arith/IR/ArithOps.td |  4 ++-
 mlir/lib/Dialect/Arith/IR/ArithOps.cpp        | 29 +++++++++++++++++++
 mlir/test/Dialect/Arith/vscale_constants.mlir | 14 +++++++++
 3 files changed, 46 insertions(+), 1 deletion(-)
 create mode 100644 mlir/test/Dialect/Arith/vscale_constants.mlir

diff --git a/mlir/include/mlir/Dialect/Arith/IR/ArithOps.td b/mlir/include/mlir/Dialect/Arith/IR/ArithOps.td
index c9df50d0395d1f..fe9f9539cdb5da 100644
--- a/mlir/include/mlir/Dialect/Arith/IR/ArithOps.td
+++ b/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.
   }];
diff --git a/mlir/lib/Dialect/Arith/IR/ArithOps.cpp b/mlir/lib/Dialect/Arith/IR/ArithOps.cpp
index 0f71c19c23b654..4d39aa7efb91f7 100644
--- a/mlir/lib/Dialect/Arith/IR/ArithOps.cpp
+++ b/mlir/lib/Dialect/Arith/IR/ArithOps.cpp
@@ -400,6 +400,35 @@ 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";
+  };
+
+  // Name `base * vscale` or `vscale * base` as `c<base_value>_vscale`.
+  IntegerAttr baseValue;
+  if (matchPattern(getLhs(), m_Constant(&baseValue)) &&
+      isVscale(getRhs().getDefiningOp())) {
+    // base * vscale
+  } else if (isVscale(getLhs().getDefiningOp()) &&
+             matchPattern(getRhs(), m_Constant(&baseValue))) {
+    // vscale * base
+  } else {
+    return;
+  }
+
+  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);
diff --git a/mlir/test/Dialect/Arith/vscale_constants.mlir b/mlir/test/Dialect/Arith/vscale_constants.mlir
new file mode 100644
index 00000000000000..568e20ae668455
--- /dev/null
+++ b/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() {
+  %0 = vector.vscale
+  %1 = arith.constant 8 : index
+  %2 = arith.constant 10 : index
+  // CHECK: %c8_vscale = arith.muli %c8, %vscale : index
+  %3 = arith.muli %1, %0 : index
+  // CHECK: %c10_vscale = arith.muli %vscale, %c10 : index
+  %4 = arith.muli %0, %2 : index
+  return
+}