Ensure flexibility in matmul definition

[zasdfgbnm]

This issue lists a few cases of valid matmul fusion definition that we are currently not supporting but eventually should support it. These cases are not prioritized, but when building out our segmenter and scheduler, we should keep them in mind to not contradict with them. Generally our design should support these cases without extra work. We should avoid having a long list of if statements to pattern these cases and special handling them.

[zasdfgbnm]

Example 1: Inputs are already broadcasted

TEST_F(NVFuserTest, SegmenterAlreadyBroadcasted_CUDA) {
  // skip until we have Hopper support
  NVFUSER_TEST_CUDA_ARCH_RANGE_GUARD(8, 0, 9, 0);
  const int M = 504, N = 136, K = 2048;

  auto fusion = std::make_unique<Fusion>();
  FusionGuard fg(fusion.get());

  auto a = makeContigConcreteTensor({M, K, 1}, DataType::Half);
  auto b = makeContigConcreteTensor({1, K, N}, DataType::Half);
  fusion->addInput(a);
  fusion->addInput(b);
  auto a_fp32 = castOp(DataType::Float, a);
  auto b_fp32 = castOp(DataType::Float, b);
  auto prod = mul(a_fp32, b_fp32);
  auto d = sum(prod, {1});
  fusion->addOutput(d);

  const auto fusion_layout = mma_utils::getMatmulLayout(fusion.get());
  TORCH_CHECK(
      fusion_layout.isValid(),
      "failed to get decide matmul layout through fusion definition");
  TORCH_CHECK(
      fusion_layout.getData() == MmaOptions::MmaLayout::TN,
      "mismatch between test layout (",
      toString(MmaOptions::MmaLayout::TN),
      ") and layout inferred from fusion definition (",
      toString(fusion_layout.getData()),
      ")");

  auto options = at::TensorOptions().dtype(at::kHalf).device(at::kCUDA, 0);
  at::Tensor t0 = at::randn({M, K, 1}, options);
  at::Tensor t1 = at::randn({1, K, N}, options);
  at::Tensor tref = t0.to(at::kFloat).mul(t1.to(at::kFloat)).sum({1});

  FusionExecutorCache executor_cache(std::move(fusion));

  auto outputs = executor_cache.runFusionWithInputs({t0, t1});

  TORCH_CHECK(
      !executor_cache.getMostRecentKernelRuntime()->isSegmented(),
      "fusion got segmented, expected to match whole fusion with single segment");

  TORCH_CHECK(
      isSchedulerInUse(
          executor_cache.getMostRecentKernelRuntime(),
          ScheduleHeuristic::Matmul),
      "matmul scheduler was not used to handle prepared fusion");

  TORCH_CHECK(outputs[0].allclose(tref, 0.001, 0.001));
}

[zasdfgbnm]

Example 2: AxA

TEST_F(NVFuserTest, SegmenterAxA_CUDA) {
  // skip until we have Hopper support
  NVFUSER_TEST_CUDA_ARCH_RANGE_GUARD(8, 0, 9, 0);
  const int K = 2048;

  auto fusion = std::make_unique<Fusion>();
  FusionGuard fg(fusion.get());

  auto a = makeContigConcreteTensor({K, K}, DataType::Half);
  fusion->addInput(a);
  auto a_broadcast = broadcast(a, {false, false, true});
  auto b_broadcast = broadcast(a, {true, false, false});
  auto a_fp32 = castOp(DataType::Float, a_broadcast);
  auto b_fp32 = castOp(DataType::Float, b_broadcast);
  auto prod = mul(a_fp32, b_fp32);
  auto d = sum(prod, {1});
  fusion->addOutput(d);

  const auto fusion_layout = mma_utils::getMatmulLayout(fusion.get());
  TORCH_CHECK(
      fusion_layout.isValid(),
      "failed to get decide matmul layout through fusion definition");
  TORCH_CHECK(
      fusion_layout.getData() == MmaOptions::MmaLayout::TN,
      "mismatch between test layout (",
      toString(MmaOptions::MmaLayout::TN),
      ") and layout inferred from fusion definition (",
      toString(fusion_layout.getData()),
      ")");

  auto options = at::TensorOptions().dtype(at::kHalf).device(at::kCUDA, 0);
  at::Tensor t0 = at::randn({K, K}, options);
  at::Tensor tref = t0.matmul(t0);

  FusionExecutorCache executor_cache(std::move(fusion));

  auto outputs = executor_cache.runFusionWithInputs({t0});

  TORCH_CHECK(
      !executor_cache.getMostRecentKernelRuntime()->isSegmented(),
      "fusion got segmented, expected to match whole fusion with single segment");

  TORCH_CHECK(
      isSchedulerInUse(
          executor_cache.getMostRecentKernelRuntime(),
          ScheduleHeuristic::Matmul),
      "matmul scheduler was not used to handle prepared fusion");

  TORCH_CHECK(outputs[0].allclose(tref, 0.001, 0.001));
}

