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

[Benjamin Maxwell]

https://github.com/MacDue updated https://github.com/llvm/llvm-project/pull/83565

>From 2390d3408a0af1a225e3a508a1b8db6808a249ef 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 1/2] [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 ead19c69a0831c..4e4c6fd601777b 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 0d466795fac0de..9b916699f9f99b 100644
--- a/mlir/lib/Dialect/Arith/IR/ArithOps.cpp
+++ b/mlir/lib/Dialect/Arith/IR/ArithOps.cpp
@@ -423,6 +423,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
+}

>From ba65363b3bf484e2a72b1c0422b9226ba533a942 Mon Sep 17 00:00:00 2001
From: Benjamin Maxwell <benjamin.maxwell at arm.com>
Date: Tue, 2 Apr 2024 15:25:28 +0000
Subject: [PATCH 2/2] Tidy up test

---
 mlir/test/Dialect/Arith/vscale_constants.mlir | 14 +++++++-------
 1 file changed, 7 insertions(+), 7 deletions(-)

diff --git a/mlir/test/Dialect/Arith/vscale_constants.mlir b/mlir/test/Dialect/Arith/vscale_constants.mlir
index 568e20ae668455..324766f49980f4 100644
--- a/mlir/test/Dialect/Arith/vscale_constants.mlir
+++ b/mlir/test/Dialect/Arith/vscale_constants.mlir
@@ -3,12 +3,12 @@
 // 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
+  %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
 }