gpu: jit: ir: fix 2d send alignment on Xe2

EMinsight · Feb 12, 2024 · b43f619 · b43f619
1 parent 763b1ae
commit b43f619
Show file tree

Hide file tree

Showing 4 changed files with 12 additions and 9 deletions.
diff --git a/src/gpu/jit/ir/block_2d_utils.hpp b/src/gpu/jit/ir/block_2d_utils.hpp
@@ -1,5 +1,5 @@
 /*******************************************************************************
-* Copyright 2022-2023 Intel Corporation
+* Copyright 2022-2024 Intel Corporation
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
@@ -33,8 +33,9 @@ inline int block_2d_base_alignment(const hw_t &hw) {
     return 64;
 }
 
-inline int block_2d_x_alignment(int type_size) {
-    return std::max(4, type_size) / type_size;
+inline int block_2d_x_alignment(const hw_t &hw, int type_size) {
+    int min_alignment = (hw == ngen::HW::Xe2) ? 16 : 4;
+    return std::max(min_alignment, type_size) / type_size;
 }
 
 inline bool block_2d_width_ok(int width, int type_size) {

diff --git a/src/gpu/jit/ir/message.cpp b/src/gpu/jit/ir/message.cpp
@@ -1,5 +1,5 @@
 /*******************************************************************************
-* Copyright 2021-2023 Intel Corporation
+* Copyright 2021-2024 Intel Corporation
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
@@ -794,8 +794,9 @@ bool access_builder_t::try_build_2d(send_params_t &send_params) {
         }
 
         if (!skip_send) {
-            if (!ir_ctx_->cset().can_prove(
-                        x % block_2d_x_alignment(send_type.size()) == 0)) {
+            auto x_alignment
+                    = block_2d_x_alignment(ir_ctx_->hw(), send_type.size());
+            if (!ir_ctx_->cset().can_prove(x % x_alignment == 0)) {
                 ok = false;
                 return;
             }

diff --git a/src/gpu/jit/ir/send_plan.cpp b/src/gpu/jit/ir/send_plan.cpp
@@ -1738,7 +1738,7 @@ class send_2d_helper_t {
             const tdim_info_t &w_tdim, const mod_info_t &mod_info) const {
         if (!params_.use_xy) return true;
         auto x_mod = w_tdim.offset(mod_info.vmods(), w_tdim.base_mod());
-        int align = block_2d_x_alignment(params_.type.size());
+        int align = block_2d_x_alignment(info_.hw(), params_.type.size());
         int x_align = align;
         if (!x_mod.is_divisible(x_align) != 0)
             return fail_2d("Unsupported x alignment: ", x_mod);

diff --git a/src/gpu/jit/v2/ir/send.hpp b/src/gpu/jit/v2/ir/send.hpp
@@ -18,6 +18,7 @@
 #define GPU_JIT_V2_IR_SEND_HPP
 
 #include "gpu/jit/ir/block_2d_utils.hpp"
+#include "gpu/jit/ir/message.hpp"
 #include "gpu/jit/v2/ir/plan_utils.hpp"
 #include "gpu/jit/v2/ir/reqs.hpp"
 #include "gpu/jit/v2/ir/tensor.hpp"
@@ -466,13 +467,13 @@ struct send_2d_desc_t {
     }
 
     bool is_supported(const view_t &view, const prover_t &prover) const {
-        if (w % block_2d_x_alignment(type.size()) != 0) return false;
+        if (w % block_2d_x_alignment(hw, type.size()) != 0) return false;
 
         auto &plane = view.plane();
         auto width_bytes = W * type.size();
         auto pitch_bytes = P * type.size();
         int base_align = block_2d_base_alignment(hw);
-        int x_align = block_2d_x_alignment(type.size());
+        int x_align = block_2d_x_alignment(hw, type.size());
         if (!prover.prove(width_bytes >= 64)) return false;
         if (!prover.prove(width_bytes <= (1 << 24))) return false;
         if (!prover.prove(width_bytes % std::max(4, type.size()) == 0))