[StableHLO Builder] Add API to set frontend attributes

PiperOrigin-RevId: 820455957

Author: GleasonK

stablehlo/stablehlo/integrations/cpp/builder/StablehloBuilder.cpp

@@ -29,10 +29,36 @@ limitations under the License.
 #include "stablehlo/dialect/TypeInference.h"
 #include "stablehlo/integrations/cpp/builder/AttrTypeBuilderUtil.h"
 #include "stablehlo/integrations/cpp/builder/MlirBuilder.h"
+#include "third_party/llvm/llvm-project/mlir/include/mlir/IR/Attributes.h"
 
 namespace mlir {
 namespace stablehlo {
 
+///////////////
+// Dialect Helpers
+///////////////
+
+MlirOp AttachFrontendAttribute(MlirBuilder& builder, MlirOp op, StringRef name,
+                               Attribute value) {
+  constexpr char kFrontendAttrName[] = "mhlo.frontend_attributes";
+  Operation* mlirOp = unwrap(op).getDefiningOp();
+  SmallVector<NamedAttribute> attrs;
+  DictionaryAttr frontendAttr =
+      mlirOp->getAttrOfType<DictionaryAttr>(kFrontendAttrName);
+  if (frontendAttr) {
+    for (NamedAttribute attr : frontendAttr.getValue()) {
+      // Populate all non-conflicting names.
+      if (attr.getName() != name) {
+        attrs.push_back(attr);
+      }
+    }
+  }
+  attrs.emplace_back(name, value);
+  mlirOp->setAttr(kFrontendAttrName,
+                  DictionaryAttr::get(&builder.getContext(), attrs));
+  return op;
+}
+
 /////////////////
 // MANUAL APIs
 /////////////////

stablehlo/stablehlo/integrations/cpp/builder/StablehloBuilder.h

@@ -27,10 +27,23 @@ limitations under the License.
 #include "stablehlo/dialect/StablehloOps.h"
 #include "stablehlo/integrations/cpp/builder/AttrTypeBuilderUtil.h"
 #include "stablehlo/integrations/cpp/builder/MlirBuilder.h"
+#include "third_party/llvm/llvm-project/mlir/include/mlir/IR/Attributes.h"
 
 namespace mlir {
 namespace stablehlo {
 
+///////////////
+// Dialect Helpers
+///////////////
+
+// Appends or overwrites an entry in the `mhlo.frontend_attributes` attribute
+//
+// of the given op.
+// Ex:
+//   stablehlo.abs %0 { mhlo.frontend_attributes = { "foo" = 123 } }
+MlirOp AttachFrontendAttribute(MlirBuilder& builder, MlirOp op, StringRef name,
+                               Attribute value);
+
 /////////////////
 // MANUAL APIs
 /////////////////

stablehlo/stablehlo/integrations/cpp/builder/StablehloBuilderTest.cpp

@@ -1592,5 +1592,57 @@ TEST(MlirBuilderTest, ResultAccuracyAttrTolerance) {
   EXPECT_EQ(expected, debugString(*module));
 }
 
+TEST(MlirBuilderTest, FrontendAttributesAppend) {
+  std::string expected = R"mlir(module {
+  func.func @main(%arg0: tensor<2xf32>) -> tensor<2xf32> {
+    %0 = stablehlo.exponential %arg0 {mhlo.frontend_attributes = {bar = "hello", foo = 123 : i32}} : tensor<2xf32>
+    return %0 : tensor<2xf32>
+  }
+})mlir";
+
+  StablehloModuleBuilder mb;
+  {
+    Location funcLoc = fileLineColLoc(mb->getContext(), "main.mlir", 1, 1);
+    func::FunctionBuilder fb(mb.get(), "main", funcLoc);
+    auto type = makeTensorType(fb.getContext(), {2}, ElementType::F32);
+    auto arg0 = func::Argument(fb, type);
+    auto exp = Exp(arg0);
+    stablehlo::AttachFrontendAttribute(
+        fb, exp, "foo", fb.getOpBuilder().getI32IntegerAttr(123));
+    stablehlo::AttachFrontendAttribute(
+        fb, exp, "bar", fb.getOpBuilder().getStringAttr("hello"));
+    func::Return(fb, {exp});
+  }
+
+  OwningOpRef<ModuleOp> module = mb->build();
+  EXPECT_EQ(expected, debugString(*module));
+}
+
+TEST(MlirBuilderTest, FrontendAttributesOverwrite) {
+  std::string expected = R"mlir(module {
+  func.func @main(%arg0: tensor<2xf32>) -> tensor<2xf32> {
+    %0 = stablehlo.exponential %arg0 {mhlo.frontend_attributes = {foo = 456 : i32}} : tensor<2xf32>
+    return %0 : tensor<2xf32>
+  }
+})mlir";
+
+  StablehloModuleBuilder mb;
+  {
+    Location funcLoc = fileLineColLoc(mb->getContext(), "main.mlir", 1, 1);
+    func::FunctionBuilder fb(mb.get(), "main", funcLoc);
+    auto type = makeTensorType(fb.getContext(), {2}, ElementType::F32);
+    auto arg0 = func::Argument(fb, type);
+    auto exp = Exp(arg0);
+    stablehlo::AttachFrontendAttribute(
+        fb, exp, "foo", fb.getOpBuilder().getI32IntegerAttr(123));
+    stablehlo::AttachFrontendAttribute(
+        fb, exp, "foo", fb.getOpBuilder().getI32IntegerAttr(456));
+    func::Return(fb, {exp});
+  }
+
+  OwningOpRef<ModuleOp> module = mb->build();
+  EXPECT_EQ(expected, debugString(*module));
+}
+
 }  // namespace stablehlo
 }  // namespace mlir