refine code

Author: xiaolil1

CMakeLists.txt

@@ -28,8 +28,8 @@ set(CUDA_FILES csrc/ops.cu csrc/kernels.cu)
 set(HIP_FILES csrc/ops.hip csrc/kernels.hip)
 set(MPS_FILES csrc/mps_ops.mm)
 set(METAL_FILES csrc/mps_kernels.metal)
-#set(XPU_FILES csrc/xpu_ops.cpp csrc/xpu_kernels.cpp csrc/xpu_cutlass.cpp csrc/xpu_cutlass-cute.cpp csrc/xpu_cutlass_fusion.cpp)
-set(XPU_FILES csrc/xpu_ops.cpp csrc/xpu_kernels.cpp csrc/xpu_cutlass_fusion.cpp)
+set(XPU_FILES csrc/xpu_ops.cpp csrc/xpu_kernels.cpp csrc/xpu_cutlass.cpp csrc/xpu_cutlass-cute.cpp csrc/xpu_cutlass_fusion.cpp)
+#set(XPU_FILES csrc/xpu_ops.cpp csrc/xpu_kernels.cpp csrc/xpu_cutlass_fusion.cpp)
 # C++ sources are always included
 list(APPEND SRC_FILES ${CPP_FILES})
 
@@ -321,11 +321,13 @@ if(BUILD_XPU)
       -Xs
       -options "-cl-intel-enable-auto-large-GRF-mode -cl-poison-unsupported-fp64-kernels -cl-intel-greater-than-4GB-buffer-required"
     )
-    set(SYCL_COMPILE_FLAGS "-fsycl;-fhonor-nans;-fhonor-infinities;-fno-associative-math;-fno-approx-func;-fno-sycl-instrument-device-code;--offload-compress;-fsycl-targets=intel_gpu_pvc;-Xspirv-translator;-spirv-ext=+SPV_INTEL_split_barrier;")
+    set(SYCL_COMPILE_FLAGS "-fsycl;-fhonor-nans;-fhonor-infinities;-fno-associative-math;-fno-approx-func;-fno-sycl-instrument-device-code;--offload-compress;-fsycl-targets=intel_gpu_pvc;-Xspirv-translator;-spirv-ext=+SPV_INTEL_split_barrier,+SPV_INTEL_2d_block_io,+SPV_INTEL_subgroup_matrix_multiply_accumulate;")
 
     set_property(TARGET bitsandbytes PROPERTY CXX_STANDARD 20)
     target_compile_options(bitsandbytes PRIVATE ${SYCL_COMPILE_FLAGS})
     target_link_options(bitsandbytes PRIVATE ${SYCL_LINK_FLAGS})
+    #find_package(IntelSYCL REQUIRED)
+    #target_link_libraries(bitsandbytes PRIVATE Intel::SYCL OpenCL::OpenCL)
 
 endif()
 

csrc/xpu_cutlass-cute.cpp

@@ -63,9 +63,9 @@ using TiledMma =
     typename TiledMMAHelper<MMA_Atom<XE_8x16x16_F32BF16BF16F32_TT>, Layout<TileShape>,
                                   Layout<Shape<_8, _4, _1>, Stride<_4, _1, _0>>>::TiledMMA;
 
-using DispatchPolicy = MainloopIntelPVC<Stages>; //, KernelPVC /*Schedule*/>;
+using DispatchPolicy = MainloopIntelXeXMX16<Stages>; //, KernelPVC /*Schedule*/>;
 using EpilogueOp = cutlass::epilogue::fusion::LinearCombination<float /*data_type of GEMM output*/, ElementComputeEpilogue, ElementAccumulator, ElementAccumulator, cutlass::FloatRoundStyle::round_to_nearest>;
-using FusionCallBacks = cutlass::epilogue::fusion::FusionCallbacks<cutlass::epilogue::IntelPVCEpilogue, EpilogueOp, TileShape, decltype(tile_shape(TiledMma()))>;
+using FusionCallBacks = cutlass::epilogue::fusion::FusionCallbacks<cutlass::epilogue::IntelXeXMX16, EpilogueOp, TileShape, decltype(tile_shape(TiledMma()))>;
 using SharedStorage = FusionCallBacks::SharedStorage;
 
 using ClusterShape = typename DispatchPolicy::ClusterShape;
@@ -79,7 +79,7 @@ using ClusterShape = typename DispatchPolicy::ClusterShape;
   using TileSchedulerParams = typename TileScheduler::Params;
 
   using CollectiveEpilogue = cutlass::epilogue::collective::CollectiveEpilogue<
