Add BatchNorm node, attributes, and ASM emitter support (#259)

Introduces BatchNormNode and BatchnormAttr for both inference and
training forward phases, with MLIR assembly emission via
torch.aten.native_batch_norm. Includes unit tests and lit tests for
NCHW layout.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

---------

Signed-off-by: Rob Suderman <rob.suderman@gmail.com>
Co-authored-by: Claude Sonnet 4.6 <noreply@anthropic.com>

Author: rsuderman

include/fusilli.h

@@ -31,6 +31,7 @@
 
 // Attributes / Types:
 #include "fusilli/attributes/attributes.h"           // IWYU pragma: export
+#include "fusilli/attributes/batchnorm_attributes.h" // IWYU pragma: export
 #include "fusilli/attributes/common.h"               // IWYU pragma: export
 #include "fusilli/attributes/conv_attributes.h"      // IWYU pragma: export
 #include "fusilli/attributes/custom_op_attributes.h" // IWYU pragma: export
@@ -43,6 +44,7 @@
 #include "fusilli/attributes/types.h"                // IWYU pragma: export
 
 // Nodes:
+#include "fusilli/node/batchnorm_node.h" // IWYU pragma: export
 #include "fusilli/node/conv_node.h"      // IWYU pragma: export
 #include "fusilli/node/custom_op_node.h" // IWYU pragma: export
 #include "fusilli/node/layernorm_node.h" // IWYU pragma: export

include/fusilli/attributes/batchnorm_attributes.h

@@ -0,0 +1,94 @@
+// Copyright 2026 Advanced Micro Devices, Inc.
+//
+// Licensed under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+//===----------------------------------------------------------------------===//
+//
+// This file contains attributes (compile-time constant metadata) for
+// batch normalization nodes.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef FUSILLI_ATTRIBUTES_BATCHNORM_ATTRIBUTES_H
+#define FUSILLI_ATTRIBUTES_BATCHNORM_ATTRIBUTES_H
+
+#include "fusilli/attributes/attributes.h"
+#include "fusilli/attributes/common.h"
+#include "fusilli/attributes/tensor_attributes.h"
+
+#include <cstdint>
+#include <memory>
+#include <unordered_map>
+
+namespace fusilli {
+
+class BatchnormAttr : public AttributesCRTP<BatchnormAttr> {
+public:
+  // Names for Tensor Inputs and Outputs.
+  enum class InputNames : uint8_t {
+    X,
+    SCALE,
+    BIAS,
+    MEAN,
+    VAR,
+    EPSILON,
+    MOMENTUM
+  };
+  enum class OutputNames : uint8_t { Y, SAVED_MEAN, SAVED_INV_VARIANCE };
+
+  std::unordered_map<InputNames, std::shared_ptr<TensorAttr>> inputs;
+  std::unordered_map<OutputNames, std::shared_ptr<TensorAttr>> outputs;
+
+  // Setters:
+  FUSILLI_GENERIC_INPUT_TENSOR_SETTER(BatchnormAttr, InputNames, X)
+  FUSILLI_GENERIC_INPUT_TENSOR_SETTER(BatchnormAttr, InputNames, SCALE)
+  FUSILLI_GENERIC_INPUT_TENSOR_SETTER(BatchnormAttr, InputNames, BIAS)
+  FUSILLI_GENERIC_INPUT_TENSOR_SETTER(BatchnormAttr, InputNames, MEAN)
+  FUSILLI_GENERIC_INPUT_TENSOR_SETTER(BatchnormAttr, InputNames, VAR)
+  FUSILLI_GENERIC_OUTPUT_TENSOR_SETTER(BatchnormAttr, OutputNames, Y)
+  FUSILLI_GENERIC_OUTPUT_TENSOR_SETTER(BatchnormAttr, OutputNames, SAVED_MEAN)
+  FUSILLI_GENERIC_OUTPUT_TENSOR_SETTER(BatchnormAttr, OutputNames,
+                                       SAVED_INV_VARIANCE)
+
+  BatchnormAttr &setEpsilon(const std::shared_ptr<TensorAttr> &epsilon) {
+    return setInput(InputNames::EPSILON, epsilon);
+  }
+
+  BatchnormAttr &setMomentum(const std::shared_ptr<TensorAttr> &momentum) {
+    return setInput(InputNames::MOMENTUM, momentum);
+  }
+
+  BatchnormAttr &setForwardPhase(NormFwdPhase forwardPhase) {
+    forwardPhase_ = forwardPhase;
+    return *this;
+  }
+
+  // Getters:
+  FUSILLI_GENERIC_INPUT_TENSOR_GETTER(InputNames, X)
+  FUSILLI_GENERIC_INPUT_TENSOR_GETTER(InputNames, SCALE)
+  FUSILLI_GENERIC_INPUT_TENSOR_GETTER(InputNames, BIAS)
+  FUSILLI_GENERIC_INPUT_TENSOR_GETTER(InputNames, MEAN)
+  FUSILLI_GENERIC_INPUT_TENSOR_GETTER(InputNames, VAR)
+  FUSILLI_GENERIC_OUTPUT_TENSOR_GETTER(OutputNames, Y)
+  FUSILLI_GENERIC_OUTPUT_TENSOR_GETTER(OutputNames, SAVED_MEAN)
+  FUSILLI_GENERIC_OUTPUT_TENSOR_GETTER(OutputNames, SAVED_INV_VARIANCE)
+
+  std::shared_ptr<TensorAttr> getEpsilon() const {
+    return getInput(InputNames::EPSILON);
+  }
+
+  std::shared_ptr<TensorAttr> getMomentum() const {
+    return getInput(InputNames::MOMENTUM);
+  }
+
+  NormFwdPhase getForwardPhase() const { return forwardPhase_; }
+
+private:
+  NormFwdPhase forwardPhase_ = NormFwdPhase::NOT_SET;
+};
+
+} // namespace fusilli
+
+#endif // FUSILLI_ATTRIBUTES_BATCHNORM_ATTRIBUTES_H

include/fusilli/graph/graph.h

@@ -14,6 +14,7 @@
 #ifndef FUSILLI_GRAPH_GRAPH_H
 #define FUSILLI_GRAPH_GRAPH_H
 
+#include "fusilli/attributes/batchnorm_attributes.h"
 #include "fusilli/attributes/common.h"
 #include "fusilli/attributes/conv_attributes.h"
 #include "fusilli/attributes/custom_op_attributes.h"
@@ -30,6 +31,7 @@
 #include "fusilli/backend/compile_session.h"
 #include "fusilli/backend/handle.h"
 #include "fusilli/graph/context.h"
+#include "fusilli/node/batchnorm_node.h"
 #include "fusilli/node/conv_node.h"
 #include "fusilli/node/custom_op_node.h"
 #include "fusilli/node/layernorm_node.h"
@@ -268,6 +270,13 @@ class Graph : public INode {
                                         const std::shared_ptr<TensorAttr> &w,
                                         ConvDGradAttr &attributes);
   std::array<std::shared_ptr<TensorAttr>, 3>
+  batchnorm(const std::shared_ptr<TensorAttr> &x,
+            const std::shared_ptr<TensorAttr> &scale,
+            const std::shared_ptr<TensorAttr> &bias,
+            const std::shared_ptr<TensorAttr> &mean,
+            const std::shared_ptr<TensorAttr> &var, BatchnormAttr &attributes);
+
+  std::array<std::shared_ptr<TensorAttr>, 3>
   layernorm(const std::shared_ptr<TensorAttr> &x,
             const std::shared_ptr<TensorAttr> &scale,
             const std::shared_ptr<TensorAttr> &bias, LayernormAttr &attributes);
@@ -727,6 +736,60 @@ Graph::convDGrad(const std::shared_ptr<TensorAttr> &dy,
   return dx;
 }
 
+// Create a BatchNormNode, populate it with the specified attributes, create
+// output tensors and add the node to the graph's sub nodes.
+inline std::array<std::shared_ptr<TensorAttr>, 3>
+Graph::batchnorm(const std::shared_ptr<TensorAttr> &x,
+                 const std::shared_ptr<TensorAttr> &scale,
+                 const std::shared_ptr<TensorAttr> &bias,
+                 const std::shared_ptr<TensorAttr> &mean,
+                 const std::shared_ptr<TensorAttr> &var,
+                 BatchnormAttr &batchnormAttr) {
+  // Populate names when not set.
+  if (batchnormAttr.getName().empty())
+    batchnormAttr.setName("batchnorm_" + std::to_string(subNodes_.size()));
+  if (x && x->getName().empty())
+    x->setName(batchnormAttr.getName() + "_X");
+  if (scale && scale->getName().empty())
+    scale->setName(batchnormAttr.getName() + "_SCALE");
+  if (bias && bias->getName().empty())
+    bias->setName(batchnormAttr.getName() + "_BIAS");
+  if (mean && mean->getName().empty())
+    mean->setName(batchnormAttr.getName() + "_MEAN");
+  if (var && var->getName().empty())
+    var->setName(batchnormAttr.getName() + "_VAR");
+  auto eps = batchnormAttr.getEpsilon();
+  if (eps && eps->getName().empty())
+    eps->setName(batchnormAttr.getName() + "_EPSILON");
+  auto mom = batchnormAttr.getMomentum();
+  if (mom && mom->getName().empty())
+    mom->setName(batchnormAttr.getName() + "_MOMENTUM");
+
+  FUSILLI_LOG_LABEL_ENDL("INFO: Adding BatchNorm '" << batchnormAttr.getName()
+                                                    << "' to Graph");
+
+  // Set inputs.
+  batchnormAttr.setX(x).setSCALE(scale).setBIAS(bias).setMEAN(mean).setVAR(var);
+
+  // Set outputs.
+  std::shared_ptr<TensorAttr> y = outputTensor(batchnormAttr.getName() + "_Y");
+  std::shared_ptr<TensorAttr> savedMean = nullptr;
+  std::shared_ptr<TensorAttr> savedInvVar = nullptr;
+  if (batchnormAttr.getForwardPhase() == NormFwdPhase::TRAINING) {
+    savedMean = outputTensor(batchnormAttr.getName() + "_SAVED_MEAN");
+    savedInvVar = outputTensor(batchnormAttr.getName() + "_SAVED_INV_VARIANCE");
+  }
+  batchnormAttr.setY(y);
+  batchnormAttr.setSAVED_MEAN(savedMean);
+  batchnormAttr.setSAVED_INV_VARIANCE(savedInvVar);
+
+  // Create node and add to Graph's subNodes_.
+  subNodes_.emplace_back(
+      std::make_unique<BatchNormNode>(std::move(batchnormAttr), context));
+
+  return {std::move(y), std::move(savedMean), std::move(savedInvVar)};
+}
+
 // Create a LayerNormNode, populate it with the specified attributes, create
 // output tensors and add the node to the graph's sub nodes
 inline std::array<std::shared_ptr<TensorAttr>, 3>