cmd/internal/obj/riscv: add support for FENCE operands and FENCE.TSO

Add support for fine-grained memory ordering flags in the RISC-V FENCE
instruction to the assembler. This implements instruction validation and
encoding for predecessor and successor flags (I, O, R, W) rather than
always falling back to a full memory barrier. This allows more precise
memory barriers like FENCE R, RW or FENCE W, W.

Additionally, this adds assembly support for the FENCE.TSO, which is
encoded as FENCE RW, RW with the fm field set to 1000.

Change-Id: Ie9c6c8cd24b38b08013032972bd54515eaedd637
Reviewed-on: https://go-review.googlesource.com/c/go/+/758000
Reviewed-by: Meng Zhuo <mengzhuo1203@gmail.com>
Reviewed-by: Junyang Shao <shaojunyang@google.com>
Reviewed-by: Dmitri Shuralyov <dmitshur@google.com>
Reviewed-by: Joel Sing <joel@sing.id.au>
LUCI-TryBot-Result: Go LUCI <golang-scoped@luci-project-accounts.iam.gserviceaccount.com>

3 files changed, 96 insertions, 4 deletions

diff --git a/src/cmd/internal/obj/riscv/cpu.go b/src/cmd/internal/obj/riscv/cpu.go
index 64a1a02a53..433c8e1e89 100644
--- a/src/cmd/internal/obj/riscv/cpu.go
+++ b/src/cmd/internal/obj/riscv/cpu.go
@@ -1552,7 +1552,10 @@ var rmSuffixSet = map[string]uint8{
 	"RMM": RM_RMM,
 }
 
