LinCombPerRowBias & EVT Changes #152

joeatodd · 2024-11-08T15:55:49Z

This is a very rough first attempt to get this working. It compiles and runs, but the validation fails.

Testing issues:

The validation in the example doesn't do a per-row-bias so of course it fails...
Need to investigate using existing unit testing infrastructure

Impl issues:

Need to think more about the EpilogueTile shape (currently arbitrarily set inside CollectiveEpilogue)
Consider passing through an Alignment template arg which is more appropriate for XE

As a precursor to properly testing #152, I have adapted the gemm_testbed_3x.hpp infrastructure slightly for Xe and added a single test case for LinCombEltAct, as that's something we know should work. --------- Co-authored-by: Finlay <[email protected]>

Not totally sure these are correct yet, but it's compiling...

joeatodd · 2024-12-05T11:49:26Z

include/cutlass/gemm/kernel/xe_gemm.hpp

@@ -262,7 +262,7 @@ class GemmUniversal<
    CollectiveEpilogue epilogue{params.epilogue, shared_storage.epilogue};
    epilogue(
      problem_shape_MNKL,
-      subgroup_shape,
+      subgroup_shape, // TODO(joe): Inconsistency here w/ blk_coord_mnkl


@aacostadiaz the subgroup shape is being passed in here instead of the WorkgroupTileShape. However, the coordinate blk_coord_mnkl is the work-group coordinate. This doesn't matter for this PR because I'm computing things inside xe_epilogue.hpp but it feels inconsistent. Potentially something to address as part of your ongoing work on the coordinates?

Yes, that makes sense. I can look into that as part of the fix for coordinates.

maybe change to TODO(codeplay): ?

aacostadiaz

Great!!!

I left some minor comments. I'm not sure if the verification in the example is checking the right thing. Could you please double check that?

examples/sycl/pvc/pvc_gemm_with_per_row_bias.cpp

aacostadiaz · 2024-12-05T14:10:54Z

examples/sycl/pvc/pvc_gemm_with_per_row_bias.cpp

+    // Check if output from CUTLASS kernel and reference kernel are equal or not
+    bool passed = cutlass::reference::device::BlockCompareEqual(
+      block_ref_D.get(), block_D.get(), block_D.size());


Is this checking the output with the epilogue or just the output of GEMM?

examples/sycl/pvc/pvc_gemm_with_per_row_bias.cpp

aacostadiaz · 2024-12-05T14:16:29Z

include/cutlass/gemm/kernel/xe_gemm.hpp

@@ -262,7 +262,7 @@ class GemmUniversal<
    CollectiveEpilogue epilogue{params.epilogue, shared_storage.epilogue};
    epilogue(
      problem_shape_MNKL,
-      subgroup_shape,
+      subgroup_shape, // TODO(joe): Inconsistency here w/ blk_coord_mnkl


Yes, that makes sense. I can look into that as part of the fix for coordinates.

test/unit/gemm/device/xe_gemm_bf16_bf16_fp32_tensor_op_fp32_evt.cpp

FMarno · 2024-12-05T14:24:35Z

examples/sycl/pvc/pvc_gemm_with_per_row_bias.cpp

+  using ElementOutput = typename CollectiveEpilogue::ElementOutput;
+  using ElementCompute = typename CollectiveEpilogue::ElementCompute;
+  using ElementAccumulator = typename CollectiveEpilogue::ElementAccumulator;
+  using ElementBias = typename CollectiveEpilogue::ThreadEpilogueOp::ElementBias; //TODO(joe) Is this right?


Joe TODO.
I'm not sure but I don't see why not.

Removed - cheers

FMarno · 2024-12-05T14:24:59Z

examples/sycl/pvc/pvc_gemm_with_per_row_bias.cpp

+    // cutlass::reference::device::TensorScaleBiasGemm<
+    //   ElementAccumulator, typename B2bGemm::ElementC, typename B2bGemm::LayoutC,
+    //   ElementCompute, typename B2bGemm::LayoutC
+    // > (
+    //   {M, N, K},
+    //   ref_Z.at(i),
+    //   ref_ref_D0.at(i),
+    //   alpha0,
+    //   ref_Scale0.at(i),
+    //   ref_Bias0.at(i)
+    // );


FMarno · 2024-12-05T14:25:18Z

examples/sycl/pvc/pvc_gemm_with_per_row_bias.cpp

+
+    syclcompat::wait();
+
+    // TODO(joe): Add the right epilogue here (PerRowBias) for testing


FMarno · 2024-12-05T14:26:22Z

examples/sycl/pvc/pvc_gemm_with_per_row_bias.cpp

+    block_C.reset(M * N * L);
+    block_D.reset(M * N * L);
+    block_ref_D.reset(M * N * L);
+    bias.reset(N * L); //TODO(joe) or N?


examples/sycl/pvc/pvc_gemm_with_per_row_bias.cpp

FMarno · 2024-12-05T14:42:04Z

include/cutlass/gemm/kernel/xe_gemm.hpp

@@ -262,7 +262,7 @@ class GemmUniversal<
    CollectiveEpilogue epilogue{params.epilogue, shared_storage.epilogue};
    epilogue(
      problem_shape_MNKL,
-      subgroup_shape,
+      subgroup_shape, // TODO(joe): Inconsistency here w/ blk_coord_mnkl


maybe change to TODO(codeplay): ?

test/unit/gemm/device/xe_gemm_bf16_bf16_fp32_tensor_op_fp32_evt.cpp

include/cutlass/epilogue/collective/xe_epilogue.hpp

examples/sycl/pvc/pvc_gemm_with_per_row_bias.cpp

t4c1 · 2024-12-05T14:52:00Z

examples/sycl/pvc/pvc_gemm_with_per_row_bias.cpp

+    // cutlass::reference::device::TensorScaleBiasGemm<
+    //   ElementAccumulator, typename B2bGemm::ElementC, typename B2bGemm::LayoutC,
+    //   ElementCompute, typename B2bGemm::LayoutC
+    // > (
+    //   {M, N, K},
+    //   ref_Z.at(i),
+    //   ref_ref_D0.at(i),
+    //   alpha0,
+    //   ref_Scale0.at(i),
+    //   ref_Bias0.at(i)
+    // );


Commented code

t4c1 · 2024-12-05T14:52:29Z

examples/sycl/pvc/pvc_gemm_with_per_row_bias.cpp

+
+    syclcompat::wait();
+
+    // TODO(joe): Add the right epilogue here (PerRowBias) for testing


t4c1 · 2024-12-05T14:53:19Z

examples/sycl/pvc/pvc_gemm_with_per_row_bias.cpp

+    block_C.reset(M * N * L);
+    block_D.reset(M * N * L);
+    block_ref_D.reset(M * N * L);
+    bias.reset(N * L); //TODO(joe) or N?


examples/sycl/pvc/pvc_gemm_with_per_row_bias.cpp

include/cutlass/epilogue/collective/xe_epilogue.hpp

t4c1 · 2024-12-05T15:06:39Z

include/cutlass/epilogue/collective/xe_epilogue.hpp

@@ -341,7 +362,7 @@ class CollectiveEpilogue<
        auto acc_frag_mn = acc_frag(_, epi_m, epi_n);

        CUTLASS_PRAGMA_UNROLL
-        for (int epi_v = 0; epi_v < FragmentSize; ++epi_v) {
+        for (int epi_v = 0; epi_v < size(trD_frag); ++epi_v) {


Can you explain what is the point to the changes you are doing in this file? Are you changing to apply epilogue per subgroup instead of workgroup?

test/unit/gemm/device/xe_gemm_bf16_bf16_fp32_tensor_op_fp32_evt.cpp

FMarno · 2024-12-05T15:50:24Z

include/cutlass/epilogue/collective/xe_epilogue.hpp

    // Get the fusion callbacks
    constexpr bool RefSrc = true;
    auto residue_mn = make_coord(M, N);
    auto cst_args = cutlass::epilogue::fusion::detail::ConsumerStoreArgs{
                      problem_shape_mnkl,
-                      TileShapeMNK{},
-                      tile_coord_mnkl,
+                      SubgroupTileShape{},


I don't think SubgroupTileShape is quite the same as tile_shape_mnk, so maybe we should have a comment there. I guess sg_coord is there to match that change too.

Co-authored-by: Finlay <[email protected]>

joeatodd force-pushed the jtodd/per-row-bias branch 2 times, most recently from 4720c36 to bdcb4bb Compare November 19, 2024 12:09

joeatodd mentioned this pull request Nov 19, 2024

Unit test for IntelPVC w/ LinCombEltAct EVT #158

Merged

3 tasks

joeatodd force-pushed the jtodd/per-row-bias branch from bdcb4bb to a704294 Compare November 19, 2024 12:33

joeatodd added the incremental Incremental changes label Nov 19, 2024

joeatodd force-pushed the jtodd/per-row-bias branch from 8d7f032 to 600cbf6 Compare December 3, 2024 10:05

joeatodd added 12 commits December 4, 2024 15:48

Define FusionCallbacks for LinCombPerRowBias

5b04a28

PVC example with per-row bias

8e2319f

Add TODO to example

ecee600

Generalise args to ConsumerStoreArgs

f737c2a

Not totally sure these are correct yet, but it's compiling...

Unit test for LinCombPerRowBias (failing)

f4217ec

Add missing initialization distribution (AllZeros)

03f5c58

Simplified/corrected (?) the thread copy/residue stuff

eb0f1ad

Fix incorrect loop

187ee9e

Hardcode EpilogueTile for now

83850ab

Pass subgroup args to ConsumerStoreArgs

d7b568e

Update test fn name

b787872

Simplify ConsumerStoreCallbacks creation

b5db222

joeatodd force-pushed the jtodd/per-row-bias branch from 17df726 to b5db222 Compare December 4, 2024 15:49

joeatodd commented Dec 5, 2024

View reviewed changes

joeatodd marked this pull request as ready for review December 5, 2024 11:49

aacostadiaz reviewed Dec 5, 2024

View reviewed changes

FMarno reviewed Dec 5, 2024

View reviewed changes

t4c1 reviewed Dec 5, 2024

View reviewed changes

FMarno reviewed Dec 5, 2024

View reviewed changes

joeatodd removed the incremental Incremental changes label Dec 11, 2024

joeatodd and others added 3 commits December 11, 2024 12:41

Address some review feedback (formatting & TODOs)

307cdc0

Apply suggestions from code review

a445f17

Co-authored-by: Finlay <[email protected]>

More review feedback

2455bb1

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

LinCombPerRowBias & EVT Changes #152

LinCombPerRowBias & EVT Changes #152

joeatodd commented Nov 8, 2024

joeatodd Dec 5, 2024

aacostadiaz Dec 5, 2024

FMarno Dec 5, 2024

joeatodd Dec 11, 2024

aacostadiaz left a comment

aacostadiaz Dec 5, 2024

aacostadiaz Dec 5, 2024

FMarno Dec 5, 2024

joeatodd Dec 11, 2024

FMarno Dec 5, 2024

FMarno Dec 5, 2024

FMarno Dec 5, 2024

FMarno Dec 5, 2024

t4c1 Dec 5, 2024

t4c1 Dec 5, 2024

t4c1 Dec 5, 2024

t4c1 Dec 5, 2024

FMarno Dec 5, 2024


		syclcompat::wait();

		// TODO(joe): Add the right epilogue here (PerRowBias) for testing

LinCombPerRowBias & EVT Changes #152

Are you sure you want to change the base?

LinCombPerRowBias & EVT Changes #152

Conversation

joeatodd commented Nov 8, 2024

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

aacostadiaz left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment