Repository: mmcloughlin/avo
Pull request #162 (2020-12-19 00:00:00)

Added support for AVX512 opmask registers (K0-K7)
Here's the result of a couple of hours of work learning how the register allocation code works and implementing the AVX512 opmask registers within the existing framework as I understand it. So this is a proposed stepping stone towards https://github.com/mmcloughlin/avo/issues/20

@lukechampine If merged, this code should also remove the need for you to define and explicitly manage opmask registers as in this code: 

https://github.com/lukechampine/blake3/blob/59a82aea36b67f99eb83ea2d12bb378392e234a1/avo/gen.go#L543-L545

The only real question I had with this code is how to properly account for the special treatment of the K0 register in AVX512.

My understanding from the Intel docs is that K0 is a real register just like K1-K7, and can be used for memory operations and to hold constants or temp intermediate values in calculations involving the other opmask registers, but the value stored in K0 can never be directly used as an actual opmask for general AVX512 SIMD instructions, because in that context specifying K0 means "all bits set"/"no mask" (which is also the default when no opmask register is specified at all).

So in this PR I treated it in much the same way as the SP register, that is, K0 is defined both as a pseudo-register (intended when used as an opmask), and a restricted physical register (when explicitly named and used for opmask ALU-type computations). However, because it is restricted, it will never be allocated as a virtual register.

Please review and let me know what you think.

mmcloughlin: Awesome! Thanks for the contribution, this should be good to land with some tweaks. I'll leave some comments on the code but a couple questions for now.

I'd prefer to land this on an `avx512` branch of the `avo` repo for now, since it's likely a large change overall and I can't anticipate all the consequences right now. Can you see how to do that in Github from your side? I guess I'd have to create the branch on my side and you'd edit the PR on your side.

Where are you at with this? I know I've said this before, but I really was planning on working on this, like literally now. I finished my grad school applications this week and I'm off work since yesterday so I should have the time, especially being locked in my apartment due to COVID. If you have made a lot of progress, maybe we could have a chat about it over zoom or something? If you have not made a lot of progress, are you okay with me picking up from here?

Thanks again!
mmcloughlin: We may also need an equivalent of PeachPy's `MaskedRegister` type.

https://github.com/Maratyszcza/PeachPy/blob/84e60a525e6f0788e88e090a11c5eb3b0e51a185/peachpy/x86_64/registers.py#L913

vsivsi: Answering the final Q above. Yes, I’m totally fine with @mmcloughlin  running off ahead with this. I’m technically on vacation through the end of the year, and while like (hopefully) most people I’m not actually going anywhere, I do hope to spend at least some of the break not staring at screens all day (presently failing at that!)

That said, I’m happy to pitch in with this in any way you find helpful. I may just not be as responsive as usual if I’m say, I don’t know, walking, reading an actual book or taking a glorious midday nap. 
mmcloughlin: Great, thanks @vsivsi! I really appreciate this contribution and your input on the design aspects.

I'll incorporate this into the `avx512` branch and when it's ready to land I'll use a `Co-authored-by` trailer to ensure you get credit. https://github.blog/2018-01-29-commit-together-with-co-authors/

Enjoy your break. I'm hoping to do some of that non-screen stuff, but we'll see 😄 

Patch:
diff --git a/reg/collection.go b/reg/collection.go
index d35c3a0..0930ea7 100644
--- a/reg/collection.go
+++ b/reg/collection.go
@@ -52,3 +52,6 @@ func (c *Collection) ZMM() VecVirtual { return c.Vec(S512) }
 
 // Vec allocates and returns a vector register of the given width.
 func (c *Collection) Vec(s Spec) VecVirtual { return newvecv(c.VirtualRegister(KindVector, s)) }
+
+// OM allocates and returns an opmask register.
+func (c *Collection) OM(s Spec) OMVirtual { return newomv(c.VirtualRegister(KindOpmask, s)) }
diff --git a/reg/x86.go b/reg/x86.go
index a1ec94c..2e1bfb8 100644
--- a/reg/x86.go
+++ b/reg/x86.go
@@ -5,6 +5,7 @@ const (
 	KindPseudo Kind = iota
 	KindGP
 	KindVector
+	KindOpmask
 )
 
 // Declare register families.
@@ -12,11 +13,13 @@ var (
 	Pseudo         = &Family{Kind: KindPseudo}
 	GeneralPurpose = &Family{Kind: KindGP}
 	Vector         = &Family{Kind: KindVector}
+	Opmask         = &Family{Kind: KindOpmask}
 
 	Families = []*Family{
 		Pseudo,
 		GeneralPurpose,
 		Vector,
+		Opmask,
 	}
 )
 