-const rmSuffixBit uint8 = 1 << 7
+const (
+	fenceTsoSuffixBit uint8 = 1 << 0
+	rmSuffixBit       uint8 = 1 << 7
+)
 
 func rmSuffixEncode(s string) (uint8, error) {
 	if s == "" {
@@ -1621,7 +1624,28 @@ const (
 	SPOP_CSR_BEGIN = SPOP_RVV_END
 	SPOP_CSR_END   = SPOP_CSR_BEGIN + 4096
 
-	SPOP_END = SPOP_CSR_END + 1
+	// FENCE operands. 16 special operands are reserved for FENCE flags (4 bits: IORW).
+	SPOP_FENCE_BEGIN = SPOP_CSR_END
+
+	SPOP_FENCE_W SpecialOperand = SPOP_FENCE_BEGIN + iota - 20
+	SPOP_FENCE_R
+	SPOP_FENCE_RW
+	SPOP_FENCE_O
+	SPOP_FENCE_OW
+	SPOP_FENCE_OR
+	SPOP_FENCE_ORW
+	SPOP_FENCE_I
+	SPOP_FENCE_IW
+	SPOP_FENCE_IR
+	SPOP_FENCE_IRW
+	SPOP_FENCE_IO
+	SPOP_FENCE_IOW
+	SPOP_FENCE_IOR
+	SPOP_FENCE_IORW
+
+	SPOP_FENCE_END = SPOP_FENCE_BEGIN + 16
+
+	SPOP_END = SPOP_FENCE_END + 1
 )
 
 var specialOperands = map[SpecialOperand]struct {
@@ -1646,6 +1670,22 @@ var specialOperands = map[SpecialOperand]struct {
 	SPOP_E16: {encoding: 1, name: "E16"},
 	SPOP_E32: {encoding: 2, name: "E32"},
 	SPOP_E64: {encoding: 3, name: "E64"},
+
+	SPOP_FENCE_W:    {encoding: 1, name: "W"},
+	SPOP_FENCE_R:    {encoding: 2, name: "R"},
+	SPOP_FENCE_RW:   {encoding: 3, name: "RW"},
+	SPOP_FENCE_O:    {encoding: 4, name: "O"},
+	SPOP_FENCE_OW:   {encoding: 5, name: "OW"},
+	SPOP_FENCE_OR:   {encoding: 6, name: "OR"},
+	SPOP_FENCE_ORW:  {encoding: 7, name: "ORW"},
+	SPOP_FENCE_I:    {encoding: 8, name: "I"},
+	SPOP_FENCE_IW:   {encoding: 9, name: "IW"},
+	SPOP_FENCE_IR:   {encoding: 10, name: "IR"},
+	SPOP_FENCE_IRW:  {encoding: 11, name: "IRW"},
+	SPOP_FENCE_IO:   {encoding: 12, name: "IO"},
+	SPOP_FENCE_IOW:  {encoding: 13, name: "IOW"},
+	SPOP_FENCE_IOR:  {encoding: 14, name: "IOR"},
+	SPOP_FENCE_IORW: {encoding: 15, name: "IORW"},
 }
 
 func (so SpecialOperand) encode() uint32 {
@@ -1660,6 +1700,11 @@ func (so SpecialOperand) encode() uint32 {
 		if _, ok := csrs[csrNum]; ok {
 			return uint32(csrNum)
 		}
+	case so > SPOP_FENCE_BEGIN && so < SPOP_FENCE_END:
+		op, ok := specialOperands[so]
+		if ok {
+			return op.encoding
+		}
 	}
 	return 0
 }
@@ -1676,6 +1721,11 @@ func (so SpecialOperand) String() string {
 		if csrName, ok := csrs[uint16(so-SPOP_CSR_BEGIN)]; ok {
 			return csrName
 		}
+	case so > SPOP_FENCE_BEGIN && so < SPOP_FENCE_END:
+		op, ok := specialOperands[so]
+		if ok {
+			return op.name
+		}
 	}
 	return ""
 }
diff --git a/src/cmd/internal/obj/riscv/list.go b/src/cmd/internal/obj/riscv/list.go
index da703c02ad..2e064eb823 100644
--- a/src/cmd/internal/obj/riscv/list.go
+++ b/src/cmd/internal/obj/riscv/list.go
@@ -37,6 +37,9 @@ func RegName(r int) string {
 }
 
 func opSuffixString(s uint8) string {
+	if s == fenceTsoSuffixBit {
+		return ".TSO"
+	}
 	if s&rmSuffixBit == 0 {
 		return ""
 	}
diff --git a/src/cmd/internal/obj/riscv/obj.go b/src/cmd/internal/obj/riscv/obj.go
index a55c6faf65..50c687f722 100644
--- a/src/cmd/internal/obj/riscv/obj.go
+++ b/src/cmd/internal/obj/riscv/obj.go
@@ -3700,6 +3700,16 @@ func (ins *instruction) compress() {
 	}
 }
 
+func encodeFenceOperand(a *obj.Addr) (uint32, bool) {
+	if a.Type == obj.TYPE_SPECIAL && a.Offset > int64(SPOP_FENCE_BEGIN) && a.Offset < int64(SPOP_FENCE_END) {
+		return SpecialOperand(a.Offset).encode(), true
+	}
+	if a.Type == obj.TYPE_NONE {
+		return SPOP_FENCE_IORW.encode(), true
+	}
+	return 0, false
+}
+
 // instructionForProg returns the default *obj.Prog to instruction mapping.
 func instructionForProg(p *obj.Prog) *instruction {
 	ins := &instruction{
@@ -4391,7 +4401,25 @@ func instructionsForProg(p *obj.Prog, compress bool) []*instruction {
 
 	case AFENCE:
 		ins.rd, ins.rs1, ins.rs2 = REG_ZERO, REG_ZERO, obj.REG_NONE
-		ins.imm = 0x0ff
+		if p.Scond == fenceTsoSuffixBit {
+			if p.From.Type != obj.TYPE_NONE || p.To.Type != obj.TYPE_NONE {
+				p.Ctxt.Diag("FENCE.TSO must not have operands: %v", p)
+			}
+			// FENCE.TSO is encoded as a FENCE instruction with fm=1000(8), pred=RW(3), succ=RW(3)
+			ins.imm = signExtend((8<<8)|(3<<4)|3, 12)
+		} else {
+			pred, ok := encodeFenceOperand(&p.From)
+			if !ok {
+				p.Ctxt.Diag("invalid FENCE predecessor operand: %v", p)
+			}
+			succ, ok := encodeFenceOperand(&p.To)
+			if !ok {
+				p.Ctxt.Diag("invalid FENCE successor operand: %v", p)
+			}
+			// FENCE pred, succ
+			// pred(4 bits), succ(4 bits)
+			ins.imm = int64((pred << 4) | succ)
+		}
 
 	case AFCVTWS, AFCVTLS, AFCVTWUS, AFCVTLUS, AFCVTWD, AFCVTLD, AFCVTWUD, AFCVTLUD:
 		// Set the default rounding mode in funct3 to round to zero.
@@ -5042,9 +5070,20 @@ func isUnsafePoint(p *obj.Prog) bool {
 }
 
 func ParseSuffix(prog *obj.Prog, cond string) (err error) {
+	cond = strings.TrimPrefix(cond, ".")
 	switch prog.As {
 	case AFCVTWS, AFCVTLS, AFCVTWUS, AFCVTLUS, AFCVTWD, AFCVTLD, AFCVTWUD, AFCVTLUD:
-		prog.Scond, err = rmSuffixEncode(strings.TrimPrefix(cond, "."))
+		prog.Scond, err = rmSuffixEncode(cond)
+	case AFENCE:
+		if cond == "TSO" {
+			prog.Scond = fenceTsoSuffixBit
+		} else {
+			err = fmt.Errorf("unrecognized suffix .%q", cond)
+		}
+	default:
+		if cond != "" {
+			err = fmt.Errorf("unrecognized suffix .%q", cond)
+		}
 	}
 	return
 }