From c32b424b981ec5b45003050e77e0232bd9722f31 Mon Sep 17 00:00:00 2001
From: Frederik Engels <frederik.engels@codeplay.com>
Date: Tue, 25 Jan 2022 17:51:29 +0000
Subject: [PATCH 01/21] recognize the patterns used for the RHS matrix

---
 src/ExtractTileOperations.cpp | 109 ++++++++++++++++++++++++++++++++++
 1 file changed, 109 insertions(+)

diff --git a/src/ExtractTileOperations.cpp b/src/ExtractTileOperations.cpp
index 54a32d756afd..4e9797764726 100644
--- a/src/ExtractTileOperations.cpp
+++ b/src/ExtractTileOperations.cpp
@@ -39,6 +39,18 @@ Type amx_op_type_result_type(AMXOpType op_ty) {
     }
 }
 
+int amx_op_type_size(AMXOpType op_ty) {
+    switch (op_ty) {
+    case AMXOpType::Int8:
+        return 1;
+    case AMXOpType::Bfloat16:
+        return 2;
+    default:
+        internal_error << "Unexpected";
+        return -1;
+    }
+}
+
 const auto wild_i32 = Variable::make(Int(32), "*");
 const auto wild_i32x = Variable::make(Int(32, 0), "*");
 
@@ -143,6 +155,103 @@ Tile<3> get_3d_tile_index(const Expr &e) {
     return {true, base, {x_stride, 0, r_stride}, {x_tile, y_tile, r_tile}};
 }
 
