Add AMAX, AVG, NORM1, NORM2, MUL, MUL_NO_ZEROS reduction modes (#325)

Enable the remaining cuDNN reduction modes in ReductionAttr and add the
corresponding MLIR schemas to the asm emitter:

- NORM1 lowers to abs + sum.dim_IntList.
- AMAX lowers to abs + amax.
- AVG lowers to mean.dim (float dtypes only — torch.aten.mean.dim is not
defined on integer tensors, so the sample skips int32 for AVG).
- NORM2 lowers to mul + sum.dim_IntList + sqrt.
- MUL lowers directly to torch.prims.prod.
- MUL_NO_ZEROS uses aten.ne.Scalar to build an i1 mask, then
aten.where.ScalarOther to substitute 1 for zero entries before feeding
the result to torch.prims.prod, so zero inputs are excluded from the
product.

Extend samples/reduction/reduction_ops.cpp to exercise every new mode.
Input data is built by a per-mode generateReductionInputData helper so
MUL/MUL_NO_ZEROS get a non-trivial pattern (mostly 1s with a 2 and a 3,
plus injected zeros for MUL_NO_ZEROS) that stays in range for
fp16/int32, and the expected value is computed by the existing reference
reduction loop rather than hardcoded.

Add lit tests for each new mode under tests/lit/ and register them in
tests/CMakeLists.txt.

---------

Signed-off-by: Rob Suderman <rob.suderman@gmail.com>
Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

Author: rsuderman

include/fusilli/attributes/reduction_attributes.h

@@ -26,16 +26,16 @@ namespace fusilli {
 
 #define FUSILLI_REDUCTION_MODES(OP)                                            \
   OP(NOT_SET)                                                                  \
-  OP(SUM)                                                                      \
-  /* OP(ADD) */                                                                \
-  /* OP(MUL) */                                                                \
+  OP(ADD)                                                                      \
+  OP(SUM) /* Remove once SUM is deprecated  */                                 \
+  OP(MUL)                                                                      \
   OP(MIN)                                                                      \
   OP(MAX)                                                                      \
-  /* OP(AMAX) */                                                               \
-  /* OP(AVG) */                                                                \
-  /* OP(NORM1) */                                                              \
-  /* OP(NORM2) */                                                              \
-  /* OP(MUL_NO_ZEROS) */
+  OP(AMAX)                                                                     \
+  OP(AVG)                                                                      \
+  OP(NORM1)                                                                    \
+  OP(NORM2)                                                                    \
+  OP(MUL_NO_ZEROS)
 
 class ReductionAttr : public AttributesCRTP<ReductionAttr> {
 public:

include/fusilli/attributes/types.h

@@ -45,6 +45,24 @@ enum class DataType : uint8_t {
 #undef DEFINE_ENUM
 };
 
+inline bool isIntegerType(DataType type) {
+  switch (type) {
+  case DataType::Uint8:
+  case DataType::Int4:
+  case DataType::Int8:
+  case DataType::Int16:
+  case DataType::Int32:
+  case DataType::Int64:
+    return true;
+  default:
+    return false;
+  }
+}
+
+inline bool isIntegralOrBoolType(DataType type) {
+  return isIntegerType(type) || type == DataType::Boolean;
+}
+
 // Map from Fusilli types to MLIR types.
 static const std::unordered_map<DataType, std::string> kDataTypeToMlirTypeAsm =
     {

include/fusilli/node/reduction_node.h

@@ -119,6 +119,13 @@ class ReductionNode : public NodeCRTP<ReductionNode> {
         ErrorCode::InvalidAttribute,
         "Reduction input and output must have the same rank");
 
+    FUSILLI_RETURN_ERROR_IF(
+        reductionAttr.getMode() == ReductionAttr::Mode::AVG &&
+            (isIntegralOrBoolType(xTensor->getDataType()) ||
+             isIntegralOrBoolType(yTensor->getDataType())),
+        ErrorCode::InvalidAttribute,
+        "Reduction AVG is not supported for integral or boolean tensors");
+
     // Validate reduction dimensions - if Y[i] differs from X[i], Y[i] must be 1
     const auto &xDim = xTensor->getDim();
     const auto &yDim = yTensor->getDim();

include/fusilli/support/asm_emitter.h

@@ -2155,6 +2155,56 @@ inline std::string ReductionNode::emitNodePreAsm() const {
     {7}
     )";
 
+  constexpr std::string_view kNorm1Schema = R"(
+    {0}
+    {1}
+    %abs_{2} = torch.aten.abs {4} : {5} -> {5}
+    %keepdim_{2} = torch.constant.bool true
+    %dtype_{2} = torch.constant.none
+    {3}_{2}_perm = torch.aten.sum.dim_IntList %abs_{2}, %reduction_dims_{2}, %keepdim_{2}, %dtype_{2} : {5}, !torch.list<int>, !torch.bool, !torch.none -> {6}
+    {7}
+    )";
+
+  constexpr std::string_view kNorm2Schema = R"(
+    {0}
+    {1}
+    %sq_{2} = torch.aten.mul.Tensor {4}, {4} : {5}, {5} -> {5}
+    %keepdim_{2} = torch.constant.bool true
+    %dtype_{2} = torch.constant.none
+    %sumsq_{2} = torch.aten.sum.dim_IntList %sq_{2}, %reduction_dims_{2}, %keepdim_{2}, %dtype_{2} : {5}, !torch.list<int>, !torch.bool, !torch.none -> {6}
+    {3}_{2}_perm = torch.aten.sqrt %sumsq_{2} : {6} -> {6}
+    {7}
+    )";
+
+  constexpr std::string_view kAbsAmaxSchema = R"(
+    {0}
+    {1}
+    %abs_{2} = torch.aten.abs {4} : {5} -> {5}
+    %keepdim_{2} = torch.constant.bool true
+    {3}_{2}_perm = torch.aten.amax %abs_{2}, %reduction_dims_{2}, %keepdim_{2} : {5}, !torch.list<int>, !torch.bool -> {6}
+    {7}
+    )";
+
+  constexpr std::string_view kMulSchema = R"(
+    {0}
+    {1}
+    %dtype_{2} = torch.constant.none
+    {3}_{2}_perm = torch.prims.prod {4}, %reduction_dims_{2}, %dtype_{2} : {5}, !torch.list<int>, !torch.none -> {6}
+    {7}
+    )";
+
+  constexpr std::string_view kMulNoZerosSchema = R"(
+    {0}
+    {1}
+    %zero_{2} = torch.constant.int 0
+    %mask_{2} = torch.aten.ne.Scalar {4}, %zero_{2} : {5}, !torch.int -> {8}
+    %one_{2} = torch.constant.int 1
+    %fixed_{2} = torch.aten.where.ScalarOther %mask_{2}, {4}, %one_{2} : {8}, {5}, !torch.int -> {5}
+    %dtype_{2} = torch.constant.none
+    {3}_{2}_perm = torch.prims.prod %fixed_{2}, %reduction_dims_{2}, %dtype_{2} : {5}, !torch.list<int>, !torch.none -> {6}
+    {7}
+    )";
+
 #define FUSILLI_DECLARE_REDUCTION_EMITTER(MODE, SCHEMA, OPIR)                  \
   case ReductionAttr::Mode::MODE: {                                            \
     return std::format(SCHEMA, permuteX,     /* {0} */                         \
@@ -2175,11 +2225,55 @@ inline std::string ReductionNode::emitNodePreAsm() const {
 #define FUSILLI_DECLARE_KEEPDIM_DTYPE_REDUCTION_EMITTER(MODE, OPIR)            \
   FUSILLI_DECLARE_REDUCTION_EMITTER(MODE, kKeepdimDtypeReductionSchema, OPIR)
 
+#define FUSILLI_DECLARE_CUSTOM_REDUCTION_EMITTER(MODE, SCHEMA)                 \
+  case ReductionAttr::Mode::MODE: {                                            \
+    return std::format(SCHEMA, permuteX,     /* {0} */                         \
+                       dimListOss.str(),     /* {1} */                         \
+                       suffix,               /* {2} */                         \
+                       getResultNamesAsm(),  /* {3} */                         \
+                       getOperandNamesAsm(), /* {4} */                         \
+                       getOperandTypesAsm(), /* {5} */                         \
+                       getResultTypesAsm(),  /* {6} */                         \
+                       permuteY              /* {7} */                         \
+    );                                                                         \
+  }
+
   switch (reductionAttr.getMode()) {
+  case ReductionAttr::Mode::ADD:
     FUSILLI_DECLARE_KEEPDIM_DTYPE_REDUCTION_EMITTER(SUM,
                                                     torch.aten.sum.dim_IntList)
     FUSILLI_DECLARE_KEEPDIM_REDUCTION_EMITTER(MIN, torch.aten.amin)
     FUSILLI_DECLARE_KEEPDIM_REDUCTION_EMITTER(MAX, torch.aten.amax)
+    FUSILLI_DECLARE_KEEPDIM_DTYPE_REDUCTION_EMITTER(AVG, torch.aten.mean.dim)
+    FUSILLI_DECLARE_CUSTOM_REDUCTION_EMITTER(NORM1, kNorm1Schema)
+    FUSILLI_DECLARE_CUSTOM_REDUCTION_EMITTER(NORM2, kNorm2Schema)
+    FUSILLI_DECLARE_CUSTOM_REDUCTION_EMITTER(AMAX, kAbsAmaxSchema)
+    FUSILLI_DECLARE_CUSTOM_REDUCTION_EMITTER(MUL, kMulSchema)
+
+  case ReductionAttr::Mode::MUL_NO_ZEROS: {
+    // Build a bool tensor type matching the (logical) shape of X. This is
+    // used for the mask produced by `torch.aten.ne.Scalar`.
+    std::ostringstream boolTypeOss;
+    boolTypeOss << "!torch.vtensor<[";
+    const auto &xDims = xT->getDim();
+    interleave(
+        xDims.begin(), xDims.end(),
+        [&](int64_t d) { boolTypeOss << std::to_string(d); },
+        [&] { boolTypeOss << ","; });
+    boolTypeOss << "],i1>";
+    std::string boolType = boolTypeOss.str();
+
+    return std::format(kMulNoZerosSchema, permuteX, /* {0} */
+                       dimListOss.str(),            /* {1} */
+                       suffix,                      /* {2} */
+                       getResultNamesAsm(),         /* {3} */
+                       getOperandNamesAsm(),        /* {4} */
+                       getOperandTypesAsm(),        /* {5} */
+                       getResultTypesAsm(),         /* {6} */
+                       permuteY,                    /* {7} */
+                       boolType                     /* {8} */
+    );
+  }
   default:
     assert(false && "Unsupported reduction mode");
     return "";
@@ -2189,6 +2283,7 @@ inline std::string ReductionNode::emitNodePreAsm() const {
 #undef FUSILLI_DECLARE_REDUCTION_EMITTER
 #undef FUSILLI_DECLARE_KEEPDIM_REDUCTION_EMITTER
 #undef FUSILLI_DECLARE_KEEPDIM_DTYPE_REDUCTION_EMITTER
+#undef FUSILLI_DECLARE_CUSTOM_REDUCTION_EMITTER
 
 //===----------------------------------------------------------------------===//
 //

samples/reduction/reduction_ops.cpp

@@ -23,6 +23,46 @@
 
 using namespace fusilli;
 
+// Builds the input tensor data for a given reduction mode. The reduction
+// in this sample is checked at output index (d0=0, d1=0), which corresponds
+// to xData indices [0, 64). Each branch picks values that produce a
+// non-trivial, deterministic expected value at that output index while
+// staying in range for fp16/int32.
+template <typename T>
+static std::vector<T> generateReductionInputData(ReductionAttr::Mode mode,
+                                                 int64_t xSize) {
+  std::vector<T> xData(xSize);
+  switch (mode) {
+  case ReductionAttr::Mode::MUL: {
+    // Mostly 1s with a 2 and a 3 inside the (d0=0, d1=0) block so the
+    // product over that block is 6 — non-trivial enough to verify the
+    // multiplication actually happens (rather than degenerating to 1).
+    std::fill(xData.begin(), xData.end(), T(1));
+    xData[10] = T(2);
+    xData[50] = T(3);
+    break;
+  }
+  case ReductionAttr::Mode::MUL_NO_ZEROS: {
+    // Same shape as MUL but with zeros sprinkled in. The expected
+    // product over the (d0=0, d1=0) block is still 6 because zeros
+    // are excluded from the reduction.
+    std::fill(xData.begin(), xData.end(), T(1));
+    xData[10] = T(2);
+    xData[50] = T(3);
+    xData[5] = T(0);
+    xData[20] = T(0);
+    break;
+  }
+  default: {
+    // Values from -50 to 49.
+    for (int64_t i = 0; i < xSize; ++i)
+      xData[i] = static_cast<T>(i % 100 - 50);
+    break;
+  }
+  }
+  return xData;
+}
+
 // Based on parameters, generates a unique name for the graph
 static std::string generateName(ReductionAttr::Mode mode, DataType type,
                                 const std::vector<int64_t> &xDim,
@@ -46,9 +86,16 @@ TEST_CASE("Reduction ops", "[reduction][graph]") {
 
   // clang-format off
   const auto mode = GENERATE(
+      ReductionAttr::Mode::ADD,
       ReductionAttr::Mode::SUM,
       ReductionAttr::Mode::MIN,
-      ReductionAttr::Mode::MAX);
+      ReductionAttr::Mode::MAX,
+      ReductionAttr::Mode::AMAX,
+      ReductionAttr::Mode::AVG,
+      ReductionAttr::Mode::NORM1,
+      ReductionAttr::Mode::NORM2,
+      ReductionAttr::Mode::MUL,
+      ReductionAttr::Mode::MUL_NO_ZEROS);
   // clang-format on
 
   auto execute = [&]<typename T>(Handle &handle, DataType dt, T initValue) {
@@ -80,11 +127,7 @@ TEST_CASE("Reduction ops", "[reduction][graph]") {
     for (auto d : xDims)
       xSize *= d;
 
-    std::vector<T> xData(xSize);
-    for (int64_t i = 0; i < xSize; ++i) {
-      // Values from -50 to 49
-      xData[i] = static_cast<T>(i % 100 - 50);
-    }
+    std::vector<T> xData = generateReductionInputData<T>(mode, xSize);
 
     FUSILLI_REQUIRE_ASSIGN(auto xBuf, allocateBufferOfType(handle, xT, xData));
     FUSILLI_REQUIRE_ASSIGN(auto yBuf,
@@ -118,12 +161,14 @@ TEST_CASE("Reduction ops", "[reduction][graph]") {
 
     // Initialize to same value as output buffer
     T expectedValue = initValue;
+    double sumSq = 0.0;
 
     // Perform reduction
     for (int64_t d2 = 0; d2 < 8; ++d2) {
       for (int64_t d3 = 0; d3 < 8; ++d3) {
         int64_t inIdx = ((d0 * 16 + d1) * 8 + d2) * 8 + d3;
         switch (mode) {
+        case ReductionAttr::Mode::ADD:
         case ReductionAttr::Mode::SUM:
           expectedValue = expectedValue + xData[inIdx];
           break;
@@ -133,11 +178,43 @@ TEST_CASE("Reduction ops", "[reduction][graph]") {
         case ReductionAttr::Mode::MAX:
           expectedValue = std::max(expectedValue, xData[inIdx]);
           break;
+        case ReductionAttr::Mode::NORM1:
+          expectedValue =
+              static_cast<T>(static_cast<double>(expectedValue) +
+                             std::abs(static_cast<double>(xData[inIdx])));
+          break;
+        case ReductionAttr::Mode::AMAX:
+          expectedValue = std::max(
+              expectedValue,
+              static_cast<T>(std::abs(static_cast<double>(xData[inIdx]))));
+          break;
+        case ReductionAttr::Mode::AVG:
+          expectedValue = expectedValue + xData[inIdx];
+          break;
+        case ReductionAttr::Mode::NORM2: {
+          double v = static_cast<double>(xData[inIdx]);
+          sumSq += v * v;
+          break;
+        }
+        case ReductionAttr::Mode::MUL:
+          expectedValue = expectedValue * xData[inIdx];
+          break;
+        case ReductionAttr::Mode::MUL_NO_ZEROS:
+          if (xData[inIdx] != T(0))
+            expectedValue = expectedValue * xData[inIdx];
+          break;
         default:
           break;
         }
       }
     }
+    if (mode == ReductionAttr::Mode::NORM2)
+      expectedValue = static_cast<T>(std::sqrt(sumSq));
+
+    // Finalize AVG by dividing the accumulated sum by the number of
+    // reduced elements (8 * 8 = 64). Integer types use integer division.
+    if (mode == ReductionAttr::Mode::AVG)
+      expectedValue = expectedValue / T(64);
 
     // Read output buffer
     std::vector<T> result;
@@ -146,10 +223,19 @@ TEST_CASE("Reduction ops", "[reduction][graph]") {
     // Validate output size and check the first value (d0=0, d1=0)
     REQUIRE(result.size() == ySize);
     int64_t checkIdx = d0 * 16 + d1;
-    if constexpr (std::is_floating_point_v<T>)
+    if constexpr (std::is_floating_point_v<T>) {
       REQUIRE(std::abs(result[checkIdx] - expectedValue) < T(0.01));
-    else
-      REQUIRE(result[checkIdx] == expectedValue);
+    } else {
+      if (mode == ReductionAttr::Mode::NORM2) {
+        // sqrt may round; allow off-by-one for integer types.
+        T diff = result[checkIdx] > expectedValue
+                     ? result[checkIdx] - expectedValue
+                     : expectedValue - result[checkIdx];
+        REQUIRE(diff <= T(1));
+      } else {
+        REQUIRE(result[checkIdx] == expectedValue);
+      }
+    }
   };
 
   // Create handle for the target backend.
@@ -158,19 +244,28 @@ TEST_CASE("Reduction ops", "[reduction][graph]") {
   // Determine initial value based on reduction mode
   auto getInitValue = [&]<typename T>() -> T {
     switch (mode) {
-    case ReductionAttr::Mode::SUM:
-      return T(0);
     case ReductionAttr::Mode::MIN:
       return std::numeric_limits<T>::max();
     case ReductionAttr::Mode::MAX:
       return std::numeric_limits<T>::lowest();
+    case ReductionAttr::Mode::MUL:
+    case ReductionAttr::Mode::MUL_NO_ZEROS:
+      return T(1);
+    case ReductionAttr::Mode::ADD:
+    case ReductionAttr::Mode::SUM:
+    case ReductionAttr::Mode::AMAX:
+    case ReductionAttr::Mode::AVG:
+    case ReductionAttr::Mode::NORM1:
+    case ReductionAttr::Mode::NORM2:
     default:
       return T(0);
     }
   };
 
-  // int32
-  execute(handle, DataType::Int32, getInitValue.template operator()<int>());
+  // torch.aten.mean.dim is not defined on integer tensors, so AVG only
+  // exercises the floating-point paths.
+  if (mode != ReductionAttr::Mode::AVG)
+    execute(handle, DataType::Int32, getInitValue.template operator()<int>());
   // fp16
   execute(handle, DataType::Half, getInitValue.template operator()<half>());
   // fp32

tests/CMakeLists.txt

@@ -223,9 +223,15 @@ add_fusilli_lit_tests(
     lit/test_sdpa_asm_emitter_gqa_hk_ne_hv.cpp
     lit/test_sdpa_asm_emitter_custom_scale.cpp
     lit/test_sdpa_asm_emitter_cross_attn.cpp
-    lit/test_reduction_asm_emitter_sum.cpp
+    lit/test_reduction_asm_emitter_add.cpp
     lit/test_reduction_asm_emitter_min.cpp
+    lit/test_reduction_asm_emitter_amax.cpp
     lit/test_reduction_asm_emitter_max.cpp
+    lit/test_reduction_asm_emitter_mul.cpp
+    lit/test_reduction_asm_emitter_mul_no_zeros.cpp
+    lit/test_reduction_asm_emitter_norm1.cpp
+    lit/test_reduction_asm_emitter_avg.cpp
+    lit/test_reduction_asm_emitter_norm2.cpp
   DEPS
     libfusilli
     libutils