@@ -39,6 +42,7 @@ var (
 	ProgramCounter = Pseudo.define(S0, 0, "PC")
 	StaticBase     = Pseudo.define(S0, 0, "SB")
 	StackPointer   = Pseudo.define(S0, 0, "SP")
+	TrueOpmask     = Pseudo.define(S0, 0, "K0") // K0 can't be used as a predicate mask
 )
 
 // GP provides additional methods for general purpose registers.
@@ -329,3 +333,56 @@ var (
 	Z30 = vec(S512, 30, "Z30")
 	Z31 = vec(S512, 31, "Z31")
 )
+
+// OM provides additional methods for opmask registers.
+type OM interface {
+}
+
+// OMPhysical is a opmask physical register.
+type OMPhysical interface {
+	Physical
+	OM
+}
+
+type omp struct {
+	Physical
+}
+
+func newomp(r Physical) OMPhysical { return omp{Physical: r} }
+
+// OMVirtual is a virtual opmask register.
+type OMVirtual interface {
+	Virtual
+	OM
+}
+
+type omv struct {
+	Virtual
+}
+
+func newomv(v Virtual) OMVirtual { return omv{Virtual: v} }
+
+func om(s Spec, id Index, name string, flags ...Info) OMPhysical {
+	r := newomp(newregister(Opmask, s, id, name, flags...))
+	Opmask.add(r)
+	return r
+}
+
+// Opmask registers
+// Note! While K0 is a physical opmask register (it is a valid opmask src &
+// dest operand), it cannot be used as an opmask predicate value because in
+// that context K0 means "all true" or "no mask" regardless of the
+// actual contents of the physical register. For that reason, K0 should never
+// be assigned as a "general purpose" opmask register. However, it *can* be
+// explicitly operated upon by name as non-predicate operand (e.g. to hold a
+// constant or temporary value during calculations on other opmask registers).
+var (
+	K0 = om(S64, 0, "K0", Restricted) // Don't allocate K0
+	K1 = om(S64, 1, "K1")
+	K2 = om(S64, 2, "K2")
+	K3 = om(S64, 3, "K3")
+	K4 = om(S64, 4, "K4")
+	K5 = om(S64, 5, "K5")
+	K6 = om(S64, 6, "K6")
+	K7 = om(S64, 7, "K7")
+)


Test patch:
diff --git a/reg/reg_test.go b/reg/reg_test.go
index 07e0b8c..aec40d9 100644
--- a/reg/reg_test.go
+++ b/reg/reg_test.go
@@ -19,6 +19,7 @@ func TestIDIsVirtual(t *testing.T) {
 	cases := []Virtual{
 		GeneralPurpose.Virtual(42, S64),
 		Vector.Virtual(42, S128),
+		Opmask.Virtual(42, S64),
 	}
 	for _, r := range cases {
 		if !r.ID().IsVirtual() {
@@ -28,7 +29,7 @@ func TestIDIsVirtual(t *testing.T) {
 }
 
 func TestIDIsPhysical(t *testing.T) {
-	cases := []Physical{AL, AH, AX, EAX, RAX, X1, Y2, Z31}
+	cases := []Physical{AL, AH, AX, EAX, RAX, X1, Y2, Z31, K1}
 	for _, r := range cases {
 		if !r.ID().IsPhysical() {
 			t.FailNow()
@@ -98,6 +99,9 @@ func TestFamilyLookup(t *testing.T) {
 		{Vector, 27, S512, Z27},
 		{Vector, 1, S16, nil},
 		{Vector, 299, S256, nil},
+		{Opmask, 1, S64, K1},
+		{Opmask, 8, S64, nil},
+		{Opmask, 0, S64, K0},
 	}
 	for _, c := range cases {
 		got := c.Family.Lookup(c.ID, c.Spec)
@@ -156,6 +160,8 @@ func TestLookupPhysical(t *testing.T) {
 		{KindVector, 7, S128, X7},
 		{KindVector, 17, S256, Y17},
 		{KindVector, 27, S512, Z27},
+
+		{KindOpmask, 1, S64, K1},
 	}
 	for _, c := range cases {
 		if got := LookupPhysical(c.Kind, c.Index, c.Spec); !Equal(got, c.Expect) {
@@ -165,7 +171,7 @@ func TestLookupPhysical(t *testing.T) {
 }
 
 func TestLookupIDSelf(t *testing.T) {
-	cases := []Physical{AL, AH, AX, EAX, RAX, X1, Y2, Z31}
+	cases := []Physical{AL, AH, AX, EAX, RAX, X1, Y2, Z31, K1}
 	for _, r := range cases {
 		if got := LookupID(r.ID(), r.spec()); !Equal(got, r) {
 			t.FailNow()