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>

6 files changed, 111 insertions, 4 deletions

diff --git a/src/cmd/asm/internal/arch/riscv64.go b/src/cmd/asm/internal/arch/riscv64.go
index 891d2be0e5..c7bb5cb55c 100644
--- a/src/cmd/asm/internal/arch/riscv64.go
+++ b/src/cmd/asm/internal/arch/riscv64.go
@@ -69,6 +69,10 @@ func RISCV64SpecialOperand(name string) riscv.SpecialOperand {
 			}
 			riscv64SpecialOperand[csrName] = riscv.SpecialOperand(int(csrCode) + int(riscv.SPOP_CSR_BEGIN))
 		}
+		// Add the FENCE operands
+		for opd := riscv.SPOP_FENCE_BEGIN + 1; opd < riscv.SPOP_FENCE_END; opd++ {
+			riscv64SpecialOperand[opd.String()] = opd
+		}
 	}
 	if opd, ok := riscv64SpecialOperand[name]; ok {
 		return opd
diff --git a/src/cmd/asm/internal/asm/testdata/riscv64.s b/src/cmd/asm/internal/asm/testdata/riscv64.s
index 94d1476a6d..65ccf09066 100644
--- a/src/cmd/asm/internal/asm/testdata/riscv64.s
+++ b/src/cmd/asm/internal/asm/testdata/riscv64.s
@@ -140,6 +140,10 @@ start:
 
 	// 2.7: Memory Ordering Instructions
 	FENCE						// 0f00f00f
+	FENCE	W, W					// 0f001001
+	FENCE	I, O					// 0f004008
+	FENCE	IORW, IORW				// 0f00f00f
+	FENCE.TSO					// 0f003083
 
 	// 4.2: Integer Computational Instructions (RV64I)
 	ADDIW	$1, X5, X6				// 1b831200
diff --git a/src/cmd/asm/internal/asm/testdata/riscv64error.s b/src/cmd/asm/internal/asm/testdata/riscv64error.s
index 3c09770d2a..e765fd66ca 100644
--- a/src/cmd/asm/internal/asm/testdata/riscv64error.s
+++ b/src/cmd/asm/internal/asm/testdata/riscv64error.s
@@ -68,6 +68,13 @@ TEXT errors(SB),$0
 	SD	X5, 4294967296(X6)		// ERROR "constant 4294967296 too large"
 	FNES	F1, (X5)			// ERROR "needs an integer register output"
 
+	// Memory Ordering Instructions
+	FENCE		X1, R				// ERROR "invalid FENCE predecessor operand"
+	FENCE		R, X2				// ERROR "invalid FENCE successor operand"
+	FENCE		$1, R				// ERROR "invalid FENCE predecessor operand"
+	FENCE		R, $2				// ERROR "invalid FENCE successor operand"
+	FENCE.TSO	R, R				// ERROR "FENCE.TSO must not have operands"
+
 	//
 	// "V" Standard Extension for Vector Operations, Version 1.0
 	//
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
 }