PR1: add layout op

1. adding Fusion node `GroupedBlockScalingFactorLayoutOp`

GroupedBlockScalingFactorLayoutOp
  [output]
    Val* output (2d tensor)

  [input]
    TensorView* input (2d tensor)
    TensorView* expert_offsets (vector)
    TensorView* sf_offsets (vector)
    Val* k  (scalar)
    Val* g  (scalar)

  [attribute]
    BlockScalingFactorLayout layout

2. adding cpp api `groupedBlockSfLayout`

TensorView* groupedBlockSfLayout(
    TensorView* input,
    TensorView* expert_offsets,
    TensorView* sf_offsets,
    BlockScalingFactorLayout layout);

Author: jjsjann123

CMakeLists.txt

@@ -1138,6 +1138,9 @@ if(BUILD_TEST)
   add_test(test_reshape "${NVFUSER_ROOT}/tests/cpp/test_reshape.cpp" "")
   list(APPEND TEST_BINARIES test_reshape)
 
+  add_test(test_layout_op ${NVFUSER_ROOT}/tests/cpp/test_layout_op.cpp "")
+  list(APPEND TEST_BINARIES test_layout_op)
+
   set(MATMUL_TEST_SRCS)
   list(APPEND MATMUL_TEST_SRCS
     ${NVFUSER_ROOT}/tests/cpp/test_cutlass_scheduler.cpp

csrc/device_lower/utils.cpp

@@ -125,6 +125,7 @@ bool isTvOp(const Expr* expr) {
           SliceOp,
           CatOp,
           ScanOp,
+          GroupedBlockScalingFactorLayoutOp,
           kir::AllocTMem,
           kir::GridReduction,
           kir::GroupedGridReduction,

csrc/dispatch.h

@@ -68,59 +68,60 @@ class Val;
 #define DISPATCH_FOR_ALL_KIR_VALS(f) f(Predicate) f(TensorIndex)
 #define DISPATCH_FOR_ALL_HIR_VALS(f) f(Stream)
 
-#define DISPATCH_FOR_ALL_EXPRS(f) \
-  f(FullOp);                      \
-  f(IotaOp);                      \
-  f(EyeOp);                       \
-  f(UnaryOp);                     \
-  f(BinaryOp);                    \
-  f(TernaryOp);                   \
-  f(ArrayConstruct);              \
-  f(StructConstruct);             \
-  f(GetAttr);                     \
-  f(GetItem);                     \
-  f(ReverseArray);                \
-  f(GetMetaData);                 \
-  f(TensorConstruct);             \
-  f(SelectOp);                    \
-  f(IndexSelectOp);               \
-  f(IndexPutAccumulateOp);        \
-  f(GatherOp);                    \
-  f(ScatterOp);                   \
-  f(RNGOp);                       \
-  f(ReductionOp);                 \
-  f(GroupedReductionOp);          \
-  f(WelfordOp);                   \
-  f(GroupedWelfordOp);            \
-  f(LoadStoreOp);                 \
-  f(MmaOp);                       \
-  f(BroadcastOp);                 \
-  f(SqueezeOp);                   \
-  f(ExpandOp);                    \
-  f(RepeatOp);                    \
-  f(ViewAsScalar);                \
-  f(ReshapeOp);                   \
-  f(CatOp);                       \
-  f(PadOp);                       \
-  f(SliceOp);                     \
-  f(Split);                       \
-  f(ArgsortOp);                   \
-  f(GroupedMmaOp);                \
-  f(ScaledMmaOp);                 \
-  f(CutlassNvfp4GroupedMmaOp);    \
-  f(TopKOp);                      \
-  f(ScanOp);                      \
-  f(Merge);                       \
-  f(Swizzle);                     \
-  f(Swizzle2D);                   \
-  f(Resize);                      \
-  f(MatmulOp);                    \
-  f(LinearOp);                    \
-  f(SdpaFwdOp);                   \
-  f(SdpaBwdOp);                   \
-  f(EmbeddingFwdOp);              \
-  f(Communication);               \
-  f(ForLoop);                     \
+#define DISPATCH_FOR_ALL_EXPRS(f)       \
+  f(FullOp);                            \
+  f(IotaOp);                            \
+  f(EyeOp);                             \
+  f(UnaryOp);                           \
+  f(BinaryOp);                          \
+  f(TernaryOp);                         \
+  f(ArrayConstruct);                    \
+  f(StructConstruct);                   \
+  f(GetAttr);                           \
+  f(GetItem);                           \
+  f(ReverseArray);                      \
+  f(GetMetaData);                       \
+  f(TensorConstruct);                   \
+  f(SelectOp);                          \
+  f(IndexSelectOp);                     \
+  f(IndexPutAccumulateOp);              \
+  f(GatherOp);                          \
+  f(ScatterOp);                         \
+  f(RNGOp);                             \
+  f(ReductionOp);                       \
+  f(GroupedReductionOp);                \
+  f(WelfordOp);                         \
+  f(GroupedWelfordOp);                  \
+  f(LoadStoreOp);                       \
+  f(MmaOp);                             \
+  f(BroadcastOp);                       \
+  f(SqueezeOp);                         \
+  f(ExpandOp);                          \
+  f(RepeatOp);                          \
+  f(ViewAsScalar);                      \
+  f(ReshapeOp);                         \
+  f(CatOp);                             \
+  f(PadOp);                             \
+  f(SliceOp);                           \
+  f(Split);                             \
+  f(ArgsortOp);                         \
+  f(GroupedMmaOp);                      \
+  f(ScaledMmaOp);                       \
+  f(CutlassNvfp4GroupedMmaOp);          \
+  f(GroupedBlockScalingFactorLayoutOp); \
+  f(TopKOp);                            \
+  f(ScanOp);                            \
+  f(Merge);                             \
+  f(Swizzle);                           \
+  f(Swizzle2D);                         \
+  f(Resize);                            \
+  f(MatmulOp);                          \
+  f(LinearOp);                          \
+  f(SdpaFwdOp);                         \
+  f(SdpaBwdOp);                         \
+  f(EmbeddingFwdOp);                    \
+  f(Communication);                     \
+  f(ForLoop);                           \
   f(P2PCommunication);
 #define DISPATCH_FOR_ALL_KIR_EXPRS(f) \
   f(Allocate);                        \

csrc/ir/internal_base_nodes.h

@@ -458,7 +458,8 @@ class NVF_API TensorDomain : public Val {
       std::vector<IterDomain*> loop_domain,
       std::optional<std::vector<IterDomain*>> alternate_loop_domain,
       std::vector<std::optional<bool>> contiguity = {},
-      std::vector<IterDomain*> additional_ids = {});
+      std::vector<IterDomain*> additional_ids = {},
+      bool skip_checks = false);
 
   TensorDomain(IrBuilderPasskey, const TensorDomain* src);
 

csrc/ir/internal_nodes.h

@@ -3442,4 +3442,61 @@ class CutlassNvfp4GroupedMmaOp : public Expr {
   }
 };
 