-          cutlass::epilogue::IntelPVCEpilogue,
+          cutlass::epilogue::IntelXeXMX16,
           TileShape,
           ElementAccumulator,
           cutlass::gemm::TagToStrideC_t<cutlass::layout::RowMajor>, // Convert CUTLASS 2.x to CUTLASS 3.x representation
@@ -280,8 +280,8 @@ class kgemv_4bit_inference_cutlass_cute {
     constexpr auto workgroup_shape = WorkgroupTileShape{}; 
     constexpr auto subgroup_shape = SubgroupTileShape{};                
 
-    Tensor mA_mkl = cute::get_pvc_tensor(make_shape(M,K,L));   //(m,k,l)
-    Tensor mB_nkl = cute::get_pvc_tensor(make_shape(N,K,L));   //(n,k,l)
+    Tensor mA_mkl = cute::get_xe_tensor(make_shape(M,K,L));   //(m,k,l)
+    Tensor mB_nkl = cute::get_xe_tensor(make_shape(N,K,L));   //(n,k,l)
 
     Tensor gA = local_tile(mA_mkl, select<0,2>(blk_shape), make_coord(m_coord,_,l_coord));
     Tensor gB = local_tile(mB_nkl, select<1,2>(blk_shape), make_coord(n_coord,_,l_coord));	

csrc/xpu_cutlass.cpp

@@ -27,10 +27,10 @@ void gemv_4bit_inference_cutlass(int m, int n, int k, T *A, T *B,
   // Create the Epilogue
   using EpilogueOp = cutlass::epilogue::fusion::LinearCombination<float /*data_type of GEMM output*/, ElementComputeEpilogue, ElementAccumulator, ElementAccumulator, cutlass::FloatRoundStyle::round_to_nearest>;
 
-  using FusionCallBacks = cutlass::epilogue::fusion::FusionCallbacks<cutlass::epilogue::IntelPVCEpilogue, EpilogueOp, TileShape,
+  using FusionCallBacks = cutlass::epilogue::fusion::FusionCallbacks<cutlass::epilogue::IntelXeXMX16, EpilogueOp, TileShape,
           decltype(tile_shape(TiledMma()))>;
   using CollectiveEpilogue = cutlass::epilogue::collective::CollectiveEpilogue<
-          cutlass::epilogue::IntelPVCEpilogue,
+          cutlass::epilogue::IntelXeXMX16,
           TileShape,
           ElementAccumulator,
           cutlass::gemm::TagToStrideC_t<cutlass::layout::RowMajor>, // Convert CUTLASS 2.x to CUTLASS 3.x representation
@@ -44,7 +44,7 @@ void gemv_4bit_inference_cutlass(int m, int n, int k, T *A, T *B,
 
   // GEMM Mainloop - iteration over blocks in K dimension
   using CollectiveMainloop = cutlass::gemm::collective::CollectiveMma<
-          cutlass::gemm::MainloopIntelPVC<2>, //use PipelineStages = 2
+          cutlass::gemm::MainloopIntelXeXMX16<2>, //use PipelineStages = 2
           TileShape,
           bfloat16_t, // data_type of input: A
           cutlass::gemm::TagToStrideA_t<cutlass::layout::RowMajor>, // Convert CUTLASS 2.x to CUTLASS 3.x representation

run_case.sh

@@ -28,7 +28,7 @@
 
 
 #gdb -args python -m pytest -vs tests/test_xpu.py::TestXPU::test_gemm_4bit
-#pytest -vs tests/test_xpu.py::TestXPU::test_gemm_4bit
-python tests/test_xpu_db.py
+pytest -vs tests/test_xpu.py::TestXPU::test_gemm_4bit
+#python tests/test_xpu_db.py
 #gdb -args python tests/test_xpu_db.py
 #pytest tests/test_functional.py::TestQuantize4BitFunctional::test_gemv_4bit[dim=256-uint8-bf16-fc1-nf4-DQ_True-xpu]

tests/test_xpu.py

@@ -93,8 +93,8 @@ def test_gemm_4bit(self, device, dim, dtype, storage_type, quant_storage, double
         #pdb.set_trace()
         diff = abs(C2-C3)
         print("diff = ", diff.sum())
-        print(C3[0])
-        print(C2[0])
+        #print(C3[0])
+        #print(C2[0])
         #print(C3)
         #print(C2)
         #A.requires_grad = True