fix vhlo DictionaryAttr's printer with separated comma (#1813)

openxla · Oct 25, 2023 · 89ea02e · 89ea02e
1 parent 43e51a7
commit 89ea02e
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 3 deletions.
diff --git a/stablehlo/dialect/VhloOps.cpp b/stablehlo/dialect/VhloOps.cpp
@@ -112,9 +112,8 @@ Attribute IntegerV1Attr::parse(AsmParser& parser, mlir::Type) {
 static void printAttributeDictionary(
     AsmPrinter& os, ArrayRef<std::pair<Attribute, Attribute>> values) {
   os << '{';
-  for (auto nvp : values) {
-    os << nvp.first << " = " << nvp.second;
-  }
+  llvm::interleaveComma(
+      values, os, [&](auto nvp) { os << nvp.first << " = " << nvp.second; });
   os << '}';
 }
 

diff --git a/stablehlo/tests/stablehlo_legalize_to_vhlo.mlir b/stablehlo/tests/stablehlo_legalize_to_vhlo.mlir
@@ -157,6 +157,12 @@ func.func @attr_custom_call_api_version_status_returning_unified(%arg0: tensor<f
 }
 // CHECK-LABEL: "attr_custom_call_api_version_status_returning_unified"
 
+// CHECK: #vhlo.dict_v1<{#vhlo.string_v1<"attr1"> = #vhlo.integer_v1<1 : i32>, #vhlo.string_v1<"attr2"> = #vhlo.integer_v1<2 : i32>}
+func.func @attr_dict() attributes {stablehlo.attr = {attr1 = 1 : i32, attr2 = 2 : i32}} {
+  return
+}
+// CHECK-LABEL: "attr_dict"
+
 // DotDimensionNumbers aka #stablehlo.dot is covered below.
 
 func.func @attr_fft_type_fft(%arg0: tensor<16xcomplex<f32>>) -> tensor<16xcomplex<f32>> {