+class GroupedBlockScalingFactorLayoutOp : public Expr {
+ public:
+  using Expr::Expr;
+
+  GroupedBlockScalingFactorLayoutOp(
+      IrBuilderPasskey,
+      Val* output,
+      Val* input,
+      Val* expert_offsets,
+      Val* sf_offsets,
+      BlockScalingFactorLayout layout,
+      Val* k,
+      Val* g);
+
+  NVFUSER_DECLARE_CLONE_AND_CREATE
+
+  const char* getOpString() const override {
+    return "GroupedBlockScalingFactorLayoutOp";
+  }
+
+  std::string toString(int indent_size = 0) const override;
+  std::string toInlineString(int indent_size = 0) const override;
+  std::vector<PolymorphicValue> evaluate(
+      const ExpressionEvaluator& ee,
+      const std::vector<PolymorphicValue>& inputs) const override;
+
+  // Get output block scaling factor
+  Val* out() const {
+    return output(0);
+  }
+
+  // Get input block scaling factor
+  Val* in() const {
+    return input(0);
+  }
+
+  TensorView* expertOffsets() const {
+    return input(1)->as<TensorView>();
+  }
+
+  TensorView* scalingFactorOffsets() const {
+    return input(2)->as<TensorView>();
+  }
+
+  Val* k() const {
+    return input(3);
+  }
+
+  Val* g() const {
+    return input(4);
+  }
+
+  BlockScalingFactorLayout layout() const {
+    return attribute<BlockScalingFactorLayout>(0);
+  }
+};
+
 } // namespace nvfuser

csrc/ir/nodes.cpp

@@ -3349,7 +3349,8 @@ TensorDomain::TensorDomain(
     std::vector<IterDomain*> loop_domain,
     std::optional<std::vector<IterDomain*>> alternate_loop_domain,
     std::vector<std::optional<bool>> contiguity,
-    std::vector<IterDomain*> additional_ids)
+    std::vector<IterDomain*> additional_ids,
+    bool skip_checks)
     : Val(passkey, ValType::TensorDomain, DataType::Null),
       root_domain_(std::move(root_domain)),
       logical_domain_(std::move(logical_domain)),