[zasdfgbnm]

Example 3: Fusion generated inputs:

TEST_F(NVFuserTest, SegmenterGeneratedInput_CUDA) {
  // skip until we have Hopper support
  NVFUSER_TEST_CUDA_ARCH_RANGE_GUARD(8, 0, 9, 0);
  const int M = 504, N = 136, K = 2048;

  auto fusion = std::make_unique<Fusion>();
  FusionGuard fg(fusion.get());

  auto value = IrBuilder::newScalar(DataType::Double);
  fusion->addInput(value);
  auto a = full(
      {IrBuilder::newConstant(M, DataType::Int),
       IrBuilder::newConstant(K, DataType::Int)},
      value,
      DataType::Half);
  auto b = full(
      {IrBuilder::newConstant(K, DataType::Int),
       IrBuilder::newConstant(N, DataType::Int)},
      value,
      DataType::Half);
  auto a_broadcast = broadcast(a, {false, false, true});
  auto b_broadcast = broadcast(b, {true, false, false});
  auto a_fp32 = castOp(DataType::Float, a_broadcast);
  auto b_fp32 = castOp(DataType::Float, b_broadcast);
  auto prod = mul(a_fp32, b_fp32);
  auto d = sum(prod, {1});
  fusion->addOutput(d);

  const auto fusion_layout = mma_utils::getMatmulLayout(fusion.get());
  TORCH_CHECK(
      fusion_layout.isValid(),
      "failed to get decide matmul layout through fusion definition");
  TORCH_CHECK(
      fusion_layout.getData() == MmaOptions::MmaLayout::TN,
      "mismatch between test layout (",
      toString(MmaOptions::MmaLayout::TN),
      ") and layout inferred from fusion definition (",
      toString(fusion_layout.getData()),
      ")");

  auto options = at::TensorOptions().dtype(at::kHalf).device(at::kCUDA, 0);
  at::Tensor t0 = at::full({M, K}, 1.0, options);
  at::Tensor t1 = at::full({K, N}, 1.0, options);
  at::Tensor tref = t0.matmul(t1);

  FusionExecutorCache executor_cache(std::move(fusion));

  auto outputs = executor_cache.runFusionWithInputs({1.0});

  TORCH_CHECK(
      !executor_cache.getMostRecentKernelRuntime()->isSegmented(),
      "fusion got segmented, expected to match whole fusion with single segment");

  TORCH_CHECK(
      isSchedulerInUse(
          executor_cache.getMostRecentKernelRuntime(),
          ScheduleHeuristic::Matmul),
      "matmul scheduler was not used to handle prepared fusion");

  TORCH_CHECK(outputs[0].allclose(tref, 0.001, 0.001));
}

[jacobhinkle]

We do support Example 1, but Examples 2 and 3 are not yet supported.

[kevinstephano]

@jacobhinkle I see you recently looked at this issue. Is it still relevant?

[jacobhinkle]

@jacobhinkle I see you recently looked at this issue. Is it still relevant?

I do not know of a use case off hand for either of these types of matmuls in current models: specifically I mean A@A or A.T@A and the other case is something like randn() @ B where some operands are generated from factory methods. I am not sure if we put anything below P2 (maybe just close as "not planned"?) but this seemed like lowest priority to me unless we know of a model that is hitting this limitation. But maybe @zasdfgbnm has other motivation.

[zasdfgbnm]

I am not sure if we put anything below P2 (maybe just close as "not planned"?) but this seemed like lowest priority to me unless we know of a model that is hitting this limitation.

I agree with this part. I just opened this issue to track known limitations that we found during development.