+Tile<3> get_3d_rhs_tile_index(const Expr &e, int element_width) {
+    const auto *sub = e.as<Sub>();
+    const Add *add_lhs = nullptr;
+
+    // there's not always a sub pattern
+    if (!sub) {
+        add_lhs = e.as<Add>();
+    } else {
+        add_lhs = sub->a.as<Add>();
+    }
+
+    if (!add_lhs) {
+        return {};
+    }
+
+    // (([x*y](ramp(0, 1, r) / [x*y*r](4)) + [x*y*r](rro*4)) * [x*y*r](rhs.stride.2)
+    const Mul *mul = add_lhs->a.as<Mul>();
+
+    if (!mul) {
+        return {};
+    }
+
+    // obtain the x, y, r dimensions
+    const Add *dim_expr = add_lhs->b.as<Add>();
+
+    const Broadcast *y_bc = dim_expr->b.as<Broadcast>();
+
+    int tile_y = y_bc->lanes;
+
+    const Mod *mod = dim_expr->a.as<Mod>();
+
+    const Broadcast *bc_ramp = mod->a.as<Broadcast>();
+
+    if (!bc_ramp) {
+        return {};
+    }
+
+    int tile_xy = bc_ramp->lanes;
+    int tile_x = tile_xy / tile_y;
+
+    const Ramp *r_ramp = bc_ramp->value.as<Ramp>();
+
+    if (!r_ramp) {
+        return {};
+    }
+
+    int tile_r = r_ramp->lanes;
+
+    // [x*y*r](rhs.stride.2)
+    const Broadcast *stride_bc = mul->b.as<Broadcast>();
+
+    if (!stride_bc) {
+        return {};
+    }
+
+    // get the base and stride
+    const Broadcast *base_stride_bc = add_lhs->b.as<Add>()->b.as<Broadcast>();
+
+    if (!base_stride_bc) {
+        return {};
+    }
+
+    const Ramp *base_stride_ramp = base_stride_bc->value.as<Ramp>();
+
+    if (!base_stride_ramp) {
+        return {};
+    }
+
+    Expr base = base_stride_ramp->base.as<Broadcast>()->value;
+    Expr stride;
+
+    bool found_stride = false;
+
+    // this stride pattern can occur if `tile_r` is the same size as `acc`
+    auto stride_pattern = Broadcast::make(Ramp::make(0, 1, tile_r), tile_x * tile_y) / Broadcast::make((4 / element_width), tile_x * tile_y * tile_r) * Broadcast::make(wild_i32, tile_x * tile_y * tile_r);
+
+    std::vector<Expr> results{};
+    if (expr_match(stride_pattern, add_lhs->a, results)) {
+        found_stride = true;
+        stride = std::move(results[0]);
+    }
+
+    if (!found_stride) {
+        stride_pattern = (Broadcast::make(Ramp::make(0, 1, tile_r), tile_x * tile_y) / Broadcast::make((4/ element_width), tile_x * tile_y * tile_r) + wild_i32) * Broadcast::make(wild_i32, tile_x * tile_y * tile_r);
+        if (expr_match(stride_pattern, add_lhs->a, results)) {
+            found_stride = true;
+            stride = std::move(results[1]);
+            base = std::move(results[0]) * stride + base;
+        }
+    }
+
+    if (!found_stride) {
+        return {};
+    }
+
+    return {true, base, {stride, 0, 0}, {tile_x, tile_y, tile_r}};
+}
 struct Matmul {
     bool result = false;
     Stmt stmt;

From 43931e96111dd590aa6159941812890037cfbf12 Mon Sep 17 00:00:00 2001
From: Frederik Engels <frederik.engels@codeplay.com>
Date: Tue, 25 Jan 2022 18:05:56 +0000
Subject: [PATCH 02/21] make 1d tile matcher more robust

---
 src/ExtractTileOperations.cpp | 11 ++++++++++-
 1 file changed, 10 insertions(+), 1 deletion(-)

diff --git a/src/ExtractTileOperations.cpp b/src/ExtractTileOperations.cpp
index 4e9797764726..0365bdd832dd 100644
--- a/src/ExtractTileOperations.cpp
+++ b/src/ExtractTileOperations.cpp
@@ -56,7 +56,16 @@ const auto wild_i32x = Variable::make(Int(32, 0), "*");
 
 Tile<1> get_1d_tile_index(const Expr &e) {
     if (const auto *r1 = e.as<Ramp>()) {
-        return {true, r1->base, {r1->stride}, {r1->lanes}};
+        Expr pattern = ((wild_i32 * wild_i32) + wild_i32) * wild_i32;
+
+        std::vector<Expr> results;
+        if (!expr_match(pattern, r1->base, results)) {
+            return {};
+        }
+
+        auto stride = std::move(results[1]);
+
+        return {true, std::move(r1->base), {std::move(stride)}, {r1->lanes}};
     }
 
     return {};

From f4a153ab0680990ecd76eff8f673f55adeb452fa Mon Sep 17 00:00:00 2001
From: Frederik Engels <frederik.engels@codeplay.com>
Date: Tue, 25 Jan 2022 18:08:12 +0000
Subject: [PATCH 03/21] put getting rhs tile's index into a separate func

---
 src/ExtractTileOperations.cpp | 85 ++++++++++++++++++++++++-----------
 1 file changed, 60 insertions(+), 25 deletions(-)

diff --git a/src/ExtractTileOperations.cpp b/src/ExtractTileOperations.cpp
index 0365bdd832dd..132e5c79cd21 100644
--- a/src/ExtractTileOperations.cpp
+++ b/src/ExtractTileOperations.cpp
@@ -261,6 +261,43 @@ Tile<3> get_3d_rhs_tile_index(const Expr &e, int element_width) {
 
     return {true, base, {stride, 0, 0}, {tile_x, tile_y, tile_r}};
 }
+
+struct BaseStride {
+    bool result{false};
+    Expr base{};
+    Expr stride{};
+};
+
+BaseStride get_rhs_tile_index(const Expr &index, int element_width, int tile_x, int tile_y, int tile_r) {
+    const auto rhs_tile2 = get_2d_tile_index(index);
+
+    if (!rhs_tile2.result) {
+        const auto rhs_tile1 = get_1d_tile_index(index);
+
+        if (!rhs_tile1.result) {
+            auto rhs_tile3 = get_3d_rhs_tile_index(index, element_width);
+            if (rhs_tile3.extent[0] != tile_x || rhs_tile3.extent[1] != tile_y || rhs_tile3.extent[2] != tile_r) {
+                return {};
+            }
+
+            return {true, rhs_tile3.base, rhs_tile3.stride[0] * element_width};
+        } else {
+            if (rhs_tile1.extent[0] != tile_y * tile_r) {
+                return {};
+            }
+
+            // times 4 because of the rhs layout
+            return {true, rhs_tile1.base, rhs_tile1.stride[0] * (4 / element_width)};
+        }
+    } else {
+        if (tile_y != rhs_tile2.extent[0] || tile_r != rhs_tile2.extent[1]) {
+            return {};
+        }
+
+        return {true, rhs_tile2.base, rhs_tile2.stride[0]};
+    }
+}
+
 struct Matmul {
     bool result = false;
     Stmt stmt;
@@ -315,14 +352,24 @@ Matmul convert_to_matmul(const Store *op, const string &new_name, AMXOpType op_t
 
     const auto *lhs_load = matches[0].as<Load>();
     const auto *rhs_broadcast = matches[1].as<Broadcast>();
-    if (!lhs_load || !rhs_broadcast) {
+
+    const Cast *rhs_cast = nullptr;
+
+    if (lhs_load && !rhs_broadcast) {
+        // now working on a larger k dimension
+        rhs_cast = matches[1].as<Cast>();
+    } else {
+        rhs_cast = rhs_broadcast->value.as<Cast>();
+    }
+
+    if (!lhs_load || !rhs_cast) {
         return {};
     }
-    const auto *rhs_cast = rhs_broadcast->value.as<Cast>();
+
     if (rhs_cast) {
         if (op_type == AMXOpType::Int8) {
             if (!(rhs_cast->value.type().element_of() == Int(8) || rhs_cast->value.type().element_of() == UInt(8))) {
-                user_assert(false) << "Expected rhs cast of i8/u8";
+                user_error << "Expected rhs cast of i8/u8, got " << rhs_cast->value.type();
             }
         } else {  // AMXOpType::Bfloat16
             user_assert(rhs_cast->value.type().element_of() == BFloat(16)) << "Expected rhs cast of bf16";
@@ -350,29 +397,15 @@ Matmul convert_to_matmul(const Store *op, const string &new_name, AMXOpType op_t
     Expr rhs_base;
     Expr rhs_stride;
 
-    const auto rhs_tile2 = get_2d_tile_index(rhs_load->index);
-    if (!rhs_tile2.result) {
-        const auto rhs_tile1 = get_1d_tile_index(rhs_load->index);
+    auto opt_base_stride = get_rhs_tile_index(rhs_load->index, amx_op_type_size(op_type), tile_x, tile_y, tile_r);
 
-        if (!rhs_tile1.result) {
-            return {};
-        }
-
-        if (rhs_tile1.extent[0] != tile_y * tile_r) {
-            return {};
-        }
-
-        rhs_base = rhs_tile1.base;
-        rhs_stride = rhs_tile1.stride[0];
-    } else {
-        if (tile_y != rhs_tile2.extent[0] || tile_r != rhs_tile2.extent[1]) {
-            return {};
-        }
-
-        rhs_base = rhs_tile2.base;
-        rhs_stride = rhs_tile2.stride[0];
+    if (!opt_base_stride.result) {
+        return {};
     }
 
+    rhs_base = opt_base_stride.base;
+    rhs_stride = opt_base_stride.stride;
+
     if (op->index.type().lanes() != tile_x * tile_y ||
         factor != tile_r) {
         return {};
@@ -388,7 +421,8 @@ Matmul convert_to_matmul(const Store *op, const string &new_name, AMXOpType op_t
     auto rhs_var = Variable::make(Handle(), rhs_load->name);
     const auto &rhs_load_type = rhs_load->type;
     auto rhs_type = rhs_load_type.with_lanes(1024 / element_width);
-    auto rhs = Call::make(rhs_type, "tile_load", {1, tile_y * tile_r * element_width, rhs_var, rhs_base * element_width, rhs_stride * tile_y * element_width}, Call::Intrinsic);
+
+    auto rhs = Call::make(rhs_type, "tile_load", {tile_r / (4 / element_width), tile_y * 4, rhs_var, rhs_base * element_width, rhs_stride}, Call::Intrinsic);
     auto res_type = amx_op_type_result_type(op_type);
 
     // {rows, colbytes, acc, out, lhs, rhs}
@@ -528,6 +562,7 @@ class ExtractTileOperations : public IRMutator {
             found_tile_x = matmul.tile_x;
             found_tile_y = matmul.tile_y;
             found_tile_r = matmul.tile_r;
+
             return matmul.stmt;
         }
 
@@ -542,7 +577,7 @@ class ExtractTileOperations : public IRMutator {
         }
 
         // Otherwise there is some other operation using the allocation, so we cannot use the AMX instructions
-        user_assert(false) << "Found non-tile operations for AMX tile allocation";
+        user_error << "Found non-tile operations for AMX tile allocation";
         return op;
     }
 };

From 38a8e96a2cbd2ac0dd3ffdc75ca83b9c3f8d2814 Mon Sep 17 00:00:00 2001
From: Frederik Engels <frederik.engels@codeplay.com>
Date: Tue, 25 Jan 2022 18:08:45 +0000
Subject: [PATCH 04/21] expand the tests used in correctness check

---
 test/correctness/tiled_matmul.cpp | 127 ++++++++++++++++++------------
 1 file changed, 77 insertions(+), 50 deletions(-)

diff --git a/test/correctness/tiled_matmul.cpp b/test/correctness/tiled_matmul.cpp
index 0d3c924fc384..d9194759b3ab 100644
--- a/test/correctness/tiled_matmul.cpp
+++ b/test/correctness/tiled_matmul.cpp
@@ -50,24 +50,45 @@ bool equal_eps(float lhs, float rhs, float eps) {
 
 struct make_uint_t {
     template<typename... Args>
-    Type operator()(Args &&...args) const {
+    Type operator()(Args &&... args) const {
         return UInt(static_cast<Args &&>(args)...);
     }
 };
 
 struct make_int_t {
     template<typename... Args>
-    Type operator()(Args &&...args) const {
+    Type operator()(Args &&... args) const {
         return Int(static_cast<Args &&>(args)...);
     }
 };
 
-template<typename LhsInt8, typename RhsInt8>
-bool matmul() {
-    constexpr int row = 16;
-    constexpr int col = 16;
-    constexpr int acc = 16;
+template<typename T>
+void print_mat(const Buffer<T> &buf, int rows, int cols) {
+    using cast_T = std::conditional_t<std::is_integral_v<T>, int, T>;
+    for (int j = 0; j != rows; ++j) {
+        for (int i = 0; i != cols; ++i) {
+            std::cout << static_cast<cast_T>(buf(i, j)) << " ";
+        }
+        std::cout << std::endl;
+    }
+}
+
+template<typename T>
+void print_mat_rhs(const Buffer<T> &buf, int rows, int cols) {
+    using cast_T = std::conditional_t<std::is_integral_v<T>, int, T>;
+    for (int j = 0; j != (rows / (4 / sizeof(T))); ++j) {
+        for (int k = 0; k != (4 / sizeof(T)); ++k) {
+            for (int i = 0; i != cols; ++i) {
+                std::cout << static_cast<cast_T>(buf(k, i, j)) << " ";
+            }
+
+            std::cout << std::endl;
+        }
+    }
+}
 
+template<typename LhsInt8, typename RhsInt8>
+bool matmul(int row, int col, int acc, int tile_x, int tile_y, int tile_r) {
     Buffer<LhsInt8> A_buf(acc, row);
     Buffer<RhsInt8> B_buf(4, col, acc / 4);
 
@@ -78,10 +99,6 @@ bool matmul() {
     mm(x, y) = cast<int32_t>(0);
     mm(x, y) += cast<int32_t>(A_buf(r, y)) * cast<int32_t>(B_buf(r % 4, x, r / 4));
 
-    constexpr int tile_x = 8;
-    constexpr int tile_y = 8;
-    constexpr int tile_r = 4;
-
     Var rxi("rxi"), ryi("ryi");
     RVar rri("rri"), rro("rro");
 
@@ -118,6 +135,15 @@ bool matmul() {
 
     result.realize(out);
 
+    // uncomment to check the matrices
+    // std::cout << "Matrix A\n";
+    // print_mat(A_buf, row, acc);
+    // std::cout << "Matrix B\n";
+    // print_mat_rhs(B_buf, acc, col);
+
+    // std::cout << "result\n";
+    // print_mat(out, row, col);
+
     for (int j = 0; j < row; ++j) {
         for (int i = 0; i < col; ++i) {
             int32_t val = 0;
@@ -126,21 +152,18 @@ bool matmul() {
             }
             if (val != out(i, j)) {
                 std::cerr << "Invalid result at " << i << ", " << j << "\n"
-                          << out(i, j) << " != " << val << "\n";
+                          << out(i, j) << " != " << val << "\n"
+                          << "Matrix dims: " << row << "x" << col << "x" << acc << "\nTile dims: " << tile_x << "x" << tile_y << "x" << tile_r << "\n";
                 return false;
             }
         }
     }
 
+    std::cout << "Success\n";
     return true;
 }
 
-bool matmul_bf16() {
-    // lhs: 32x16, rhs: 16x32
-    const int row = 32;
-    const int col = 32;
-    const int acc = 16;
-
+bool matmul_bf16(int row, int col, int acc, int tile_x, int tile_y, int tile_r) {
     Var x("x"), y("y");
     Buffer<bfloat16_t> A(acc, row);
     Buffer<bfloat16_t> B(2, col, acc / 2);
@@ -151,10 +174,6 @@ bool matmul_bf16() {
     mm(x, y) = cast<float>(0);
     mm(x, y) += cast<float>(cast<float>(A(r.x, y))) * cast<float>(B(r.x % 2, x, r.x / 2));
 
-    int tile_x = 8;
-    int tile_y = 8;
-    int tile_r = 2;
-
     Var rxi("rxi"), ryi("ryi");
     RVar rri("rri"), rro("rro");
 
@@ -195,20 +214,31 @@ bool matmul_bf16() {
 
     result.realize(out);
 
+    // uncomment to check the matrices
+    // std::cout << "Matrix A\n";
+    // print_mat(A, row, acc);
+    // std::cout << "Matrix B\n";
+    // print_mat_rhs(B, acc, col);
+
+    // std::cout << "result\n";
+    // print_mat(out, row, col);
+
     for (int j = 0; j < row; ++j) {
         for (int i = 0; i < col; ++i) {
             float val = 0.f;
             for (int k = 0; k < acc; ++k) {
                 val += static_cast<float>(A(k, j)) * static_cast<float>(B(k % 2, i, k / 2));
             }
-            if (!equal_eps(val, out(i, j), 0.01f)) {
+            if (!equal_eps(val, out(i, j), 0.03f)) {
                 std::cerr << "Invalid result at " << i << ", " << j << "\n"
-                          << out(i, j) << " != " << val << "\n";
+                          << out(i, j) << " != " << val << "\n"
+                          << "Matrix dims: " << row << "x" << col << "x" << acc << "\nTile dims: " << tile_x << "x" << tile_y << "x" << tile_r << "\n";
                 return false;
             }
         }
     }
 
+    std::cout << "Success!\n";
     return true;
 }
 
@@ -217,6 +247,10 @@ auto matmul_us = &matmul<uint8_t, int8_t>;
 auto matmul_su = &matmul<int8_t, uint8_t>;
 auto matmul_uu = &matmul<uint8_t, uint8_t>;
 
+bool run_tests(bool (*fn)(int, int, int, int, int, int), int element_width) {
+    return fn(2, 2, 16, 2, 2, 8 / element_width) && fn(4, 4, 8, 4, 4, 8 / element_width) && fn(32, 32, 32, 8, 8, 8 / element_width) && fn(32, 32, 32, 8, 8, 4 / element_width);
+}
+
 int main(int argc, char **argv) {
     Target t = get_jit_target_from_environment();
     if (!t.has_feature(Target::AVX512_SapphireRapids)) {
@@ -225,38 +259,31 @@ int main(int argc, char **argv) {
     }
 
     printf("Running AMX matmul (signed/signed)\n");
-    if (!matmul_ss()) {
+    if (!run_tests(matmul_ss, 1)) {
         return -1;
-    } else {
-        printf("Success!\n");
     }
 
-    printf("Running AMX matmul (signed/unsigned)\n");
-    if (!matmul_su()) {
-        return -1;
-    } else {
-        printf("Success!\n");
-    }
+    // llvm >= 13.0 is required for unsigned and float AMX instructions
+    if (Halide::Internal::get_llvm_version() >= 130) {
+        printf("Running AMX matmul (signed/unsigned)\n");
+        if (!run_tests(matmul_su, 1)) {
+            return -1;
+        }
 
-    printf("Running AMX matmul (unsigned/signed)\n");
-    if (!matmul_us()) {
-        return -1;
-    } else {
-        printf("Success!\n");
-    }
+        printf("Running AMX matmul (unsigned/signed)\n");
+        if (!run_tests(matmul_us, 1)) {
+            return -1;
+        }
 
-    printf("Running AMX matmul (unsigned/unsigned)\n");
-    if (!matmul_uu()) {
-        return -1;
-    } else {
-        printf("Success!\n");
-    }
+        printf("Running AMX matmul (unsigned/unsigned)\n");
+        if (!run_tests(matmul_uu, 1)) {
+            return -1;
+        }
 
-    printf("Running AMX matmul (bf16)\n");
-    if (!matmul_bf16()) {
-        return -1;
-    } else {
-        printf("Success!\n");
+        printf("Running AMX matmul (bf16)\n");
+        if (!run_tests(matmul_bf16, 2)) {
+            return -1;
+        }
     }
     return 0;
 }
\ No newline at end of file

From cbdc68fdc9bd99e7fd8ce73bed5b5e6c5d1eec77 Mon Sep 17 00:00:00 2001
From: Frederik Engels <frederik.engels@codeplay.com>
Date: Tue, 25 Jan 2022 18:16:25 +0000
Subject: [PATCH 05/21] add exclamation mark

---
 test/correctness/tiled_matmul.cpp | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/test/correctness/tiled_matmul.cpp b/test/correctness/tiled_matmul.cpp
index d9194759b3ab..4fd40c3536e6 100644
--- a/test/correctness/tiled_matmul.cpp
+++ b/test/correctness/tiled_matmul.cpp
@@ -159,7 +159,7 @@ bool matmul(int row, int col, int acc, int tile_x, int tile_y, int tile_r) {
         }
     }
 
-    std::cout << "Success\n";
+    std::cout << "Success!\n";
     return true;
 }
 

From 0e3998ef75939720ebd3529fb2c388d8a3eb1624 Mon Sep 17 00:00:00 2001
From: Frederik Engels <frederik.engels@codeplay.com>
Date: Wed, 26 Jan 2022 14:42:17 +0000
Subject: [PATCH 06/21] remove unused vars

---
 src/ExtractTileOperations.cpp | 14 --------------
 1 file changed, 14 deletions(-)

diff --git a/src/ExtractTileOperations.cpp b/src/ExtractTileOperations.cpp
index 132e5c79cd21..8f31a9a177b5 100644
--- a/src/ExtractTileOperations.cpp
+++ b/src/ExtractTileOperations.cpp
@@ -179,13 +179,6 @@ Tile<3> get_3d_rhs_tile_index(const Expr &e, int element_width) {
         return {};
     }
 
-    // (([x*y](ramp(0, 1, r) / [x*y*r](4)) + [x*y*r](rro*4)) * [x*y*r](rhs.stride.2)
-    const Mul *mul = add_lhs->a.as<Mul>();
-
-    if (!mul) {
-        return {};
-    }
-
     // obtain the x, y, r dimensions
     const Add *dim_expr = add_lhs->b.as<Add>();
 
@@ -212,13 +205,6 @@ Tile<3> get_3d_rhs_tile_index(const Expr &e, int element_width) {
 
     int tile_r = r_ramp->lanes;
 
-    // [x*y*r](rhs.stride.2)
-    const Broadcast *stride_bc = mul->b.as<Broadcast>();
-
-    if (!stride_bc) {
-        return {};
-    }
-
     // get the base and stride
     const Broadcast *base_stride_bc = add_lhs->b.as<Add>()->b.as<Broadcast>();
 

From 464c2400080613b3487eb1075543be14be61405d Mon Sep 17 00:00:00 2001
From: Frederik Engels <frederik.engels@codeplay.com>
Date: Wed, 26 Jan 2022 15:17:54 +0000
Subject: [PATCH 07/21] run format and tidy

---
 src/ExtractTileOperations.cpp     | 4 ++--
 test/correctness/tiled_matmul.cpp | 4 ++--
 2 files changed, 4 insertions(+), 4 deletions(-)

diff --git a/src/ExtractTileOperations.cpp b/src/ExtractTileOperations.cpp
index 8f31a9a177b5..93a8ab96a600 100644
--- a/src/ExtractTileOperations.cpp
+++ b/src/ExtractTileOperations.cpp
@@ -65,7 +65,7 @@ Tile<1> get_1d_tile_index(const Expr &e) {
 
         auto stride = std::move(results[1]);
 
-        return {true, std::move(r1->base), {std::move(stride)}, {r1->lanes}};
+        return {true, r1->base, {std::move(stride)}, {r1->lanes}};
     }
 
     return {};
@@ -233,7 +233,7 @@ Tile<3> get_3d_rhs_tile_index(const Expr &e, int element_width) {
     }
 
     if (!found_stride) {
-        stride_pattern = (Broadcast::make(Ramp::make(0, 1, tile_r), tile_x * tile_y) / Broadcast::make((4/ element_width), tile_x * tile_y * tile_r) + wild_i32) * Broadcast::make(wild_i32, tile_x * tile_y * tile_r);
+        stride_pattern = (Broadcast::make(Ramp::make(0, 1, tile_r), tile_x * tile_y) / Broadcast::make((4 / element_width), tile_x * tile_y * tile_r) + wild_i32) * Broadcast::make(wild_i32, tile_x * tile_y * tile_r);
         if (expr_match(stride_pattern, add_lhs->a, results)) {
             found_stride = true;
             stride = std::move(results[1]);
diff --git a/test/correctness/tiled_matmul.cpp b/test/correctness/tiled_matmul.cpp
index 4fd40c3536e6..793c816b778e 100644
--- a/test/correctness/tiled_matmul.cpp
+++ b/test/correctness/tiled_matmul.cpp
@@ -50,14 +50,14 @@ bool equal_eps(float lhs, float rhs, float eps) {
 
 struct make_uint_t {
     template<typename... Args>
-    Type operator()(Args &&... args) const {
+    Type operator()(Args &&...args) const {
         return UInt(static_cast<Args &&>(args)...);
     }
 };
 
 struct make_int_t {
     template<typename... Args>
-    Type operator()(Args &&... args) const {
+    Type operator()(Args &&...args) const {
         return Int(static_cast<Args &&>(args)...);
     }
 };

From c174b981ee3e8ab0f3bdeedeedc3e0beedfa5a48 Mon Sep 17 00:00:00 2001
From: Frederik Engels <frederik.engels@codeplay.com>
Date: Fri, 28 Jan 2022 15:41:54 +0000
Subject: [PATCH 08/21] check for null before using IR in the next step

---
 src/ExtractTileOperations.cpp | 22 ++++++++++++++--------
 1 file changed, 14 insertions(+), 8 deletions(-)

diff --git a/src/ExtractTileOperations.cpp b/src/ExtractTileOperations.cpp
index 93a8ab96a600..7b719b544305 100644
--- a/src/ExtractTileOperations.cpp
+++ b/src/ExtractTileOperations.cpp
@@ -182,12 +182,24 @@ Tile<3> get_3d_rhs_tile_index(const Expr &e, int element_width) {
     // obtain the x, y, r dimensions
     const Add *dim_expr = add_lhs->b.as<Add>();
 
-    const Broadcast *y_bc = dim_expr->b.as<Broadcast>();
+    if (!dim_expr) {
+        return {};
+    }
+
+    const Broadcast *base_stride_bc = dim_expr->b.as<Broadcast>();
+
+    if (!base_stride_bc) {
+        return {};
+    }
 
-    int tile_y = y_bc->lanes;
+    int tile_y = base_stride_bc->lanes;
 
     const Mod *mod = dim_expr->a.as<Mod>();
 
+    if (!mod) {
+        return {};
+    }
+
     const Broadcast *bc_ramp = mod->a.as<Broadcast>();
 
     if (!bc_ramp) {
@@ -206,12 +218,6 @@ Tile<3> get_3d_rhs_tile_index(const Expr &e, int element_width) {
     int tile_r = r_ramp->lanes;
 
     // get the base and stride
-    const Broadcast *base_stride_bc = add_lhs->b.as<Add>()->b.as<Broadcast>();
-
-    if (!base_stride_bc) {
-        return {};
-    }
-
     const Ramp *base_stride_ramp = base_stride_bc->value.as<Ramp>();
 
     if (!base_stride_ramp) {

From f9086faf178f714eae647f397404f613f1d8e9e8 Mon Sep 17 00:00:00 2001
From: Frederik Engels <frederik.engels@codeplay.com>
Date: Fri, 28 Jan 2022 15:44:13 +0000
Subject: [PATCH 09/21] check if the broadcast was found

---
 src/ExtractTileOperations.cpp | 8 +++++++-
 1 file changed, 7 insertions(+), 1 deletion(-)

diff --git a/src/ExtractTileOperations.cpp b/src/ExtractTileOperations.cpp
index 7b719b544305..9ee2dab18531 100644
--- a/src/ExtractTileOperations.cpp
+++ b/src/ExtractTileOperations.cpp
@@ -224,7 +224,13 @@ Tile<3> get_3d_rhs_tile_index(const Expr &e, int element_width) {
         return {};
     }
 
-    Expr base = base_stride_ramp->base.as<Broadcast>()->value;
+    const Broadcast *base_bc = base_stride_ramp->base.as<Broadcast>();
+
+    if (!base_bc) {
+        return {};
+    }
+
+    Expr base = base_bc->value;
     Expr stride;
 
     bool found_stride = false;

From 92426c8eb4f484112e86322630eef3ab507bb65c Mon Sep 17 00:00:00 2001
From: Frederik Engels <frederik.engels@codeplay.com>
Date: Tue, 14 Jun 2022 16:09:29 +0100
Subject: [PATCH 10/21] llvm below 13 is no longer supported

---
 test/correctness/tiled_matmul.cpp | 38 +++++++++++++++----------------
 1 file changed, 18 insertions(+), 20 deletions(-)

diff --git a/test/correctness/tiled_matmul.cpp b/test/correctness/tiled_matmul.cpp
index 793c816b778e..2a54910fafdd 100644
--- a/test/correctness/tiled_matmul.cpp
+++ b/test/correctness/tiled_matmul.cpp
@@ -50,14 +50,14 @@ bool equal_eps(float lhs, float rhs, float eps) {
 
 struct make_uint_t {
     template<typename... Args>
-    Type operator()(Args &&...args) const {
+    Type operator()(Args &&... args) const {
         return UInt(static_cast<Args &&>(args)...);
     }
 };
 
 struct make_int_t {
     template<typename... Args>
-    Type operator()(Args &&...args) const {
+    Type operator()(Args &&... args) const {
         return Int(static_cast<Args &&>(args)...);
     }
 };
@@ -263,27 +263,25 @@ int main(int argc, char **argv) {
         return -1;
     }
 
-    // llvm >= 13.0 is required for unsigned and float AMX instructions
-    if (Halide::Internal::get_llvm_version() >= 130) {
-        printf("Running AMX matmul (signed/unsigned)\n");
-        if (!run_tests(matmul_su, 1)) {
-            return -1;
-        }
+    printf("Running AMX matmul (signed/unsigned)\n");
+    if (!run_tests(matmul_su, 1)) {
+        return -1;
+    }
 
-        printf("Running AMX matmul (unsigned/signed)\n");
-        if (!run_tests(matmul_us, 1)) {
-            return -1;
-        }
+    printf("Running AMX matmul (unsigned/signed)\n");
+    if (!run_tests(matmul_us, 1)) {
+        return -1;
+    }
 
-        printf("Running AMX matmul (unsigned/unsigned)\n");
-        if (!run_tests(matmul_uu, 1)) {
-            return -1;
-        }
+    printf("Running AMX matmul (unsigned/unsigned)\n");
+    if (!run_tests(matmul_uu, 1)) {
+        return -1;
+    }
 
-        printf("Running AMX matmul (bf16)\n");
-        if (!run_tests(matmul_bf16, 2)) {
-            return -1;
-        }
+    printf("Running AMX matmul (bf16)\n");
+    if (!run_tests(matmul_bf16, 2)) {
+        return -1;
     }
+
     return 0;
 }
\ No newline at end of file

From fba5a2e89fb4cbf094f71dec05966a01ec7bf8bb Mon Sep 17 00:00:00 2001
From: Frederik Engels <frederik.engels@codeplay.com>
Date: Wed, 22 Jun 2022 20:27:38 +0100
Subject: [PATCH 11/21] replace single pattern with commutative permutations

---
 src/ExtractTileOperations.cpp | 26 ++++++++++++++++++--------
 1 file changed, 18 insertions(+), 8 deletions(-)

diff --git a/src/ExtractTileOperations.cpp b/src/ExtractTileOperations.cpp
index 9ee2dab18531..3858658469f0 100644
--- a/src/ExtractTileOperations.cpp
+++ b/src/ExtractTileOperations.cpp
@@ -56,16 +56,26 @@ const auto wild_i32x = Variable::make(Int(32, 0), "*");
 
 Tile<1> get_1d_tile_index(const Expr &e) {
     if (const auto *r1 = e.as<Ramp>()) {
-        Expr pattern = ((wild_i32 * wild_i32) + wild_i32) * wild_i32;
 
-        std::vector<Expr> results;
-        if (!expr_match(pattern, r1->base, results)) {
-            return {};
+        const auto stride_var = Variable::make(Int(32), "stride");
+        const auto v1 = Variable::make(Int(32), "v1");
+        const auto v2 = Variable::make(Int(32), "v2");
+        const auto v3 = Variable::make(Int(32), "v3");
+
+        Expr patterns[] = {
+            ((v1 * stride_var) + v2) * v3,
+            v3 * ((v1 * stride_var) + wild_i32),
+            (v2 + (v1 * stride_var)) * v3,
+            v3 * (v2 + (v1 * stride_var)),
+        };
+
+        std::map<std::string, Expr> matches;
+        for (const auto &pattern : patterns) {
+            if (expr_match(pattern, r1->base, matches)) {
+                auto stride = std::move(matches["stride"]);
+                return {true, r1->base, {std::move(stride)}, {r1->lanes}};
+            }
         }
-
-        auto stride = std::move(results[1]);
-
-        return {true, r1->base, {std::move(stride)}, {r1->lanes}};
     }
 
     return {};

From 3764d49506a8c7b2d140fac6992e7b0f53f40afe Mon Sep 17 00:00:00 2001
From: Frederik Engels <frederik.engels@codeplay.com>
Date: Thu, 23 Jun 2022 15:59:14 +0100
Subject: [PATCH 12/21] check if the stride is an `IntImm`, otherwise reject
 pattern

---
 src/ExtractTileOperations.cpp | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/src/ExtractTileOperations.cpp b/src/ExtractTileOperations.cpp
index 3858658469f0..8eb7860001f7 100644
--- a/src/ExtractTileOperations.cpp
+++ b/src/ExtractTileOperations.cpp
@@ -73,6 +73,9 @@ Tile<1> get_1d_tile_index(const Expr &e) {
         for (const auto &pattern : patterns) {
             if (expr_match(pattern, r1->base, matches)) {
                 auto stride = std::move(matches["stride"]);
+                if (!stride.as<IntImm>()) {
+                    return {};
+                }
                 return {true, r1->base, {std::move(stride)}, {r1->lanes}};
             }
         }

From bd5cd6b471a93eb94aca669d626d0c54ab6dc445 Mon Sep 17 00:00:00 2001
From: Frederik Engels <frederik.engels@codeplay.com>
Date: Mon, 27 Jun 2022 16:26:03 +0100
Subject: [PATCH 13/21] apply clang-format-13

---
 test/correctness/tiled_matmul.cpp | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/test/correctness/tiled_matmul.cpp b/test/correctness/tiled_matmul.cpp
index 2a54910fafdd..6b5e949dd518 100644
--- a/test/correctness/tiled_matmul.cpp
+++ b/test/correctness/tiled_matmul.cpp
@@ -50,14 +50,14 @@ bool equal_eps(float lhs, float rhs, float eps) {
 
 struct make_uint_t {
     template<typename... Args>
-    Type operator()(Args &&... args) const {
+    Type operator()(Args &&...args) const {
         return UInt(static_cast<Args &&>(args)...);
     }
 };
 
 struct make_int_t {
     template<typename... Args>
-    Type operator()(Args &&... args) const {
+    Type operator()(Args &&...args) const {
         return Int(static_cast<Args &&>(args)...);
     }
 };

From ddd3a0e8a189b975829fde4374e17a0a0cf5bf2c Mon Sep 17 00:00:00 2001
From: Frederik Engels <frederik.engels@codeplay.com>
Date: Fri, 15 Jul 2022 21:58:29 +0100
Subject: [PATCH 14/21] rename wild_i32 -> v2

---
 src/ExtractTileOperations.cpp | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/ExtractTileOperations.cpp b/src/ExtractTileOperations.cpp
index 8eb7860001f7..7bf2991bff4e 100644
--- a/src/ExtractTileOperations.cpp
+++ b/src/ExtractTileOperations.cpp
@@ -64,7 +64,7 @@ Tile<1> get_1d_tile_index(const Expr &e) {
 
         Expr patterns[] = {
             ((v1 * stride_var) + v2) * v3,
-            v3 * ((v1 * stride_var) + wild_i32),
+            v3 * ((v1 * stride_var) + v2),
             (v2 + (v1 * stride_var)) * v3,
             v3 * (v2 + (v1 * stride_var)),
         };

From db12ffc23c6208a5da02278305f595e588d9386d Mon Sep 17 00:00:00 2001
From: Frederik Engels <frederik.engels@codeplay.com>
Date: Fri, 15 Jul 2022 21:59:02 +0100
Subject: [PATCH 15/21] check if v1 could be the stride value

---
 src/ExtractTileOperations.cpp | 12 +++++++++---
 1 file changed, 9 insertions(+), 3 deletions(-)

diff --git a/src/ExtractTileOperations.cpp b/src/ExtractTileOperations.cpp
index 7bf2991bff4e..3dc993942ca8 100644
--- a/src/ExtractTileOperations.cpp
+++ b/src/ExtractTileOperations.cpp
@@ -73,10 +73,16 @@ Tile<1> get_1d_tile_index(const Expr &e) {
         for (const auto &pattern : patterns) {
             if (expr_match(pattern, r1->base, matches)) {
                 auto stride = std::move(matches["stride"]);
-                if (!stride.as<IntImm>()) {
-                    return {};
+                // stride must be a constant in order to not be confused with v1
+                if (stride.as<IntImm>()) {
+                    return {true, r1->base, {std::move(stride)}, {r1->lanes}};
+                }
+
+                // if stride wasn't a constant then v1 could possibly be the stride if constant
+                auto v1_expr = std::move(matches["v1"]);
+                if (v1_expr.as<IntImm>()) {
+                    return {true, r1->base, {std::move(v1_expr)}, {r1->lanes}};
                 }
-                return {true, r1->base, {std::move(stride)}, {r1->lanes}};
             }
         }
     }

From b56aad1eff11e7ec75ba58280e406f568b136949 Mon Sep 17 00:00:00 2001
From: Frederik Engels <frederik.engels@codeplay.com>
Date: Wed, 27 Jul 2022 16:25:04 +0100
Subject: [PATCH 16/21] add more detail to a receiving a bad type

---
 src/ExtractTileOperations.cpp | 12 ++++++------
 1 file changed, 6 insertions(+), 6 deletions(-)

diff --git a/src/ExtractTileOperations.cpp b/src/ExtractTileOperations.cpp
index 3dc993942ca8..b6e3be67b42e 100644
--- a/src/ExtractTileOperations.cpp
+++ b/src/ExtractTileOperations.cpp
@@ -384,12 +384,12 @@ Matmul convert_to_matmul(const Store *op, const string &new_name, AMXOpType op_t
     }
 
     if (rhs_cast) {
-        if (op_type == AMXOpType::Int8) {
-            if (!(rhs_cast->value.type().element_of() == Int(8) || rhs_cast->value.type().element_of() == UInt(8))) {
-                user_error << "Expected rhs cast of i8/u8, got " << rhs_cast->value.type();
-            }
-        } else {  // AMXOpType::Bfloat16
-            user_assert(rhs_cast->value.type().element_of() == BFloat(16)) << "Expected rhs cast of bf16";
+        bool is_i8_u8 = rhs_cast->value.type().element_of() == Int(8) || rhs_cast->value.type().element_of() == UInt(8);
+        bool is_bf16 = rhs_cast->value.type().element_of() == BFloat(16);
+
+        if ((op_type == AMXOpType::Int8 && !is_i8_u8) || (op_type == AMXOpType::Bfloat16 && !is_bf16)) {
+            user_error << "Expected rhs type of " << (op_type == AMXOpType::Int8 ? "i8/u8" : "bf16")
+                       << ", got " << rhs_cast->value.type() << " instead.\nIn Expression: " << Expr(rhs_cast);
         }
     } else {
         return {};

From 7be1b5b6867ce22ae36bd4f15e0f07545f3c5c81 Mon Sep 17 00:00:00 2001
From: Frederik Engels <frederik.engels@codeplay.com>
Date: Wed, 27 Jul 2022 16:57:51 +0100
Subject: [PATCH 17/21] added short explanation of the right-hand matrix layout

---
 src/ExtractTileOperations.cpp | 14 ++++++++++++++
 1 file changed, 14 insertions(+)

diff --git a/src/ExtractTileOperations.cpp b/src/ExtractTileOperations.cpp
index b6e3be67b42e..c6ae0d959e34 100644
--- a/src/ExtractTileOperations.cpp
+++ b/src/ExtractTileOperations.cpp
@@ -374,6 +374,20 @@ Matmul convert_to_matmul(const Store *op, const string &new_name, AMXOpType op_t
 
     if (lhs_load && !rhs_broadcast) {
         // now working on a larger k dimension
+        // with a K dimension of 4 (or 2) with bf16 all the elements in the right-hand matrix are 
+        // layed out in a way that multiplying with a column can be done in a single dot product.
+        // Therefore the indexing can be reused with a broadcast,
+        // with higher K dimensions this can no longer be done and the broadcast won't exist.
+        // ┌──┐
+        // │1 │
+        // │2 │
+        // │3 │   ┌────────┐
+        // │4 │   │1234    │
+        // │5 │   │5678    │
+        // │6 │   └────────┘
+        // │7 │
+        // │8 │
+        // └──┘
         rhs_cast = matches[1].as<Cast>();
     } else {
         rhs_cast = rhs_broadcast->value.as<Cast>();

From 81bacbb680ef74ab47ec30700a4868b272e8f281 Mon Sep 17 00:00:00 2001
From: Frederik Engels <frederik.engels@codeplay.com>
Date: Wed, 27 Jul 2022 16:59:55 +0100
Subject: [PATCH 18/21] added explanation for where the 4 comes from

---
 src/ExtractTileOperations.cpp | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/src/ExtractTileOperations.cpp b/src/ExtractTileOperations.cpp
index c6ae0d959e34..81d8a3da863b 100644
--- a/src/ExtractTileOperations.cpp
+++ b/src/ExtractTileOperations.cpp
@@ -303,7 +303,9 @@ BaseStride get_rhs_tile_index(const Expr &index, int element_width, int tile_x,
                 return {};
             }
 
-            // times 4 because of the rhs layout
+            // times 4 because of the rhs layout, each vector used by AMX is 4 bytes in size.
+            // For the 4 gets divided by the element width which means each vector has 4 elements in u8/i8 and
+            // 2 elements for bf16.
             return {true, rhs_tile1.base, rhs_tile1.stride[0] * (4 / element_width)};
         }
     } else {

From a80fc766ff63a2002a4fc459c04b17ef0127d784 Mon Sep 17 00:00:00 2001
From: Frederik Engels <frederik.engels@codeplay.com>
Date: Wed, 27 Jul 2022 17:15:55 +0100
Subject: [PATCH 19/21] provide further documentation as to the layout of AMX

---
 src/ExtractTileOperations.cpp | 27 ++++++++++++++++++++++++++-
 1 file changed, 26 insertions(+), 1 deletion(-)

diff --git a/src/ExtractTileOperations.cpp b/src/ExtractTileOperations.cpp
index 81d8a3da863b..7f07c3b52392 100644
--- a/src/ExtractTileOperations.cpp
+++ b/src/ExtractTileOperations.cpp
@@ -5,6 +5,31 @@
 #include "IROperator.h"
 #include "Util.h"
 
+/** \file Support extraction of AMX instructions. */
+
+/**
+ * https://asciiflow.com/#/share/eJyVUkFugzAQ%2FMrKxwoRhdAkza23SmlySHvogQsBp7FkbGSbAoryiz6nr%2BlLugZDk6ghKvJhbXZmd2b3QEScUbIQBece4XFNFVmQQ0SqiCwegtCLSI1RMBtjZGhl8BIRAHh%2BeoFVbBSr4Pq36ZOiSOBpX5cDCEikSGhuipjzun0pmdnD4%2BqtwX9%2Ffg2cLmUcTML76WyO4VAtWJ%2Ff7kIkWMEJ6gbBae2%2F3q53OHBuFBz3TS1HodPqfvUO3%2F4wO7gQag07IXqVkCuZU4VzyApuWI5BAJkdZ0K1B2ZP2%2BwJ%2FEs%2BjhKY0EYViWFSaMAaO6kypBY1hLCtDRIvMTvsekmlsc2kiGgKMw2cxqkGIyEGjn%2FlzonoIMjPUibeQX5Q1bHGisbav%2FBh2kHW2ESzdlaZkqUltaFd9UZ25TnIrIOg%2Bb7vQykLnv661GysRSaSF1k78HkHcaSbntSReLAtTL%2FscOlaI9rxYaRzzgwUOTrZeOCokLzN0TDqRYvUqtFwB6Fvqco9S5r%2BBCiqsWmNLHabzny2Y7E4PyJHcvwBx0t%2BJw%3D%3D)
+ *
+ *   LHS Matrix                           RHS Matrix
+ *
+ *      K                            conceptually      with AMX
+ *  ┌────────┐
+ *  │12345678│                             N             N*4
+ *M │        │                            ┌──┐        ┌────────┐
+ *  └────────┘                            │1 │     K/4│1234    │
+ *                                        │2 │        │5678    │
+ * To properly multiply 2 matrices, the   │3 │        └────────┘
+ * AMX instructions perform many 4 byte  K│4 │
+ * dot products, this leads to a lot of   │5 │
+ * striding over 4 byte areas.            │6 │
+ * Normally the row of the LHS matrix,    │7 │
+ * 123... would multiply with the column  │8 │
+ * of the RHS matrix 123..., but with AMX └──┘
+ * this column is split up into a matrix of columns / 4 byte and rows * 4.
+ * which then results in K/4 dot products per row.
+ *
+ */
+
 namespace Halide {
 namespace Internal {
 
@@ -376,7 +401,7 @@ Matmul convert_to_matmul(const Store *op, const string &new_name, AMXOpType op_t
 
     if (lhs_load && !rhs_broadcast) {
         // now working on a larger k dimension
-        // with a K dimension of 4 (or 2) with bf16 all the elements in the right-hand matrix are 
+        // with a K dimension of 4 (or 2) with bf16 all the elements in the right-hand matrix are
         // layed out in a way that multiplying with a column can be done in a single dot product.
         // Therefore the indexing can be reused with a broadcast,
         // with higher K dimensions this can no longer be done and the broadcast won't exist.

From 314656cbb1084923ace01b727623124f2dc5ee3b Mon Sep 17 00:00:00 2001
From: Frederik Engels <frederik.engels@codeplay.com>
Date: Fri, 29 Jul 2022 01:32:10 +0100
Subject: [PATCH 20/21] add comments for expected patterns to
 get_3d_rhs_tile_index

---
 src/ExtractTileOperations.cpp | 17 +++++++++++++++++
 1 file changed, 17 insertions(+)

diff --git a/src/ExtractTileOperations.cpp b/src/ExtractTileOperations.cpp
index 7f07c3b52392..7194d20361b1 100644
--- a/src/ExtractTileOperations.cpp
+++ b/src/ExtractTileOperations.cpp
@@ -213,6 +213,7 @@ Tile<3> get_3d_rhs_tile_index(const Expr &e, int element_width) {
     const Add *add_lhs = nullptr;
 
     // there's not always a sub pattern
+    // This depends on whether we have an ImageParam or a Buffer
     if (!sub) {
         add_lhs = e.as<Add>();
     } else {
@@ -223,13 +224,17 @@ Tile<3> get_3d_rhs_tile_index(const Expr &e, int element_width) {
         return {};
     }
 
+    // The right hand side of the add expression is used for retrieving the dimensions of the matrix.
     // obtain the x, y, r dimensions
+    // this expr looks like below, the shape of `add_lhs->a` can be seen further down below
+    // broadcast(ramp(0, 1, r), x*y) % broadcast(4, x*y*r) + broadcast(ramp(broadcast(_, r), broadcast(4, r), x) , y)
     const Add *dim_expr = add_lhs->b.as<Add>();
 
     if (!dim_expr) {
         return {};
     }
 
+    // broadcast(ramp(broadcast(_, r), broadcast(4, r), x), y)
     const Broadcast *base_stride_bc = dim_expr->b.as<Broadcast>();
 
     if (!base_stride_bc) {
@@ -238,12 +243,14 @@ Tile<3> get_3d_rhs_tile_index(const Expr &e, int element_width) {
 
     int tile_y = base_stride_bc->lanes;
 
+    // broadcast(ramp(0, 1, r), x*y) % broadcast(4, x*y*r)
     const Mod *mod = dim_expr->a.as<Mod>();
 
     if (!mod) {
         return {};
     }
 
+    // broadcast(ramp(0, 1, r), x*y)
     const Broadcast *bc_ramp = mod->a.as<Broadcast>();
 
     if (!bc_ramp) {
@@ -253,6 +260,7 @@ Tile<3> get_3d_rhs_tile_index(const Expr &e, int element_width) {
     int tile_xy = bc_ramp->lanes;
     int tile_x = tile_xy / tile_y;
 
+    // ramp(0, 1, r)
     const Ramp *r_ramp = bc_ramp->value.as<Ramp>();
 
     if (!r_ramp) {
@@ -262,12 +270,14 @@ Tile<3> get_3d_rhs_tile_index(const Expr &e, int element_width) {
     int tile_r = r_ramp->lanes;
 
     // get the base and stride
+    // ramp(broadcast(_, r), broadcast(4, r), x)
     const Ramp *base_stride_ramp = base_stride_bc->value.as<Ramp>();
 
     if (!base_stride_ramp) {
         return {};
     }
 
+    // broadcast(_, r)
     const Broadcast *base_bc = base_stride_ramp->base.as<Broadcast>();
 
     if (!base_bc) {
@@ -279,6 +289,10 @@ Tile<3> get_3d_rhs_tile_index(const Expr &e, int element_width) {
 
     bool found_stride = false;
 
+    // the following pattern will match the following shape
+    // broadcast(ramp(0, 1, k), x*y) / broadcast(4, x*y*k) * broadcast(_, x*y*k)
+    // where the stride is marked by _.
+
     // this stride pattern can occur if `tile_r` is the same size as `acc`
     auto stride_pattern = Broadcast::make(Ramp::make(0, 1, tile_r), tile_x * tile_y) / Broadcast::make((4 / element_width), tile_x * tile_y * tile_r) * Broadcast::make(wild_i32, tile_x * tile_y * tile_r);
 
@@ -288,6 +302,9 @@ Tile<3> get_3d_rhs_tile_index(const Expr &e, int element_width) {
         stride = std::move(results[0]);
     }
 
+    // This pattern is similar to the above except with an additional offset to iterate over the tiles in the k dimension
+    // (broadcast(ramp(0, 1, k), m * n) / broadcast(4, m*n*k) + _) * broadcast(_, m*n*k)
+    // here the first _ marks the base and the second _ the stride.
     if (!found_stride) {
         stride_pattern = (Broadcast::make(Ramp::make(0, 1, tile_r), tile_x * tile_y) / Broadcast::make((4 / element_width), tile_x * tile_y * tile_r) + wild_i32) * Broadcast::make(wild_i32, tile_x * tile_y * tile_r);
         if (expr_match(stride_pattern, add_lhs->a, results)) {

From 9b3067996efa7460978a0fa15a94798086dda5af Mon Sep 17 00:00:00 2001
From: Frederik Engels <frederik.engels@codeplay.com>
Date: Mon, 1 Aug 2022 16:24:51 +0100
Subject: [PATCH 21/21] Document the matched pattern

---
 src/ExtractTileOperations.cpp | 14 +++++++++++++-
 1 file changed, 13 insertions(+), 1 deletion(-)

diff --git a/src/ExtractTileOperations.cpp b/src/ExtractTileOperations.cpp
index 7194d20361b1..8fdcea73f34b 100644
--- a/src/ExtractTileOperations.cpp
+++ b/src/ExtractTileOperations.cpp
@@ -208,6 +208,18 @@ Tile<3> get_3d_tile_index(const Expr &e) {
     return {true, base, {x_stride, 0, r_stride}, {x_tile, y_tile, r_tile}};
 }
 
+/**
+ * \brief Get the 3d rhs tile index configuration
+ *
+ * \param e index expression
+ * \param element_width the width of the elements, 1 for u8/i8, 2 for bf16
+ * \return Tile<3> the tile configuration found
+ *
+ * The pattern which is getting matched looks roughly like
+ * `broadcast(ramp(0, 1, r), x*y) / broadcast(4, x*y*r) + optional(broadcast(base, x*y*r)) * broadcast(8, x*y*r) +
+ *  broadcast(ramp(0, 1, r), x*y) % broadcast(4, x*y*r) +
+ *  broadcast(ramp(broadcast(_, r), broadcast(4, r), x) , y)`
+ */
 Tile<3> get_3d_rhs_tile_index(const Expr &e, int element_width) {
     const auto *sub = e.as<Sub>();
     const Add *add_lhs = nullptr;
@@ -227,7 +239,7 @@ Tile<3> get_3d_rhs_tile_index(const Expr &e, int element_width) {
     // The right hand side of the add expression is used for retrieving the dimensions of the matrix.
     // obtain the x, y, r dimensions
     // this expr looks like below, the shape of `add_lhs->a` can be seen further down below
-    // broadcast(ramp(0, 1, r), x*y) % broadcast(4, x*y*r) + broadcast(ramp(broadcast(_, r), broadcast(4, r), x) , y)
+    // broadcast(ramp(0, 1, r), x*y) % broadcast(4, x*y*r) + broadcast(ramp(broadcast(base, r), broadcast(4, r), x) , y)
     const Add *dim_expr = add_lhs->b.as<Add>();
 
     if (!dim_expr) {