@@ -3366,18 +3367,21 @@ TensorDomain::TensorDomain(
   NVF_CHECK(
       loop_domain_.empty() == logical_domain_.empty(),
       "logical domain and loop domain can only be both empty or neither empty");
-  validateLoopDomain(logical_domain_, loop_domain_, additional_ids_);
-  if (!root_domain_.empty()) {
-    ir_utils::validateDomainEquivalence(
-        logical_domain_, root_domain_, additional_ids_);
-  }
-  if (!allocation_domain_.empty()) {
-    ir_utils::validateDomainEquivalence(
-        logical_domain_, allocation_domain_, additional_ids_);
-  }
-  if (alternate_loop_domain_.has_value()) {
-    validateLoopDomain(
-        logical_domain_, alternate_loop_domain_.value(), additional_ids_);
+
+  if (!skip_checks) {
+    validateLoopDomain(logical_domain_, loop_domain_, additional_ids_);
+    if (!root_domain_.empty()) {
+      ir_utils::validateDomainEquivalence(
+          logical_domain_, root_domain_, additional_ids_);
+    }
+    if (!allocation_domain_.empty()) {
+      ir_utils::validateDomainEquivalence(
+          logical_domain_, allocation_domain_, additional_ids_);
+    }
+    if (alternate_loop_domain_.has_value()) {
+      validateLoopDomain(
+          logical_domain_, alternate_loop_domain_.value(), additional_ids_);
+    }
   }
 
   // resetDomains initializes other member variables, required by clang-tidy
@@ -6551,4 +6555,60 @@ std::vector<PolymorphicValue> CutlassNvfp4GroupedMmaOp::evaluate(
 
 NVFUSER_DEFINE_CLONE_AND_CREATE(CutlassNvfp4GroupedMmaOp)
 
+GroupedBlockScalingFactorLayoutOp::GroupedBlockScalingFactorLayoutOp(
+    IrBuilderPasskey passkey,
+    Val* output,
+    Val* input,
+    Val* expert_offsets,
+    Val* sf_offsets,
+    BlockScalingFactorLayout layout,
+    Val* k,
+    Val* g)
+    : Expr(passkey) {
+  addInput(input);
+  addInput(expert_offsets);
+  addInput(sf_offsets);
+  addInput(k);
+  addInput(g);
+  addOutput(output);
+  addDataAttribute(layout);
+}
+
+std::string GroupedBlockScalingFactorLayoutOp::toString(int indent_size) const {
+  std::stringstream ss;
+  indent(ss, indent_size) << output(0)->toString() << "\n";
+  indent_size++;
+  indent(ss, indent_size) << " = grouped_block_scaling_factor_layout(\n";
+  indent_size++;
+  indent(ss, indent_size) << "input = " << in()->toString() << ",\n";
+  indent(ss, indent_size) << "expert_offsets = " << expertOffsets()->toString()
+                          << ",\n";
+  indent(ss, indent_size) << "sf_offsets = "
+                          << scalingFactorOffsets()->toString() << ",\n";
+  indent(ss, indent_size) << "layout = "
+                          << (layout() == BlockScalingFactorLayout::Block128x4
+                                  ? "Block128x4"
+                                  : "Unknown")
+                          << "\n";
+  indent_size--;
+  indent(ss, indent_size) << ")\n";
+  return ss.str();
+}
+
+std::string GroupedBlockScalingFactorLayoutOp::toInlineString(
+    int indent_size) const {
+  NVF_CHECK(
+      false, "GroupedBlockScalingFactorLayoutOp can not be printed inline");
+}
+
+std::vector<PolymorphicValue> GroupedBlockScalingFactorLayoutOp::evaluate(
+    const ExpressionEvaluator& ee,
+    const std::vector<PolymorphicValue>& inputs) const {
+  // This is a placeholder implementation - the actual implementation
+  // would depend on the specific block scaling factor layout operation
+  NVF_THROW("GroupedBlockScalingFactorLayoutOp evaluation not yet implemented");
+}
+
+NVFUSER_DEFINE_CLONE_AND_CREATE(GroupedBlockScalingFactorLayoutOp)
+
 } // namespace nvfuser

csrc/logical_domain_map.cpp

@@ -131,6 +131,17 @@ std::pair<std::unordered_set<IterDomain*>, bool> getNonMappingDomainInfo(
       non_mapping_ids.insert(producer_logical.at(topk_dim));
       has_consumer_id = true;
     }
+  } else if (
+      auto* grouped_block_sf_layout =
+          dynamic_cast<GroupedBlockScalingFactorLayoutOp*>(
+              consumer_tv->definition())) {
+    if (producer_tv != grouped_block_sf_layout->in()) {
+      auto producer_logical =
+          TensorDomain::noReductions(producer_tv->getLogicalDomain());
+      non_mapping_ids.insert(producer_logical.begin(), producer_logical.end());
+      // we are not mapping anything, `has_consumer_id` doesn't matter.
+      has_consumer_id = false;
+    }
   }
 
   return std::make_pair(non_mapping_ids, has_consumer_id);