add enzyme broadcastop

EnzymeAD · wsmoses · Dec 27, 2024 · Nov 28, 2024 · Nov 28, 2024 · Nov 28, 2024
commit 396a29b8efe1f13efdb5aaf11c452affc169ae1b
diff --git a/enzyme/Enzyme/MLIR/Dialect/EnzymeOps.td b/enzyme/Enzyme/MLIR/Dialect/EnzymeOps.td
@@ -192,4 +192,20 @@ def GenericAdjointOp : Enzyme_Op<"genericAdjoint", [AttrSizedOperandSegments]> {
 
 }
 
+def BroadcastOp : Enzyme_Op<"broadcast"> {
+  let description = [{
+  Broadcast the operand by adding an extra dimension the frond with a size equal to the width attribute.
+  For scalar operands, a one-dimensional ranked tensor is created.
+
+  NOTE: Only works for scalars and *ranked* tensors for now.
+  }];
+
+  let arguments = (ins AnyType:$input, I64Attr:$width);
+  let results = (outs AnyRankedTensor:$output);
+
+  let builders = [
+    OpBuilder<(ins "Value":$input, "int64_t":$width)>
+  ];
+}
+
 #endif // ENZYME_OPS
diff --git a/enzyme/Enzyme/MLIR/Dialect/Ops.cpp b/enzyme/Enzyme/MLIR/Dialect/Ops.cpp
@@ -191,3 +191,23 @@ LogicalResult BatchOp::verifySymbolUses(SymbolTableCollection &symbolTable) {
 
   return success();
 }
+
+//===----------------------------------------------------------------------===//
+// BroadcastOp
+//===----------------------------------------------------------------------===//
+
+void BroadcastOp::build(OpBuilder &builder, OperationState &result, Value input, int64_t width) {
+  auto widthAttr = builder.getI64IntegerAttr(width);
+  RankedTensorType output;
+  // TODO: support things other than scalars and ranked tensors, maybe reuse getShadowType here?
+  if (auto tensorType = input.getType().dyn_cast<TensorType>()) {
+      auto shape = tensorType.getShape();
+      SmallVector<int64_t, 4> newShape;
+      newShape.push_back(width);
+      newShape.append(shape.begin(), shape.end());
+      output = RankedTensorType::get(newShape, tensorType.getElementType());
+  } else {
+    output = RankedTensorType::get({width}, input.getType());
+  }
+  build(builder, result, output, input, widthAttr);
+}