Add support for AMX instructions #5818

mcleary · 2021-03-17T09:14:48Z

This pull request continues the work started by @jwlawson in #5780 with the objective of adding initial support for AMX instructions in Halide.

The main addition here is the fix for building Halide using LLVM 11. Support for AMX instructions requires LLVM 12 or newer so when building with LLVM 11 the unsupported instructions are not included.

A new LLVM module was created (x86_amx.ll) to contain all the required intrinsics to enable support for tile operations, this module is only included when LLVM >= 12 is present.

This will only be included if LLVM >= 12 is used to build Halide

steven-johnson · 2021-03-17T17:05:52Z

(Synced to head to fix some irrelevant LLVM build issues)

mcleary · 2021-03-17T17:26:35Z

(Synced to head to fix some irrelevant LLVM build issues)

That is a good excuse for me to setup LLVM 13 locally. I will try to see why it is not building with it

Recent changes in LLVM trunk made the previous calling convention deprecated (and thus compiling with warning/error)

steven-johnson · 2021-03-24T16:39:42Z

The OSX failure is unrelated (will be fixed by #5841), should be good to land

steven-johnson · 2021-03-31T16:35:27Z

You should sync this to master to force the bots to retry.

mcleary · 2021-03-31T21:33:18Z

You should sync this to master to force the bots to retry.

I'm not sure if the buildbot is still running since there is a "cancelled" message there.

steven-johnson · 2021-04-06T22:56:02Z

Please try syncing to master once again; hopefully the buildbots will finally be clean.

mcleary · 2021-04-07T09:04:28Z

Please try syncing to master once again; hopefully the buildbots will finally be clean.

Thanks, I will do that, hopefully it will all be green now.

steven-johnson · 2021-04-12T16:47:49Z

Failures are the unrelated cuda-hang failure that we still haven't diagnosed; ok to land

dependencies/llvm/CMakeLists.txt

alexreinking · 2021-07-09T20:12:21Z

src/CodeGen_X86.cpp

@@ -190,7 +200,7 @@ const x86Intrinsic intrinsic_defs[] = {

    {"dpwssdx16", Int(32, 16), "dot_product", {Int(32, 16), Int(16, 32), Int(16, 32)}, Target::AVX512_SapphireRapids},
    {"dpwssdx8", Int(32, 8), "dot_product", {Int(32, 8), Int(16, 16), Int(16, 16)}, Target::AVX512_SapphireRapids},
-    {"dpwssdx4", Int(32, 4), "dot_product", {Int(32, 4), Int(16, 8), Int(16, 8)}, Target::AVX512_SapphireRapids},
+    {"dpwssdx4", Int(32, 4), "dot_product", {Int(32, 4), Int(16, 8), Int(16, 8)}, Target::AVX512_SapphireRapids},    


nit: irrelevant whitespace?

Addressed in b1e1452

alexreinking · 2021-07-09T20:16:08Z

src/runtime/CMakeLists.txt

+if (Halide_LLVM_VERSION VERSION_GREATER_EQUAL 12.0)
+  # AMX instructions require LLVM 12 or newer
+  list(APPEND RUNTIME_LL x86_amx)
+endif ()
+


Does including this fail at build time or at runtime only?

this fails at build time with the following message

[1/7] Generating initmod.x86_amx.bc FAILED: src/runtime/initmod.x86_amx.bc /home/frederik/projects/halide/build-11/src/runtime/initmod.x86_amx.bc cd /home/frederik/projects/halide/build-11/src/runtime && /usr/lib/llvm-11/bin/llvm-as /home/frederik/projects/halide/src/runtime/x86_amx.ll -o initmod.x86_amx.bc /usr/lib/llvm-11/bin/llvm-as: /home/frederik/projects/halide/src/runtime/x86_amx.ll:3:18: error: expected type %2 = tail call x86_amx @llvm.x86.tileloadd64.internal(i16 %rows, i16 %colbytes, i8* %1, i64 %stride) nounwind readonly ^

frengels · 2021-09-01T15:43:50Z

Just letting you know we haven't lost track of this and the TensorCore PRs. We had some different priorities and annual leave. I look forward to getting this merged soon.

frengels · 2021-09-14T09:24:45Z

I don't think the test failures are related to anything in this PR

steven-johnson

LGTM so far -- aside from style nits, I think it would be good to split the new test into correctness and performance tests, as Halide does for virtually all other features.

steven-johnson · 2021-09-17T21:56:52Z

src/ExtractTileOperations.cpp

@@ -0,0 +1,414 @@
+#include "ExtractTileOperations.h"
+
+#include "IRMatch.h"  // expr_match


Nit: We don't usually add comments explaining why each header is included.

Speaking of which, it might be a good idea to run IWYU on our codebase...

steven-johnson · 2021-09-17T21:58:06Z

src/ExtractTileOperations.cpp

+
+enum class AMXOpType {
+    Int8,
+    Bf16,


Nit: I assume this is bfloat16? If so, spelling it out (eg Bfloat16) would be preferable.

steven-johnson · 2021-09-17T21:59:15Z

src/ExtractTileOperations.cpp

+    case AMXOpType::Bf16:
+        return Float(32, 256);
+    default:
+        return Type();


I assume this is a should-never-happen case, so doing something like internal_error << "Unexpected"; would be appropriate.

steven-johnson · 2021-09-17T22:00:18Z

src/ExtractTileOperations.cpp

+const auto wild_i32 = Variable::make(Int(32), "*");
+const auto wild_i32x = Variable::make(Int(32, 0), "*");
+
+Tile<2> is_2d_tile_index(const Expr &e) {


Nit: I'd expect a function named "is_whatever" to return bool, but this returns a struct. Something like get_2d_tile_index would be better.

steven-johnson · 2021-09-17T22:00:53Z

src/ExtractTileOperations.cpp

+    return {};
+}
+
+Tile<3> is_3d_tile_index(const Expr &e) {


Same nit here.

steven-johnson · 2021-09-17T22:04:57Z

src/ExtractTileOperations.cpp

+    // 4 bytes for i32, f32
+    auto colbytes = tile_y * 4;
+    auto matmul =
+        Call::make(res_type, "tile_matmul", {tile_x, colbytes, tile_r, out, lhs, rhs}, Call::Intrinsic);


No need to split this into two lines

steven-johnson · 2021-09-17T22:06:32Z

src/ExtractTileOperations.cpp

+                op_type = AMXOpType::Bf16;
+            }
+
+            user_assert(!in_allocate) << "Found two possible tile allocations for AMX allocation";


Would it be helpful to append amx_name or tile_name to the error message, for debugging purposes?

steven-johnson · 2021-09-17T22:06:56Z

src/ExtractTileOperations.cpp

+            }
+
+            auto alloc_type = amx_op_type_result_type(op_type);
+


No need for this blank line

steven-johnson · 2021-09-17T22:07:40Z

src/ExtractTileOperations.cpp

+        }
+
+        auto body = mutate(op->body);
+        return ProducerConsumer::make(amx_name, op->is_producer, body);


Nit: std::move(body) ?

steven-johnson · 2021-09-17T22:10:56Z

test/performance/tiled_matmul.cpp

+        .vectorize(mmyi);
+
+    Func result = mm.in();
+    //result.print_loop_nest();


Don't check in commented-out code (unless there is a comment explaining why, as is done elsewhere in this file)

frengels · 2021-09-29T16:15:21Z

When converting to correctness tests there's a bit of a change in the pattern for the rhs load when using Buffer instead of ImageParam, so I'm having to take a bit extra time to make sure it's accepting all possible scenarios and generating correct code.

When using `Buffer` instead of `ImageParam` the `Ramp` expression generated is 1D instead of 2D, therefore we recognize this with a special case. The lanes are still matched against the dimensions of the LHS 3d tile lanes.

frengels · 2021-10-20T15:36:00Z

I think the recent commits addressed all comments, is there anything else that needs to be addressed?

abadams · 2021-10-21T17:42:14Z

lgtm. The pattern matching seems to be pretty ad-hoc and possibly brittle, but that can always be improved later. The checks for LLVM 12 will be removed pretty soon too.

steven-johnson · 2021-10-22T00:33:25Z

This is failing for LLVM11 for Makefile-based builds. I'll see if I can prep a patch.

jwlawson and others added 12 commits March 16, 2021 18:25

Add support for AMX tile instructions

036b037

Make AMX transform opt-in with memory type

2bd3453

Clean up tiled_matmul test

34e1a4c

Handle AMX intrinsic attributes better

dfeac55

Format

a9f84de

Fix test to behave like other tests

da04b0a

Add doc and missing load check

9923040

Format

1a1d10a

Throw error if user requests AMX for invalid operation

d54fe24

Add Tile lowering pass to makefile

228beda

Use spaces in Makefile

673480a

Place AMX instrinsics into a separate module (x86_amx.ll)

f91d79f

This will only be included if LLVM >= 12 is used to build Halide

steven-johnson requested a review from halidebuildbots March 17, 2021 16:26

Merge branch 'master' into pr/5818

9fc4366

steven-johnson and others added 3 commits March 17, 2021 12:56

Fix CreateAlignedLoad() call in CodeGen_X86

16a1c7b

Recent changes in LLVM trunk made the previous calling convention deprecated (and thus compiling with warning/error)

Merge branch 'master' into tile_matmul

b3b1dfa

Merge branch 'master' into tile_matmul

03f894a

Merge branch 'master' into tile_matmul

97890dd

Merge branch 'master' into tile_matmul

293ca90

Merge branch 'master' into tile_matmul

4b950e5

frengels added 2 commits April 20, 2021 11:56

fix exporting to module

d0e5123

add llvm funcs for su, us, uu amx variants

9189120

alexreinking reviewed Jul 9, 2021

View reviewed changes

dependencies/llvm/CMakeLists.txt Show resolved Hide resolved

alexreinking reviewed Jul 9, 2021

View reviewed changes

frengels added 3 commits September 2, 2021 12:09

have tile_store return i32

34557cb

make is_3d_tile_index robust to indexing changes

5ad06e0

Merge remote-tracking branch 'upstream/master' into tile_matmul_bf16

f0f9f3e

frengels force-pushed the tile_matmul branch from aac8f78 to f0f9f3e Compare September 8, 2021 15:48

frengels added 3 commits September 8, 2021 17:11

apply formatting suggestions

7cab155

both first and second can be const qualified

8f83544

remove trailing whitespace in unformatted section

b1e1452

Merge branch 'master' into pr/5818

95d38f0

steven-johnson requested changes Sep 17, 2021

View reviewed changes

make requested style changes

14df0bc

frengels added 5 commits October 1, 2021 14:31

rename NewMatmul -> Matmul

8b63d77

fix warning about missing return value

6a5eeaa

use get_1d_tile_index to handle special case

cc7c97d

When using `Buffer` instead of `ImageParam` the `Ramp` expression generated is 1D instead of 2D, therefore we recognize this with a special case. The lanes are still matched against the dimensions of the LHS 3d tile lanes.

add correctness test for AMX instructions

014f0c6

correctness part has been separated out

655dbdf

frengels force-pushed the tile_matmul branch from 5280253 to 655dbdf Compare October 6, 2021 09:00

remove unused variables

abc660b

abadams approved these changes Oct 21, 2021

View reviewed changes

steven-johnson self-requested a review October 21, 2021 17:51

steven-johnson approved these changes Oct 21, 2021

View reviewed changes

steven-johnson merged commit 0078880 into halide:master Oct 21, 2021

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Add support for AMX instructions #5818

Add support for AMX instructions #5818

mcleary commented Mar 17, 2021

steven-johnson commented Mar 17, 2021

mcleary commented Mar 17, 2021

steven-johnson commented Mar 24, 2021

steven-johnson commented Mar 31, 2021

mcleary commented Mar 31, 2021

steven-johnson commented Apr 6, 2021

mcleary commented Apr 7, 2021 •

edited

Loading

steven-johnson commented Apr 12, 2021

alexreinking Jul 9, 2021

frengels Sep 9, 2021

alexreinking Jul 9, 2021

frengels Sep 9, 2021

frengels commented Sep 1, 2021

frengels commented Sep 14, 2021

steven-johnson left a comment

steven-johnson Sep 17, 2021

alexreinking Sep 17, 2021

steven-johnson Sep 17, 2021

steven-johnson Sep 17, 2021

steven-johnson Sep 17, 2021

steven-johnson Sep 17, 2021

steven-johnson Sep 17, 2021

steven-johnson Sep 17, 2021

steven-johnson Sep 17, 2021

steven-johnson Sep 17, 2021

steven-johnson Sep 17, 2021

steven-johnson Sep 17, 2021

frengels commented Sep 29, 2021

frengels commented Oct 20, 2021

abadams commented Oct 21, 2021

steven-johnson commented Oct 22, 2021

		@@ -0,0 +1,414 @@
		#include "ExtractTileOperations.h"

		#include "IRMatch.h" // expr_match

Add support for AMX instructions #5818

Add support for AMX instructions #5818

Conversation

mcleary commented Mar 17, 2021

steven-johnson commented Mar 17, 2021

mcleary commented Mar 17, 2021

steven-johnson commented Mar 24, 2021

steven-johnson commented Mar 31, 2021

mcleary commented Mar 31, 2021

steven-johnson commented Apr 6, 2021

mcleary commented Apr 7, 2021 • edited Loading

steven-johnson commented Apr 12, 2021

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

frengels commented Sep 1, 2021

frengels commented Sep 14, 2021

steven-johnson 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

frengels commented Sep 29, 2021

frengels commented Oct 20, 2021

abadams commented Oct 21, 2021

steven-johnson commented Oct 22, 2021

mcleary commented Apr 7, 2021 •

edited

Loading