补充tla.mmad中hf32模式支持-catlass-AtomGit

Mpython/tla_dsl/catlass/_mlir_bindings/tla_ops_gen.py+17-3

	@@ -1830,7 +1830,7 @@ class MmadOp(_ods_ir.OpView):
1830		1830
1831	_ODS_REGIONS = (0, True)	1831	_ODS_REGIONS = (0, True)
1832		1832
1833	- def __init__(self, acc, lhs, rhs, init_c, unit_flag, compute_order, *, loc=None, ip=None):	1833	+ def __init__(self, acc, lhs, rhs, init_c, unit_flag, compute_order, hf32_mode, *, loc=None, ip=None):
1834	operands = []	1834	operands = []
1835	results = []	1835	results = []
1836	attributes = {}	1836	attributes = {}
	@@ -1845,6 +1845,10 @@ class MmadOp(_ods_ir.OpView):
1845	isinstance(compute_order, _ods_ir.Attribute) or	1845	isinstance(compute_order, _ods_ir.Attribute) or
1846	not _ods_ir.AttrBuilder.contains('Tla_ComputeOrderAttr')) else	1846	not _ods_ir.AttrBuilder.contains('Tla_ComputeOrderAttr')) else
1847	_ods_ir.AttrBuilder.get('Tla_ComputeOrderAttr')(compute_order, context=_ods_context))	1847	_ods_ir.AttrBuilder.get('Tla_ComputeOrderAttr')(compute_order, context=_ods_context))
		1848	+ attributes["hf32_mode"] = (hf32_mode if (
		1849	+ isinstance(hf32_mode, _ods_ir.Attribute) or
		1850	+ not _ods_ir.AttrBuilder.contains('Tla_HF32ModeAttr')) else
		1851	+ _ods_ir.AttrBuilder.get('Tla_HF32ModeAttr')(hf32_mode, context=_ods_context))
1848	_ods_successors = None	1852	_ods_successors = None
1849	super().__init__(self.build_generic(attributes=attributes, results=results, operands=operands, successors=_ods_successors, regions=regions, loc=loc, ip=ip))	1853	super().__init__(self.build_generic(attributes=attributes, results=results, operands=operands, successors=_ods_successors, regions=regions, loc=loc, ip=ip))
1850		1854
	@@ -1878,8 +1882,18 @@ class MmadOp(_ods_ir.OpView):
1878	raise ValueError("'None' not allowed as value for mandatory attributes")	1882	raise ValueError("'None' not allowed as value for mandatory attributes")
1879	self.operation.attributes["compute_order"] = value	1883	self.operation.attributes["compute_order"] = value
1880		1884
1881	-def mmad(acc, lhs, rhs, init_c, unit_flag, compute_order, *, loc=None, ip=None) -> _ods_ir.Operation:	1885	+ @builtins.property
1882	- return _get_op_result_or_op_results(MmadOp(acc=acc, lhs=lhs, rhs=rhs, init_c=init_c, unit_flag=unit_flag, compute_order=compute_order, loc=loc, ip=ip))	1886	+ def hf32_mode(self):
		1887	+ return self.operation.attributes["hf32_mode"]
		1888	+
		1889	+ @hf32_mode.setter
		1890	+ def hf32_mode(self, value):
		1891	+ if value is None:
		1892	+ raise ValueError("'None' not allowed as value for mandatory attributes")
		1893	+ self.operation.attributes["hf32_mode"] = value
		1894	+
		1895	+def mmad(acc, lhs, rhs, init_c, unit_flag, compute_order, hf32_mode, *, loc=None, ip=None) -> _ods_ir.Operation:
		1896	+ return _get_op_result_or_op_results(MmadOp(acc=acc, lhs=lhs, rhs=rhs, init_c=init_c, unit_flag=unit_flag, compute_order=compute_order, hf32_mode=hf32_mode, loc=loc, ip=ip))
1883		1897
1884	@_ods_cext.register_operation(_Dialect)	1898	@_ods_cext.register_operation(_Dialect)
1885	class MulOp(_ods_ir.OpView):	1899	class MulOp(_ods_ir.OpView):

Mpython/tla_dsl/catlass/core_api.py+22-1

	@@ -62,7 +62,16 @@ from .types import (
62	dtype_size_bytes,	62	dtype_size_bytes,
63	_replace_flat_leaves_in_tree,	63	_replace_flat_leaves_in_tree,
64	)	64	)
65	-from .params import CopyParams, CopyL0C2DstParams, QuantMode, L0C2UBMode, AtomicMode, ComputeOrder, MemType	65	+from .params import (
		66	+ AtomicMode,
		67	+ HF32Mode,
		68	+ ComputeOrder,
		69	+ CopyL0C2DstParams,
		70	+ CopyParams,
		71	+ L0C2UBMode,
		72	+ MemType,
		73	+ QuantMode,
		74	+)
66		75
67		76
68	_PIPE_VALUES = {	77	_PIPE_VALUES = {
	@@ -5129,6 +5138,7 @@ def mmad(
5129	init_c: bool \| Bool \| None = None,	5138	init_c: bool \| Bool \| None = None,
5130	unit_flag: IndexLike \| None = None,	5139	unit_flag: IndexLike \| None = None,
5131	compute_order: ComputeOrder = ComputeOrder.M_FIRST,	5140	compute_order: ComputeOrder = ComputeOrder.M_FIRST,
		5141	+ hf32_mode: HF32Mode = HF32Mode.HF32_DISABLE,
5132	loc: mlir_ir.Location \| None = None,	5142	loc: mlir_ir.Location \| None = None,
5133	**extra_kwargs: object,	5143	**extra_kwargs: object,
5134	) -> None:	5144	) -> None:
	@@ -5145,6 +5155,8 @@ Description:
5145	- `unit_flag` (`IndexLike \| None`): Unit-flag control bits; defaults to `0`	5155	- `unit_flag` (`IndexLike \| None`): Unit-flag control bits; defaults to `0`
5146	when omitted. Optional, default `None`.	5156	when omitted. Optional, default `None`.
5147	- `compute_order` (`ComputeOrder`): M/N compute-direction priority; default `M_FIRST`.	5157	- `compute_order` (`ComputeOrder`): M/N compute-direction priority; default `M_FIRST`.
		5158	+ - `hf32_mode` (`HF32Mode`): HF32 rounding mode for FP32 operands in L0A/L0B
		5159	+ before the matrix multiply. Optional, default `HF32_DISABLE`.
5148		5160
5149	Constraints:	5161	Constraints:
5150	- Must be called inside a `@tla.kernel`-decorated kernel function.	5162	- Must be called inside a `@tla.kernel`-decorated kernel function.
	@@ -5198,10 +5210,18 @@ Description:
5198	"tla.mmad attribute 'compute_order' must be a "	5210	"tla.mmad attribute 'compute_order' must be a "
5199	f"{ComputeOrder}, got {type(compute_order).__name__}"	5211	f"{ComputeOrder}, got {type(compute_order).__name__}"
5200	)	5212	)
		5213	+ if not isinstance(hf32_mode, HF32Mode):
		5214	+ raise TlaLoweringError(
		5215	+ "tla.mmad attribute 'hf32_mode' must be a "
		5216	+ f"{HF32Mode}, got {type(hf32_mode).__name__}"
		5217	+ )
5201	ctx = loc.context if loc is not None else mlir_ir.Context.current	5218	ctx = loc.context if loc is not None else mlir_ir.Context.current
5202	compute_order_attr = mlir_ir.Attribute.parse(	5219	compute_order_attr = mlir_ir.Attribute.parse(
5203	f"#tla.compute_order<{str(compute_order)}>", context=ctx	5220	f"#tla.compute_order<{str(compute_order)}>", context=ctx
5204	)	5221	)
		5222	+ hf32_mode_attr = mlir_ir.Attribute.parse(
		5223	+ f"#tla.hf32_mode<{str(hf32_mode)}>", context=ctx
		5224	+ )
5205		5225
5206	acc_value = _as_value(acc)	5226	acc_value = _as_value(acc)
5207	lhs_value = _as_value(lhs)	5227	lhs_value = _as_value(lhs)
	@@ -5215,6 +5235,7 @@ Description:
5215	unit_flag_value,	5235	unit_flag_value,
5216	loc=loc,	5236	loc=loc,
5217	compute_order=compute_order_attr,	5237	compute_order=compute_order_attr,
		5238	+ hf32_mode=hf32_mode_attr,
5218	)	5239	)
5219		5240
5220		5241

Mpython/tla_dsl/catlass/params.py+7-0

	@@ -30,6 +30,13 @@ class ComputeOrder(enum.IntEnum):
30	def __str__(self):	30	def __str__(self):
31	return self.name # "M_FIRST"/"N_FIRST"	31	return self.name # "M_FIRST"/"N_FIRST"
32		32
		33	+class HF32Mode(enum.IntEnum):
		34	+ HF32_DISABLE = 0
		35	+ HF32_NEAREST_ZERO = 1
		36	+ HF32_NEAREST_EVEN = 2
		37	+
		38	+ def __str__(self):
		39	+ return self.name # "HF32_DISABLE"/"HF32_NEAREST_ZERO"/"HF32_NEAREST_EVEN"
33		40
34	class AtomicMode(enum.Enum):	41	class AtomicMode(enum.Enum):
35	"""Atomic operation mode"""	42	"""Atomic operation mode"""

Mpython/tla_dsl/csrc/mlir/include/Dialect/Tla/IR/Tla.td+31-1

	@@ -382,6 +382,27 @@ def Tla_ComputeOrderAttr : Tla_EnumAttr<"ComputeOrder", "compute_order"> {
382	}];	382	}];
383	}	383	}
384		384
		385	+// HF32Mode: mmad HF32 rounding mode (CTRL[46] enable, CTRL[47] rounding).
		386	+// HF32_DISABLE(0) is the hardware default (no HF32 rounding).
		387	+def Tla_HF32Mode_Disable : I32EnumAttrCase<"HF32_DISABLE", 0>;
		388	+def Tla_HF32Mode_NearestZero : I32EnumAttrCase<"HF32_NEAREST_ZERO", 1>;
		389	+def Tla_HF32Mode_NearestEven : I32EnumAttrCase<"HF32_NEAREST_EVEN", 2>;
		390	+def Tla_HF32Mode_Enum : Tla_I32Enum<"HF32Mode",
		391	+ "mmad HF32 rounding mode", [
		392	+ Tla_HF32Mode_Disable, Tla_HF32Mode_NearestZero, Tla_HF32Mode_NearestEven
		393	+]>;
		394	+def Tla_HF32ModeAttr : Tla_EnumAttr<"HF32Mode", "hf32_mode"> {
		395	+ let parameters = (ins EnumParameter<Tla_HF32Mode_Enum>:$value);
		396	+ let assemblyFormat = "`<` params `>`";
		397	+ let description = [{
		398	+ mmad HF32 rounding mode, lowered to SPR.CTRL[46]/CTRL[47]:
		399	+
		400	+ - HF32_DISABLE (0): HF32 mode disabled (hardware default)
		401	+ - HF32_NEAREST_ZERO (1): f32 rounded to HF32 towards zero
		402	+ - HF32_NEAREST_EVEN (2): f32 rounded to HF32 to nearest even
		403	+ }];
		404	+}
		405	+
385	// SatMode: overflow behaviour of the cast (AVE `sat` BoolAttr).	406	// SatMode: overflow behaviour of the cast (AVE `sat` BoolAttr).
386	def Tla_SatMode_Unknown : I32EnumAttrCase<"unknown", 0>;	407	def Tla_SatMode_Unknown : I32EnumAttrCase<"unknown", 0>;
387	def Tla_SatMode_Sat : I32EnumAttrCase<"sat", 1>;	408	def Tla_SatMode_Sat : I32EnumAttrCase<"sat", 1>;
	@@ -857,10 +878,19 @@ def Tla_MmadOp : Tla_Op<"mmad", []> {
857	every mmad in a function agrees on the value it is emitted once at the	878	every mmad in a function agrees on the value it is emitted once at the
858	function entry; if a function mixes values it is emitted ahead of each mmad	879	function entry; if a function mixes values it is emitted ahead of each mmad
859	runtime call instead.	880	runtime call instead.
		881	+
		882	+ ``hf32_mode`` is always supplied by the frontend. The lowering emits
		883	+ ``hivm.set_ctrl`` on ``CTRL[46]`` (enable HF32 rounding) and ``CTRL[47]``
		884	+ (rounding mode, ``HF32_NEAREST_ZERO`` -> 1 / ``HF32_NEAREST_EVEN`` -> 0,
		885	+ mirroring ``AscendC::SetHF32Mode`` / ``AscendC::SetHF32TransMode``): when
		886	+ every mmad in a function agrees on the value it is emitted once at the
		887	+ function entry; if a function mixes values it is emitted ahead of each mmad
		888	+ runtime call instead.
860	}];	889	}];
861	let arguments = (ins Tla_TensorType:$acc, Tla_TensorType:$lhs, Tla_TensorType:$rhs,	890	let arguments = (ins Tla_TensorType:$acc, Tla_TensorType:$lhs, Tla_TensorType:$rhs,
862	I1:$init_c, I64:$unit_flag,	891	I1:$init_c, I64:$unit_flag,
863	- Tla_ComputeOrderAttr:$compute_order);	892	+ Tla_ComputeOrderAttr:$compute_order,
		893	+ Tla_HF32ModeAttr:$hf32_mode);
864	let assemblyFormat = [{	894	let assemblyFormat = [{
865	$acc `,` $lhs `,` $rhs `,` $init_c `,` $unit_flag attr-dict `:`	895	$acc `,` $lhs `,` $rhs `,` $init_c `,` $unit_flag attr-dict `:`
866	qualified(type($acc)) `,` qualified(type($lhs)) `,` qualified(type($rhs)) `,` type($init_c) `,` type($unit_flag)	896	qualified(type($acc)) `,` qualified(type($lhs)) `,` qualified(type($rhs)) `,` type($init_c) `,` type($unit_flag)

Mpython/tla_dsl/csrc/mlir/lib/Dialect/Tla/IR/TlaDialect.cpp+14-0

	@@ -215,6 +215,20 @@ Attribute TlaDialect::parseAttribute(DialectAsmParser &parser, Type type) const
215	return ::tla::ComputeOrderAttr::get(getContext(), *symbolized);	215	return ::tla::ComputeOrderAttr::get(getContext(), *symbolized);
216	}	216	}
217		217
		218	+ if (attrTag == "hf32_mode") {
		219	+ if (parser.parseLess())
		220	+ return {};
		221	+ StringRef modeKeyword;
		222	+ if (parser.parseKeyword(&modeKeyword) \|\| parser.parseGreater())
		223	+ return {};
		224	+ auto symbolized = ::symbolizeHF32Mode(modeKeyword);
		225	+ if (!symbolized) {
		226	+ parser.emitError(parser.getNameLoc()) << "invalid tla.hf32_mode value: " << modeKeyword;
		227	+ return {};
		228	+ }
		229	+ return ::tla::HF32ModeAttr::get(getContext(), *symbolized);
		230	+ }
		231	+
218	StringRef mnemonic = attrTag;	232	StringRef mnemonic = attrTag;
219	Attribute value;	233	Attribute value;
220	OptionalParseResult parseResult = generatedAttributeParser(parser, &mnemonic, type, value);	234	OptionalParseResult parseResult = generatedAttributeParser(parser, &mnemonic, type, value);

Mpython/tla_dsl/csrc/mlir/lib/Dialect/Tla/IR/TlaOps.cpp+23-0

	@@ -279,6 +279,29 @@ template <typename AncestorOp> static bool hasEnclosingRegion(mlir::Operation *o
279	mlir::LogicalResult MmadOp::verify() {	279	mlir::LogicalResult MmadOp::verify() {
280	if (!hasEnclosingRegion<CubeOp>(getOperation()))	280	if (!hasEnclosingRegion<CubeOp>(getOperation()))
281	return emitOpError("must be nested inside a tla.cube region");	281	return emitOpError("must be nested inside a tla.cube region");
		282	+
		283	+ // HF32 rounding only applies to f32 L0A/L0B operands; requesting it for any
		284	+ // other source element type (f16/bf16/...) is meaningless and must be rejected.
		285	+ HF32Mode mode = getHf32Mode().getValue();
		286	+ if (mode != HF32Mode::HF32_DISABLE) {
		287	+ auto checkF32SourceOperand = [&](TlaTensorType operandType,
		288	+ llvm::StringRef operandName)
		289	+ -> mlir::LogicalResult {
		290	+ mlir::Type elementType = operandType.getPtr().getPointee();
		291	+ if (!elementType.isF32())
		292	+ return emitOpError()
		293	+ << "hf32_mode " << stringifyHF32Mode(mode)
		294	+ << " requires f32 source operands, but " << operandName
		295	+ << " operand has element type " << elementType;
		296	+ return mlir::success();
		297	+ };
		298	+
		299	+ if (failed(checkF32SourceOperand(getLhs().getType(), "lhs")))
		300	+ return mlir::failure();
		301	+ if (failed(checkF32SourceOperand(getRhs().getType(), "rhs")))
		302	+ return mlir::failure();
		303	+ }
		304	+
282	return mlir::success();	305	return mlir::success();
283	}	306	}
284		307

Mpython/tla_dsl/csrc/mlir/lib/Passes/TlaCubeRegionPass.cpp+49-3

	@@ -20,16 +20,25 @@ namespace {
20	// CTRL[51] selects the mmad M/N compute-direction priority	20	// CTRL[51] selects the mmad M/N compute-direction priority
21	static constexpr unsigned int ComputeOrderBit = 51;	21	static constexpr unsigned int ComputeOrderBit = 51;
22		22
		23	+// CTRL[46] enables the mmad HF32 rounding mode
		24	+static constexpr unsigned int HF32ModeBit = 46;
		25	+
		26	+// CTRL[47] selects the mmad HF32 rounding mode:
		27	+// 0 = NEAREST_EVEN (hardware default), 1 = NEAREST_ZERO
		28	+static constexpr unsigned int HF32TransModeBit = 47;
		29	+
23	struct LowerTlaMmadPattern : public OpRewritePattern<::tla::MmadOp> {	30	struct LowerTlaMmadPattern : public OpRewritePattern<::tla::MmadOp> {
24	LowerTlaMmadPattern(MLIRContext *ctx,	31	LowerTlaMmadPattern(MLIRContext *ctx,
25	DenseMap<Value, TensorDescriptor> &tensorDescriptorByValue,	32	DenseMap<Value, TensorDescriptor> &tensorDescriptorByValue,
26	SmallVectorImpl<Operation *> &toErase,	33	SmallVectorImpl<Operation *> &toErase,
27	DenseMap<Value, Value> &loweredMemrefByValue,	34	DenseMap<Value, Value> &loweredMemrefByValue,
28	- bool funcLevelComputeOrderSet)	35	+ bool funcLevelComputeOrderSet,
		36	+ bool funcLevelHF32Set)
29	: OpRewritePattern<::tla::MmadOp>(ctx),	37	: OpRewritePattern<::tla::MmadOp>(ctx),
30	tensorDescriptorByValue(tensorDescriptorByValue), toErase(toErase),	38	tensorDescriptorByValue(tensorDescriptorByValue), toErase(toErase),
31	loweredMemrefByValue(loweredMemrefByValue),	39	loweredMemrefByValue(loweredMemrefByValue),
32	- funcLevelComputeOrderSet(funcLevelComputeOrderSet) {}	40	+ funcLevelComputeOrderSet(funcLevelComputeOrderSet),
		41	+ funcLevelHF32Set(funcLevelHF32Set) {}
33		42
34	LogicalResult matchAndRewrite(::tla::MmadOp op, PatternRewriter &rewriter) const override {	43	LogicalResult matchAndRewrite(::tla::MmadOp op, PatternRewriter &rewriter) const override {
35	if (op->getNumOperands() < 3)	44	if (op->getNumOperands() < 3)
	@@ -193,6 +202,14 @@ static constexpr unsigned int ComputeOrderBit = 51;
193	bool isNFirst = computeOrderAttr.getValue() == ComputeOrder::N_FIRST;	202	bool isNFirst = computeOrderAttr.getValue() == ComputeOrder::N_FIRST;
194	rewriter.create<hivm::SetCtrlOp>(op.getLoc(), isNFirst, ComputeOrderBit);	203	rewriter.create<hivm::SetCtrlOp>(op.getLoc(), isNFirst, ComputeOrderBit);
195	}	204	}
		205	+ if (!funcLevelHF32Set) {
		206	+ auto modeAttr = op->getAttrOfType<::tla::HF32ModeAttr>("hf32_mode");
		207	+ HF32Mode mode = modeAttr.getValue();
		208	+ bool enableHF32 = mode != HF32Mode::HF32_DISABLE;
		209	+ bool nearestZero = mode == HF32Mode::HF32_NEAREST_ZERO;
		210	+ rewriter.create<hivm::SetCtrlOp>(op.getLoc(), enableHF32, HF32ModeBit);
		211	+ rewriter.create<hivm::SetCtrlOp>(op.getLoc(), nearestZero, HF32TransModeBit);
		212	+ }
196	rewriter.create<func::CallOp>(op.getLoc(), callee, operands);	213	rewriter.create<func::CallOp>(op.getLoc(), callee, operands);
197	toErase.push_back(op.getOperation());	214	toErase.push_back(op.getOperation());
198	return success();	215	return success();
	@@ -203,6 +220,7 @@ static constexpr unsigned int ComputeOrderBit = 51;
203	SmallVectorImpl<Operation *> &toErase;	220	SmallVectorImpl<Operation *> &toErase;
204	DenseMap<Value, Value> &loweredMemrefByValue;	221	DenseMap<Value, Value> &loweredMemrefByValue;
205	bool funcLevelComputeOrderSet;	222	bool funcLevelComputeOrderSet;
		223	+ bool funcLevelHF32Set;
206	};	224	};
207		225
208	struct LowerTlaCopyPattern : public OpRewritePattern<::tla::CopyOp> {	226	struct LowerTlaCopyPattern : public OpRewritePattern<::tla::CopyOp> {
	@@ -534,9 +552,37 @@ public:
534	builder.create<hivm::SetCtrlOp>(funcOp.getLoc(), isNFirst, ComputeOrderBit);	552	builder.create<hivm::SetCtrlOp>(funcOp.getLoc(), isNFirst, ComputeOrderBit);
535	funcLevelComputeOrderSet = true;	553	funcLevelComputeOrderSet = true;
536	}	554	}
		555	+
		556	+ // CTRL[46]/CTRL[47] (mmad HF32 rounding mode) are global and persist once
		557	+ // set, so when every mmad in this function agrees on hf32_mode they are set
		558	+ // once at the function entry. If the function mixes values the per-mmad path
		559	+ // in LowerTlaMmadPattern is used instead.
		560	+ std::optional<HF32Mode> funcLevelHF32Mode;
		561	+ bool HF32Conflict = false;
		562	+ root->walk([&](::tla::MmadOp op) {
		563	+ auto attr = op->getAttrOfType<::tla::HF32ModeAttr>("hf32_mode");
		564	+ HF32Mode mode = attr.getValue();
		565	+ if (funcLevelHF32Mode && *funcLevelHF32Mode != mode)
		566	+ HF32Conflict = true;
		567	+ else if (!funcLevelHF32Mode)
		568	+ funcLevelHF32Mode = mode;
		569	+ });
		570	+ bool funcLevelHF32Set = false;
		571	+ if (funcLevelHF32Mode && !HF32Conflict) {
		572	+ Block &entry = funcOp.getBody().front();
		573	+ PatternRewriter builder(funcOp.getContext());
		574	+ builder.setInsertionPointToStart(&entry);
		575	+ HF32Mode mode = *funcLevelHF32Mode;
		576	+ bool enableHF32 = mode != HF32Mode::HF32_DISABLE;
		577	+ bool nearestZero = mode == HF32Mode::HF32_NEAREST_ZERO;
		578	+ builder.create<hivm::SetCtrlOp>(funcOp.getLoc(), enableHF32, HF32ModeBit);
		579	+ builder.create<hivm::SetCtrlOp>(funcOp.getLoc(), nearestZero, HF32TransModeBit);
		580	+ funcLevelHF32Set = true;
		581	+ }
537	LowerTlaMmadPattern lowerMmad(&getContext(), tensorDescriptorByValue, toErase,	582	LowerTlaMmadPattern lowerMmad(&getContext(), tensorDescriptorByValue, toErase,
538	lowering.loweredMemrefByValue,	583	lowering.loweredMemrefByValue,
539	- funcLevelComputeOrderSet);	584	+ funcLevelComputeOrderSet,
		585	+ funcLevelHF32Set);
540	SmallVector<Operation *, 16> mmadOps;	586	SmallVector<Operation *, 16> mmadOps;
541	root->walk([&](Operation *op) {	587	root->walk([&](Operation *op) {
542	if (llvm::isa<::tla::MmadOp>(op))	588	if (llvm::isa<::tla::MmadOp>(op))

Mpython/tla_dsl/csrc/mlir/lib/Passes/TlaPrologueEpiloguePass.cpp+14-3

	@@ -17,6 +17,12 @@ static constexpr unsigned int MaskControlBit = 56;
17	// CTRL[51] is the mmad M/N compute-direction priority bit.	17	// CTRL[51] is the mmad M/N compute-direction priority bit.
18	static constexpr unsigned int ComputeOrderBit = 51;	18	static constexpr unsigned int ComputeOrderBit = 51;
19		19
		20	+// CTRL[46] is the mmad HF32 rounding-mode enable bit.
		21	+static constexpr unsigned int HF32ModeBit = 46;
		22	+
		23	+// CTRL[47] is the mmad HF32 rounding-mode select bit.
		24	+static constexpr unsigned int HF32TransModeBit = 47;
		25	+
20	class TlaPrologueEpiloguePass	26	class TlaPrologueEpiloguePass
21	: public PassWrapper<TlaPrologueEpiloguePass, OperationPass<ModuleOp>> {	27	: public PassWrapper<TlaPrologueEpiloguePass, OperationPass<ModuleOp>> {
22	public:	28	public:
	@@ -52,10 +58,12 @@ public:
52	Operation *lastBodyOp = terminator ? terminator->getPrevNode() : nullptr;	58	Operation *lastBodyOp = terminator ? terminator->getPrevNode() : nullptr;
53	if (auto barrier = llvm::dyn_cast_or_null<hivm::PipeBarrierOp>(lastBodyOp)) {	59	if (auto barrier = llvm::dyn_cast_or_null<hivm::PipeBarrierOp>(lastBodyOp)) {
54	if (barrier.getPipe().getPipe() == hivm::PIPE::PIPE_ALL) {	60	if (barrier.getPipe().getPipe() == hivm::PIPE::PIPE_ALL) {
55	- // Already has a trailing PIPE_ALL barrier: restore CTRL[51] to the	61	+ // Already has a trailing PIPE_ALL barrier: restore CTRL[51] / CTRL[46] /
56	- // hardware default just before it.	62	+ // CTRL[47] to the hardware defaults just before it.
57	builder.setInsertionPoint(barrier);	63	builder.setInsertionPoint(barrier);
58	builder.create<hivm::SetCtrlOp>(loc, /enable=/false, ComputeOrderBit);	64	builder.create<hivm::SetCtrlOp>(loc, /enable=/false, ComputeOrderBit);
		65	+ builder.create<hivm::SetCtrlOp>(loc, /enable=/false, HF32ModeBit);
		66	+ builder.create<hivm::SetCtrlOp>(loc, /enable=/false, HF32TransModeBit);
59	continue;	67	continue;
60	}	68	}
61	}	69	}
	@@ -64,8 +72,11 @@ public:
64	builder.setInsertionPoint(terminator);	72	builder.setInsertionPoint(terminator);
65	else	73	else
66	builder.setInsertionPointToEnd(&entry);	74	builder.setInsertionPointToEnd(&entry);
67	- // Restore CTRL[51] to the hardware default, then add the trailing barrier.	75	+ // Restore CTRL[51] / CTRL[46] / CTRL[47] to the hardware defaults, then add
		76	+ // the trailing barrier.
68	builder.create<hivm::SetCtrlOp>(loc, /enable=/false, ComputeOrderBit);	77	builder.create<hivm::SetCtrlOp>(loc, /enable=/false, ComputeOrderBit);
		78	+ builder.create<hivm::SetCtrlOp>(loc, /enable=/false, HF32ModeBit);
		79	+ builder.create<hivm::SetCtrlOp>(loc, /enable=/false, HF32TransModeBit);
69	builder.create<hivm::PipeBarrierOp>(loc, pipeAll);	80	builder.create<hivm::PipeBarrierOp>(loc, pipeAll);
70	}	81	}
71	}	82	}

Mpython/tla_dsl/docs/en/kernel_api_reference.md+59-57

	@@ -42,7 +42,7 @@ Construction and views for front-end structured values such as Shape / Coord / S
42		42
43	### `make_shape`	43	### `make_shape`
44		44
45	-Source: [`catlass.core_api.make_shape`](../../catlass/core_api.py#L3508)	45	+Source: [`catlass.core_api.make_shape`](../../catlass/core_api.py#L3517)
46		46
47	Description:	47	Description:
48		48
	@@ -86,7 +86,7 @@ zn_shape = tla.make_shape((16, 8), (16, 4))
86		86
87	### `make_coord`	87	### `make_coord`
88		88
89	-Source: [`catlass.core_api.make_coord`](../../catlass/core_api.py#L3549)	89	+Source: [`catlass.core_api.make_coord`](../../catlass/core_api.py#L3558)
90		90
91	Description:	91	Description:
92		92
	@@ -117,7 +117,7 @@ coord = tla.make_coord(block_row, 0)
117		117
118	### `make_stride`	118	### `make_stride`
119		119
120	-Source: [`catlass.core_api.make_stride`](../../catlass/core_api.py#L3578)	120	+Source: [`catlass.core_api.make_stride`](../../catlass/core_api.py#L3587)
121		121
122	Description:	122	Description:
123		123
	@@ -179,7 +179,7 @@ nz_stride = tla.make_stride((1, 1024), (16, 256))
179		179
180	### `make_layout`	180	### `make_layout`
181		181
182	-Source: [`catlass.core_api.make_layout`](../../catlass/core_api.py#L3638)	182	+Source: [`catlass.core_api.make_layout`](../../catlass/core_api.py#L3647)
183		183
184	Description:	184	Description:
185		185
	@@ -242,7 +242,7 @@ zn = tla.make_layout(
242		242
243	### `tile_view`	243	### `tile_view`
244		244
245	-Source: [`catlass.core_api.tile_view`](../../catlass/core_api.py#L3807)	245	+Source: [`catlass.core_api.tile_view`](../../catlass/core_api.py#L3816)
246		246
247	Description:	247	Description:
248		248
	@@ -277,7 +277,7 @@ tile = tla.tile_view(
277		277
278	### `make_tensor`	278	### `make_tensor`
279		279
280	-Source: [`catlass.core_api.make_tensor`](../../catlass/core_api.py#L3854)	280	+Source: [`catlass.core_api.make_tensor`](../../catlass/core_api.py#L3863)
281		281
282	Description:	282	Description:
283		283
	@@ -322,7 +322,7 @@ tensor = tla.make_tensor(ptr, layout, coord=tla.make_coord(0, 0))
322		322
323	### `make_tensor_like`	323	### `make_tensor_like`
324		324
325	-Source: [`catlass.core_api.make_tensor_like`](../../catlass/core_api.py#L4051)	325	+Source: [`catlass.core_api.make_tensor_like`](../../catlass/core_api.py#L4060)
326		326
327	Description:	327	Description:
328		328
	@@ -356,7 +356,7 @@ dst = tla.make_tensor_like(ptr, like=src_tile, layoutTag=tla.arch.RowMajor)
356		356
357	### `make_ptr`	357	### `make_ptr`
358		358
359	-Source: [`catlass.core_api.make_ptr`](../../catlass/core_api.py#L6979)	359	+Source: [`catlass.core_api.make_ptr`](../../catlass/core_api.py#L7000)
360		360
361	Description:	361	Description:
362		362
	@@ -391,7 +391,7 @@ ptr = tla.make_ptr(tla.Float16, addr, mem_space=tla.AddressSpace.gm)
391		391
392	### `recast_ptr`	392	### `recast_ptr`
393		393
394	-Source: [`catlass.core_api.recast_ptr`](../../catlass/core_api.py#L7033)	394	+Source: [`catlass.core_api.recast_ptr`](../../catlass/core_api.py#L7054)
395		395
396	Description:	396	Description:
397		397
	@@ -427,7 +427,7 @@ Tensor copies between on-chip and global memory, and UB register load/store.
427		427
428	### `copy`	428	### `copy`
429		429
430	-Source: [`catlass.core_api.copy`](../../catlass/core_api.py#L4241)	430	+Source: [`catlass.core_api.copy`](../../catlass/core_api.py#L4250)
431		431
432	Description:	432	Description:
433		433
	@@ -592,7 +592,7 @@ Cube-side matrix multiply-accumulate (`tla.mmad`).
592		592
593	### `mmad`	593	### `mmad`
594		594
595	-Source: [`catlass.core_api.mmad`](../../catlass/core_api.py#L5125)	595	+Source: [`catlass.core_api.mmad`](../../catlass/core_api.py#L5134)
596		596
597	Description:	597	Description:
598		598
	@@ -601,7 +601,7 @@ Emit matrix-multiply-accumulate on TLA tiles.
601	Prototype:	601	Prototype:
602		602
603	```python	603	```python
604	-tla.mmad(acc: Tensor, lhs: Tensor, rhs: Tensor, init_c: bool \| Bool \| None = None, unit_flag: IndexLike \| None = None, compute_order: ComputeOrder = ComputeOrder.M_FIRST, **extra_kwargs: object) -> None	604	+tla.mmad(acc: Tensor, lhs: Tensor, rhs: Tensor, init_c: bool \| Bool \| None = None, unit_flag: IndexLike \| None = None, compute_order: ComputeOrder = ComputeOrder.M_FIRST, hf32_mode: HF32Mode = HF32Mode.HF32_DISABLE, **extra_kwargs: object) -> None
605	```	605	```
606		606
607	Parameters:	607	Parameters:
	@@ -614,6 +614,8 @@ Parameters:
614	- `unit_flag` (`IndexLike \| None`): Unit-flag control bits; defaults to `0`	614	- `unit_flag` (`IndexLike \| None`): Unit-flag control bits; defaults to `0`
615	when omitted. Optional, default `None`.	615	when omitted. Optional, default `None`.
616	- `compute_order` (`ComputeOrder`): M/N compute-direction priority; default `M_FIRST`.	616	- `compute_order` (`ComputeOrder`): M/N compute-direction priority; default `M_FIRST`.
		617	+- `hf32_mode` (`HF32Mode`): HF32 rounding mode for FP32 operands in L0A/L0B
		618	+ before the matrix multiply. Optional, default `HF32_DISABLE`.
617		619
618	Constraints:	620	Constraints:
619		621
	@@ -643,7 +645,7 @@ Mask creation and tail-mask updates.
643		645
644	#### `create_mask`	646	#### `create_mask`
645		647
646	-Source: [`catlass.core_api.create_mask`](../../catlass/core_api.py#L7291)	648	+Source: [`catlass.core_api.create_mask`](../../catlass/core_api.py#L7312)
647		649
648	Description:	650	Description:
649		651
	@@ -699,7 +701,7 @@ with tla.vec.func(mode="simd"):
699		701
700	#### `update_mask`	702	#### `update_mask`
701		703
702	-Source: [`catlass.core_api.update_mask`](../../catlass/core_api.py#L7357)	704	+Source: [`catlass.core_api.update_mask`](../../catlass/core_api.py#L7378)
703		705
704	Description:	706	Description:
705		707
	@@ -736,7 +738,7 @@ Element-wise arithmetic and unary math ops. `VectorSSA` overloads `+` / `-` / `*
736		738
737	#### `exp`	739	#### `exp`
738		740
739	-Source: [`catlass.core_api.exp`](../../catlass/core_api.py#L5710)	741	+Source: [`catlass.core_api.exp`](../../catlass/core_api.py#L5731)
740		742
741	Description:	743	Description:
742		744
	@@ -769,7 +771,7 @@ with tla.vec.func(mode="simd"):
769		771
770	#### `log`	772	#### `log`
771		773
772	-Source: [`catlass.core_api.log`](../../catlass/core_api.py#L5732)	774	+Source: [`catlass.core_api.log`](../../catlass/core_api.py#L5753)
773		775
774	Description:	776	Description:
775		777
	@@ -802,7 +804,7 @@ with tla.vec.func(mode="simd"):
802		804
803	#### `sqrt`	805	#### `sqrt`
804		806
805	-Source: [`catlass.core_api.sqrt`](../../catlass/core_api.py#L5754)	807	+Source: [`catlass.core_api.sqrt`](../../catlass/core_api.py#L5775)
806		808
807	Description:	809	Description:
808		810
	@@ -835,7 +837,7 @@ with tla.vec.func(mode="simd"):
835		837
836	#### `abs`	838	#### `abs`
837		839
838	-Source: [`catlass.core_api.abs`](../../catlass/core_api.py#L5776)	840	+Source: [`catlass.core_api.abs`](../../catlass/core_api.py#L5797)
839		841
840	Description:	842	Description:
841		843
	@@ -868,7 +870,7 @@ with tla.vec.func(mode="simd"):
868		870
869	#### `neg`	871	#### `neg`
870		872
871	-Source: [`catlass.core_api.neg`](../../catlass/core_api.py#L5798)	873	+Source: [`catlass.core_api.neg`](../../catlass/core_api.py#L5819)
872		874
873	Description:	875	Description:
874		876
	@@ -901,7 +903,7 @@ with tla.vec.func(mode="simd"):
901		903
902	#### `add`	904	#### `add`
903		905
904	-Source: [`catlass.core_api.add`](../../catlass/core_api.py#L5978)	906	+Source: [`catlass.core_api.add`](../../catlass/core_api.py#L5999)
905		907
906	Description:	908	Description:
907		909
	@@ -942,7 +944,7 @@ with tla.vec.func(mode="simd"):
942		944
943	#### `sub`	945	#### `sub`
944		946
945	-Source: [`catlass.core_api.sub`](../../catlass/core_api.py#L6024)	947	+Source: [`catlass.core_api.sub`](../../catlass/core_api.py#L6045)
946		948
947	Description:	949	Description:
948		950
	@@ -980,7 +982,7 @@ with tla.vec.func(mode="simd"):
980		982
981	#### `mul`	983	#### `mul`
982		984
983	-Source: [`catlass.core_api.mul`](../../catlass/core_api.py#L6061)	985	+Source: [`catlass.core_api.mul`](../../catlass/core_api.py#L6082)
984		986
985	Description:	987	Description:
986		988
	@@ -1019,7 +1021,7 @@ with tla.vec.func(mode="simd"):
1019		1021
1020	#### `max`	1022	#### `max`
1021		1023
1022	-Source: [`catlass.core_api.max`](../../catlass/core_api.py#L6106)	1024	+Source: [`catlass.core_api.max`](../../catlass/core_api.py#L6127)
1023		1025
1024	Description:	1026	Description:
1025		1027
	@@ -1053,7 +1055,7 @@ with tla.vec.func(mode="simd"):
1053		1055
1054	#### `min`	1056	#### `min`
1055		1057
1056	-Source: [`catlass.core_api.min`](../../catlass/core_api.py#L6146)	1058	+Source: [`catlass.core_api.min`](../../catlass/core_api.py#L6167)
1057		1059
1058	Description:	1060	Description:
1059		1061
	@@ -1087,7 +1089,7 @@ with tla.vec.func(mode="simd"):
1087		1089
1088	#### `div`	1090	#### `div`
1089		1091
1090	-Source: [`catlass.core_api.div`](../../catlass/core_api.py#L6186)	1092	+Source: [`catlass.core_api.div`](../../catlass/core_api.py#L6207)
1091		1093
1092	Description:	1094	Description:
1093		1095
	@@ -1129,7 +1131,7 @@ Bitwise and logical ops on Mask / Vector.
1129		1131
1130	#### `bitwise_not`	1132	#### `bitwise_not`
1131		1133
1132	-Source: [`catlass.core_api.bitwise_not`](../../catlass/core_api.py#L5944)	1134	+Source: [`catlass.core_api.bitwise_not`](../../catlass/core_api.py#L5965)
1133		1135
1134	Description:	1136	Description:
1135		1137
	@@ -1162,7 +1164,7 @@ with tla.vec.func(mode="simd"):
1162		1164
1163	#### `bitwise_and`	1165	#### `bitwise_and`
1164		1166
1165	-Source: [`catlass.core_api.bitwise_and`](../../catlass/core_api.py#L6555)	1167	+Source: [`catlass.core_api.bitwise_and`](../../catlass/core_api.py#L6576)
1166		1168
1167	Description:	1169	Description:
1168		1170
	@@ -1196,7 +1198,7 @@ with tla.vec.func(mode="simd"):
1196		1198
1197	#### `bitwise_or`	1199	#### `bitwise_or`
1198		1200
1199	-Source: [`catlass.core_api.bitwise_or`](../../catlass/core_api.py#L6593)	1201	+Source: [`catlass.core_api.bitwise_or`](../../catlass/core_api.py#L6614)
1200		1202
1201	Description:	1203	Description:
1202		1204
	@@ -1230,7 +1232,7 @@ with tla.vec.func(mode="simd"):
1230		1232
1231	#### `bitwise_xor`	1233	#### `bitwise_xor`
1232		1234
1233	-Source: [`catlass.core_api.bitwise_xor`](../../catlass/core_api.py#L6631)	1235	+Source: [`catlass.core_api.bitwise_xor`](../../catlass/core_api.py#L6652)
1234		1236
1235	Description:	1237	Description:
1236		1238
	@@ -1268,7 +1270,7 @@ Vector compares that produce masks, and masked select.
1268		1270
1269	#### `where`	1271	#### `where`
1270		1272
1271	-Source: [`catlass.core_api.where`](../../catlass/core_api.py#L6270)	1273	+Source: [`catlass.core_api.where`](../../catlass/core_api.py#L6291)
1272		1274
1273	Description:	1275	Description:
1274		1276
	@@ -1302,7 +1304,7 @@ with tla.vec.func(mode="simd"):
1302		1304
1303	#### `cmp`	1305	#### `cmp`
1304		1306
1305	-Source: [`catlass.core_api.cmp`](../../catlass/core_api.py#L6477)	1307	+Source: [`catlass.core_api.cmp`](../../catlass/core_api.py#L6498)
1306		1308
1307	Description:	1309	Description:
1308		1310
	@@ -1341,7 +1343,7 @@ Constant fill and lane-index sequence construction.
1341		1343
1342	#### `full`	1344	#### `full`
1343		1345
1344	-Source: [`catlass.core_api.full`](../../catlass/core_api.py#L5222)	1346	+Source: [`catlass.core_api.full`](../../catlass/core_api.py#L5243)
1345		1347
1346	Description:	1348	Description:
1347		1349
	@@ -1374,7 +1376,7 @@ with tla.vec.func(mode="simd"):
1374		1376
1375	#### `arange`	1377	#### `arange`
1376		1378
1377	-Source: [`catlass.core_api.arange`](../../catlass/core_api.py#L5295)	1379	+Source: [`catlass.core_api.arange`](../../catlass/core_api.py#L5316)
1378		1380
1379	Description:	1381	Description:
1380		1382
	@@ -1412,7 +1414,7 @@ Gather elements from a UB tensor by index.
1412		1414
1413	#### `gather`	1415	#### `gather`
1414		1416
1415	-Source: [`catlass.core_api.gather`](../../catlass/core_api.py#L6669)	1417	+Source: [`catlass.core_api.gather`](../../catlass/core_api.py#L6690)
1416		1418
1417	Description:	1419	Description:
1418		1420
	@@ -1450,7 +1452,7 @@ Interleave / deinterleave and related lane reshuffles.
1450		1452
1451	#### `interleave`	1453	#### `interleave`
1452		1454
1453	-Source: [`catlass.core_api.interleave`](../../catlass/core_api.py#L5835)	1455	+Source: [`catlass.core_api.interleave`](../../catlass/core_api.py#L5856)
1454		1456
1455	Description:	1457	Description:
1456		1458
	@@ -1483,7 +1485,7 @@ with tla.vec.func(mode="simd"):
1483		1485
1484	#### `deinterleave`	1486	#### `deinterleave`
1485		1487
1486	-Source: [`catlass.core_api.deinterleave`](../../catlass/core_api.py#L5889)	1488	+Source: [`catlass.core_api.deinterleave`](../../catlass/core_api.py#L5910)
1487		1489
1488	Description:	1490	Description:
1489		1491
	@@ -1520,7 +1522,7 @@ Compress valid lanes under a mask.
1520		1522
1521	#### `squeeze`	1523	#### `squeeze`
1522		1524
1523	-Source: [`catlass.core_api.squeeze`](../../catlass/core_api.py#L6324)	1525	+Source: [`catlass.core_api.squeeze`](../../catlass/core_api.py#L6345)
1524		1526
1525	Description:	1527	Description:
1526		1528
	@@ -1557,7 +1559,7 @@ In-core / cross-core flags, pipe barriers, mutexes, and local-memory barriers.
1557		1559
1558	### `flag`	1560	### `flag`
1559		1561
1560	-Source: [`catlass.core_api.flag`](../../catlass/core_api.py#L4433)	1562	+Source: [`catlass.core_api.flag`](../../catlass/core_api.py#L4442)
1561		1563
1562	Description:	1564	Description:
1563		1565
	@@ -1595,7 +1597,7 @@ with tla.vector():
1595		1597
1596	### `cross_flag`	1598	### `cross_flag`
1597		1599
1598	-Source: [`catlass.core_api.cross_flag`](../../catlass/core_api.py#L4486)	1600	+Source: [`catlass.core_api.cross_flag`](../../catlass/core_api.py#L4495)
1599		1601
1600	Description:	1602	Description:
1601		1603
	@@ -1627,7 +1629,7 @@ cf = tla.cross_flag("aic_aiv", mode=2)
1627		1629
1628	### `cross_core_set_flag`	1630	### `cross_core_set_flag`
1629		1631
1630	-Source: [`catlass.core_api.cross_core_set_flag`](../../catlass/core_api.py#L4558)	1632	+Source: [`catlass.core_api.cross_core_set_flag`](../../catlass/core_api.py#L4567)
1631		1633
1632	Description:	1634	Description:
1633		1635
	@@ -1662,7 +1664,7 @@ with tla.cube():
1662		1664
1663	### `cross_core_wait_flag`	1665	### `cross_core_wait_flag`
1664		1666
1665	-Source: [`catlass.core_api.cross_core_wait_flag`](../../catlass/core_api.py#L4604)	1667	+Source: [`catlass.core_api.cross_core_wait_flag`](../../catlass/core_api.py#L4613)
1666		1668
1667	Description:	1669	Description:
1668		1670
	@@ -1696,7 +1698,7 @@ with tla.vector():
1696		1698
1697	### `set_flag`	1699	### `set_flag`
1698		1700
1699	-Source: [`catlass.core_api.set_flag`](../../catlass/core_api.py#L4649)	1701	+Source: [`catlass.core_api.set_flag`](../../catlass/core_api.py#L4658)
1700		1702
1701	Description:	1703	Description:
1702		1704
	@@ -1728,7 +1730,7 @@ with tla.vector():
1728		1730
1729	### `wait_flag`	1731	### `wait_flag`
1730		1732
1731	-Source: [`catlass.core_api.wait_flag`](../../catlass/core_api.py#L4675)	1733	+Source: [`catlass.core_api.wait_flag`](../../catlass/core_api.py#L4684)
1732		1734
1733	Description:	1735	Description:
1734		1736
	@@ -1760,7 +1762,7 @@ with tla.vector():
1760		1762
1761	### `pipe_barrier`	1763	### `pipe_barrier`
1762		1764
1763	-Source: [`catlass.core_api.pipe_barrier`](../../catlass/core_api.py#L4701)	1765	+Source: [`catlass.core_api.pipe_barrier`](../../catlass/core_api.py#L4710)
1764		1766
1765	Description:	1767	Description:
1766		1768
	@@ -1792,7 +1794,7 @@ with tla.vector():
1792		1794
1793	### `mutex`	1795	### `mutex`
1794		1796
1795	-Source: [`catlass.core_api.mutex`](../../catlass/core_api.py#L4736)	1797	+Source: [`catlass.core_api.mutex`](../../catlass/core_api.py#L4745)
1796		1798
1797	Description:	1799	Description:
1798		1800
	@@ -1824,7 +1826,7 @@ mtx = tla.mutex("l1_buf", id=0)
1824		1826
1825	### `mutex_guard`	1827	### `mutex_guard`
1826		1828
1827	-Source: [`catlass.core_api.mutex_guard`](../../catlass/core_api.py#L4784)	1829	+Source: [`catlass.core_api.mutex_guard`](../../catlass/core_api.py#L4793)
1828		1830
1829	Description:	1831	Description:
1830		1832
	@@ -1856,7 +1858,7 @@ with tla.mutex_guard(mtx):
1856		1858
1857	### `mutex_lock`	1859	### `mutex_lock`
1858		1860
1859	-Source: [`catlass.core_api.mutex_lock`](../../catlass/core_api.py#L4825)	1861	+Source: [`catlass.core_api.mutex_lock`](../../catlass/core_api.py#L4834)
1860		1862
1861	Description:	1863	Description:
1862		1864
	@@ -1888,7 +1890,7 @@ tla.mutex_lock(mtx, pipe=tla.arch.MTE2)
1888		1890
1889	### `mutex_unlock`	1891	### `mutex_unlock`
1890		1892
1891	-Source: [`catlass.core_api.mutex_unlock`](../../catlass/core_api.py#L4855)	1893	+Source: [`catlass.core_api.mutex_unlock`](../../catlass/core_api.py#L4864)
1892		1894
1893	Description:	1895	Description:
1894		1896
	@@ -1920,7 +1922,7 @@ tla.mutex_unlock(mtx, pipe=tla.arch.MTE2)
1920		1922
1921	### `local_mem_bar`	1923	### `local_mem_bar`
1922		1924
1923	-Source: [`catlass.core_api.local_mem_bar`](../../catlass/core_api.py#L4884)	1925	+Source: [`catlass.core_api.local_mem_bar`](../../catlass/core_api.py#L4893)
1924		1926
1925	Description:	1927	Description:
1926		1928
	@@ -1957,7 +1959,7 @@ Architecture attributes on `tla.arch` (layout tags, pipe identifiers, block help
1957		1959
1958	### `arch`	1960	### `arch`
1959		1961
1960	-Source: [`catlass.core_api.arch`](../../catlass/core_api.py#L7149)	1962	+Source: [`catlass.core_api.arch`](../../catlass/core_api.py#L7170)
1961		1963
1962	Description:	1964	Description:
1963		1965
	@@ -2028,7 +2030,7 @@ On-chip scratch allocation via `allocate`.
2028		2030
2029	### `allocate`	2031	### `allocate`
2030		2032
2031	-Source: [`catlass.core_api.allocate`](../../catlass/core_api.py#L6917)	2033	+Source: [`catlass.core_api.allocate`](../../catlass/core_api.py#L6938)
2032		2034
2033	Description:	2035	Description:
2034		2036
	@@ -2071,7 +2073,7 @@ In-kernel scalar / tensor debug printing.
2071		2073
2072	### `print`	2074	### `print`
2073		2075
2074	-Source: [`catlass.core_api.print`](../../catlass/core_api.py#L3346)	2076	+Source: [`catlass.core_api.print`](../../catlass/core_api.py#L3355)
2075		2077
2076	Description:	2078	Description:
2077		2079
	@@ -2109,7 +2111,7 @@ Cube / Vector / `vec.func` regions and kernel-side loop ranges.
2109		2111
2110	### `range`	2112	### `range`
2111		2113
2112	-Source: [`catlass.core_api.range`](../../catlass/core_api.py#L4934)	2114	+Source: [`catlass.core_api.range`](../../catlass/core_api.py#L4943)
2113		2115
2114	Description:	2116	Description:
2115		2117
	@@ -2143,7 +2145,7 @@ for i in tla.range(0, n, 1):
2143		2145
2144	### `range_constexpr`	2146	### `range_constexpr`
2145		2147
2146	-Source: [`catlass.core_api.range_constexpr`](../../catlass/core_api.py#L4984)	2148	+Source: [`catlass.core_api.range_constexpr`](../../catlass/core_api.py#L4993)
2147		2149
2148	Description:	2150	Description:
2149		2151
	@@ -2177,7 +2179,7 @@ for k in tla.range_constexpr(0, 4):
2177		2179
2178	### `cube`	2180	### `cube`
2179		2181
2180	-Source: [`catlass.core_api.cube`](../../catlass/core_api.py#L5030)	2182	+Source: [`catlass.core_api.cube`](../../catlass/core_api.py#L5039)
2181		2183
2182	Description:	2184	Description:
2183		2185
	@@ -2209,7 +2211,7 @@ with tla.cube():
2209		2211
2210	### `vector`	2212	### `vector`
2211		2213
2212	-Source: [`catlass.core_api.vector`](../../catlass/core_api.py#L5052)	2214	+Source: [`catlass.core_api.vector`](../../catlass/core_api.py#L5061)
2213		2215
2214	Description:	2216	Description:
2215		2217
	@@ -2241,7 +2243,7 @@ with tla.vector():
2241		2243
2242	### `vec.func`	2244	### `vec.func`
2243		2245
2244	-Source: [`catlass.core_api._vec_func`](../../catlass/core_api.py#L5086)	2246	+Source: [`catlass.core_api._vec_func`](../../catlass/core_api.py#L5095)
2245		2247
2246	Description:	2248	Description:
2247		2249

Mpython/tla_dsl/examples/end_to_end/basic_mmad/basic_matmul.py+24-6

	@@ -38,6 +38,7 @@ def basic_mmad_kernel(
38	gm_b: tla.Tensor,	38	gm_b: tla.Tensor,
39	gm_c: tla.Tensor,	39	gm_c: tla.Tensor,
40	_tiling: TilingParams,	40	_tiling: TilingParams,
		41	+ hf32_mode: tla.Constexpr[tla.params.HF32Mode],
41	) -> None:	42	) -> None:
42	c0 = 0	43	c0 = 0
43	c1 = 1	44	c1 = 1
	@@ -216,7 +217,10 @@ def basic_mmad_kernel(
216	else 0b10	217	else 0b10
217	)	218	)
218	init_c = True if k_l1 == 0 and k_l0 == 0 else False	219	init_c = True if k_l1 == 0 and k_l0 == 0 else False
219	- tla.mmad(l0_c, l0_a, l0_b, init_c=init_c, unit_flag=unit_flag)	220	+ if tla.const_expr(hf32_mode != tla.params.HF32Mode.HF32_DISABLE and dtype_a == tla.Float32 and dtype_b == tla.Float32):
		221	+ tla.mmad(l0_c, l0_a, l0_b, init_c=init_c, unit_flag=unit_flag, hf32_mode=hf32_mode)
		222	+ else:
		223	+ tla.mmad(l0_c, l0_a, l0_b, init_c=init_c, unit_flag=unit_flag)
220	if l0_buf_idx == c0:	224	if l0_buf_idx == c0:
221	tla.set_flag(l0a0_available)	225	tla.set_flag(l0a0_available)
222	tla.set_flag(l0b0_available)	226	tla.set_flag(l0b0_available)
	@@ -247,6 +251,7 @@ def run(args: argparse.Namespace) -> int:
247	get_block_num,	251	get_block_num,
248	create_tla_tensor,	252	create_tla_tensor,
249	compare,	253	compare,
		254	+ to_hf32,
250	)	255	)
251		256
252	torch.npu.set_device(args.device)	257	torch.npu.set_device(args.device)
	@@ -260,12 +265,20 @@ def run(args: argparse.Namespace) -> int:
260	dtype_a = dtypes[args.dtype_a]	265	dtype_a = dtypes[args.dtype_a]
261	dtype_b = dtypes[args.dtype_b]	266	dtype_b = dtypes[args.dtype_b]
262	dtype_c = dtypes[args.dtype_c]	267	dtype_c = dtypes[args.dtype_c]
		268	+
263	a = torch.rand(args.m, args.k, dtype=dtype_a, device="cpu") * 10.0 - 5.0	269	a = torch.rand(args.m, args.k, dtype=dtype_a, device="cpu") * 10.0 - 5.0
264	b = torch.rand(args.k, args.n, dtype=dtype_b, device="cpu") * 10.0 - 5.0	270	b = torch.rand(args.k, args.n, dtype=dtype_b, device="cpu") * 10.0 - 5.0
265	c = torch.rand(args.m, args.n, dtype=dtype_c, device="cpu") * 10.0 - 5.0	271	c = torch.rand(args.m, args.n, dtype=dtype_c, device="cpu") * 10.0 - 5.0
266	- ref = a.float() @ b.float()	272	+
267	- if dtype_c in (torch.float16, torch.bfloat16):	273	+ hf32_mode = tla.params.HF32Mode.HF32_NEAREST_EVEN
268	- ref = ref.to(dtype_c).float()	274	+ enable_hf32 = (hf32_mode != tla.params.HF32Mode.HF32_DISABLE and
		275	+ dtype_a == torch.float32 and dtype_b == torch.float32)
		276	+ if enable_hf32:
		277	+ ref = to_hf32(a, hf32_mode) @ to_hf32(b, hf32_mode)
		278	+ else:
		279	+ ref = a.float() @ b.float()
		280	+ if dtype_c in (torch.float16, torch.bfloat16):
		281	+ ref = ref.to(dtype_c).float()
269		282
270	a = (	283	a = (
271	a.contiguous() if args.layout_a == "row" else a.permute(1, 0).contiguous()	284	a.contiguous() if args.layout_a == "row" else a.permute(1, 0).contiguous()
	@@ -283,14 +296,19 @@ def run(args: argparse.Namespace) -> int:
283	a_tensor,	296	a_tensor,
284	b_tensor,	297	b_tensor,
285	c_tensor,	298	c_tensor,
286	- TilingParams(), # default tiling: L1: (256, 256, 128); L0: (256, 256, 32)	299	+ TilingParams(), # default tiling: L1: (256, 256, 128); L0: (256, 256, 32)
		300	+ hf32_mode,
287	options="--npu-arch 3510",	301	options="--npu-arch 3510",
288	)	302	)
289	block_num = get_block_num(args.block_num, args.device, kind="cube")	303	block_num = get_block_num(args.block_num, args.device, kind="cube")
290	artifact(a_tensor, b_tensor, c_tensor, block_num=block_num)	304	artifact(a_tensor, b_tensor, c_tensor, block_num=block_num)
291	torch.npu.synchronize()	305	torch.npu.synchronize()
292		306
293	- passed = compare(c.detach().cpu(), ref, args.k)	307	+ result = c.detach().cpu()
		308	+ if enable_hf32:
		309	+ passed = compare(result, ref, enable_hf32=True)
		310	+ else:
		311	+ passed = compare(result, ref, args.k)
294	print(f"passed={passed} cache_key={artifact.cache_key}")	312	print(f"passed={passed} cache_key={artifact.cache_key}")
295	print(f"kernel.o={artifact.kernel_binary_path}")	313	print(f"kernel.o={artifact.kernel_binary_path}")
296	return 0 if passed else 1	314	return 0 if passed else 1

Mpython/tla_dsl/examples/end_to_end/common/__init__.py+2-1

	@@ -10,7 +10,7 @@
10		10
11	from .golden import compare, tolerance	11	from .golden import compare, tolerance
12	from .params import TilingParams, SwizzleParams	12	from .params import TilingParams, SwizzleParams
13	-from .utils import create_tla_tensor, get_block_num	13	+from .utils import create_tla_tensor, get_block_num, to_hf32
14		14
15	__all__ = [	15	__all__ = [
16	# struct-like params	16	# struct-like params
	@@ -20,6 +20,7 @@ __all__ = [
20	# helper function	20	# helper function
21	"create_tla_tensor",	21	"create_tla_tensor",
22	"get_block_num",	22	"get_block_num",
		23	+ "to_hf32",
23		24
24	# golden compare	25	# golden compare
25	"compare",	26	"compare",

Mpython/tla_dsl/examples/end_to_end/common/golden.py+40-4

	@@ -12,6 +12,12 @@ from typing import overload
12		12
13	import torch	13	import torch
14		14
		15	+# HF32 mixed tolerance thresholds.
		16	+_HF32_RTOL = 2.0 ** -9
		17	+_HF32_ATOL = 2.0 ** -10
		18	+_HF32_REQUIRED_MATCHED_RATIO = 0.99
		19	+_HF32_MAX_ABS_ERROR_LIMIT = 1e-1
		20	+
15		21
16	def tolerance(	22	def tolerance(
17	expected: torch.Tensor,	23	expected: torch.Tensor,
	@@ -32,6 +38,27 @@ def tolerance(
32	return rtol * torch.maximum(torch.full_like(expected, floor), expected.abs())	38	return rtol * torch.maximum(torch.full_like(expected, floor), expected.abs())
33		39
34		40
		41	+def _compare_hf32(
		42	+ result: torch.Tensor,
		43	+ expected: torch.Tensor,
		44	+) -> bool:
		45	+ """Compare result against an HF32-semantic golden."""
		46	+ result = result.float()
		47	+ expected = expected.float()
		48	+ diff = (result - expected).abs()
		49	+ matched = diff <= _HF32_ATOL + _HF32_RTOL * expected.abs()
		50	+ matched_ratio = matched.float().mean().item()
		51	+ max_abs_error = diff.max().item()
		52	+ ulp = torch.finfo(expected.dtype).eps * (
		53	+ 2.0 ** torch.floor(torch.log2(expected.abs().max()))
		54	+ )
		55	+ max_abs_error_limit = max(_HF32_MAX_ABS_ERROR_LIMIT, 32.0 * ulp.item())
		56	+ return bool(
		57	+ matched_ratio >= _HF32_REQUIRED_MATCHED_RATIO
		58	+ and max_abs_error <= max_abs_error_limit
		59	+ )
		60	+
		61	+
35	@overload	62	@overload
36	def compare(	63	def compare(
37	result: torch.Tensor,	64	result: torch.Tensor,
	@@ -48,6 +75,7 @@ def compare(
48	result: torch.Tensor,	75	result: torch.Tensor,
49	expected: torch.Tensor,	76	expected: torch.Tensor,
50	*,	77	*,
		78	+ enable_hf32: bool = False,
51	rtol: float = 0.0,	79	rtol: float = 0.0,
52	atol: float = 0.0,	80	atol: float = 0.0,
53	) -> bool: ...	81	) -> bool: ...
	@@ -61,25 +89,33 @@ def compare(
61	rtol: float \| None = None,	89	rtol: float \| None = None,
62	floor: float \| None = None,	90	floor: float \| None = None,
63	atol: float = 0.0,	91	atol: float = 0.0,
		92	+ enable_hf32: bool = False,
64	) -> bool:	93	) -> bool:
65	"""Compare ``result`` against ``expected`` with given threshold.	94	"""Compare ``result`` against ``expected`` with given threshold.
66		95
67	- Two call forms are supported:	96	+ Three call forms are supported:
68		97
69	1. Accumulative precision standard (``k`` given), for matmul-like operators.	98	1. Accumulative precision standard (``k`` given), for matmul-like operators.
70	2. Generic element-wise check.	99	2. Generic element-wise check.
		100	+ 3. HF32 mixed tolerance (``enable_hf32=True``).
71	"""	101	"""
		102	+ if enable_hf32:
		103	+ # Use mixed tolerance for HF32.
		104	+ return _compare_hf32(result, expected)
72	if k is not None and isinstance(k, int):	105	if k is not None and isinstance(k, int):
73	- is_bf16= (result.dtype == torch.bfloat16)	106	+ # Single precision standard.
		107	+ is_bf16 = (result.dtype == torch.bfloat16)
74	result, expected = result.float(), expected.float()	108	result, expected = result.float(), expected.float()
75	return bool(	109	return bool(
76	- ((result - expected).abs() <=	110	+ (
77	- tolerance(expected, k, bf16=is_bf16, rtol=rtol, floor=floor)	111	+ (result - expected).abs()
		112	+ <= tolerance(expected, k, bf16=is_bf16, rtol=rtol, floor=floor)
78	).all()	113	).all()
79	)	114	)
80	if result.dtype != expected.dtype:	115	if result.dtype != expected.dtype:
81	raise TypeError("the data type between the golden and the result do not match")	116	raise TypeError("the data type between the golden and the result do not match")
82	if result.dtype in (torch.float32, torch.float16, torch.bfloat16):	117	if result.dtype in (torch.float32, torch.float16, torch.bfloat16):
		118	+ # Generic element-wise precision check.
83	return bool(	119	return bool(
84	torch.isclose(	120	torch.isclose(
85	result, expected, rtol=0.0 if rtol is None else rtol, atol=atol	121	result, expected, rtol=0.0 if rtol is None else rtol, atol=atol

Mpython/tla_dsl/examples/end_to_end/common/utils.py+31-0

	@@ -37,3 +37,34 @@ def create_tla_tensor(buf, layout: str):
37	tag = tla.arch.RowMajor if layout == "row" else tla.arch.ColumnMajor	37	tag = tla.arch.RowMajor if layout == "row" else tla.arch.ColumnMajor
38	return from_dlpack(buf.contiguous(), layout_tag=tag).mark_layout_dynamic()	38	return from_dlpack(buf.contiguous(), layout_tag=tag).mark_layout_dynamic()
39		39
		40	+def to_hf32(
		41	+ x: torch.Tensor,
		42	+ hf32_mode: tla.params.HF32Mode,
		43	+) -> torch.Tensor:
		44	+ """Simulate HF32 rounding mode on f32 values.
		45	+
		46	+ HF32 keeps the FP32 sign and 8-bit exponent, and reduces the mantissa to
		47	+ 11 significant bits (10 explicit mantissa bits, close to FP16):
		48	+
		49	+ - ``HF32_NEAREST_ZERO`` rounds to nearest, ties toward zero.
		50	+ - ``HF32_NEAREST_EVEN`` rounds them to nearest-even.
		51	+ """
		52	+ if not isinstance(hf32_mode, tla.params.HF32Mode):
		53	+ raise TypeError(
		54	+ "hf32_mode must be a tla.params.HF32Mode, "
		55	+ f"got {type(hf32_mode).__name__}"
		56	+ )
		57	+ x = x.float()
		58	+ if hf32_mode == tla.params.HF32Mode.HF32_DISABLE:
		59	+ return x
		60	+
		61	+ bits = x.contiguous().view(torch.int32)
		62	+ if hf32_mode == tla.params.HF32Mode.HF32_NEAREST_ZERO:
		63	+ rounded_bits = (bits + 0x0FFF) & ~0x1FFF
		64	+ elif hf32_mode == tla.params.HF32Mode.HF32_NEAREST_EVEN:
		65	+ lsb = (bits >> 13) & 1
		66	+ rounded_bits = (bits + 0x0FFF + lsb) & ~0x1FFF
		67	+ else:
		68	+ raise ValueError(f"Unsupported HF32 mode: {hf32_mode!r}")
		69	+
		70	+ return rounded_bits.view(torch.float32)

Mpython/tla_dsl/tests/lit/tla-compile/auto-mutex-diagnostics.test+2-2

	@@ -148,7 +148,7 @@ module {
148	%c = "tla.make_tensor_like"(%pc, %gm) {layoutTag = "L0Clayout"} : (!tla.ptr<f32, l0c, 512>, !gm) -> !l0c	148	%c = "tla.make_tensor_like"(%pc, %gm) {layoutTag = "L0Clayout"} : (!tla.ptr<f32, l0c, 512>, !gm) -> !l0c
149	%init = arith.constant true	149	%init = arith.constant true
150	"tla.cube"() ({	150	"tla.cube"() ({
151	- tla.mmad %c, %a, %b, %init, %unit {compute_order = #tla.compute_order<M_FIRST>} : !l0c, !l0a, !l0b, i1, i64	151	+ tla.mmad %c, %a, %b, %init, %unit {compute_order = #tla.compute_order<M_FIRST>, hf32_mode = #tla.hf32_mode<HF32_DISABLE>} : !l0c, !l0a, !l0b, i1, i64
152	}) : () -> ()	152	}) : () -> ()
153	tla.return	153	tla.return
154	}	154	}
	@@ -176,7 +176,7 @@ module {
176	scf.yield %two : i64	176	scf.yield %two : i64
177	}	177	}
178	"tla.cube"() ({	178	"tla.cube"() ({
179	- tla.mmad %c, %a, %b, %init, %unit {compute_order = #tla.compute_order<M_FIRST>} : !mixed_l0c, !mixed_l0a, !mixed_l0b, i1, i64	179	+ tla.mmad %c, %a, %b, %init, %unit {compute_order = #tla.compute_order<M_FIRST>, hf32_mode = #tla.hf32_mode<HF32_DISABLE>} : !mixed_l0c, !mixed_l0a, !mixed_l0b, i1, i64
180	}) : () -> ()	180	}) : () -> ()
181	tla.return	181	tla.return
182	}	182	}

Mpython/tla_dsl/tests/lit/tla-compile/auto-mutex-id-spaces.test+1-1

	@@ -96,7 +96,7 @@ module {
96	%init = arith.constant true	96	%init = arith.constant true
97	%unit = arith.constant 0 : i64	97	%unit = arith.constant 0 : i64
98	"tla.cube"() ({	98	"tla.cube"() ({
99	- tla.mmad %c, %a, %b, %init, %unit {compute_order = #tla.compute_order<M_FIRST>} : !l0c, !l0a, !l0b, i1, i64	99	+ tla.mmad %c, %a, %b, %init, %unit {compute_order = #tla.compute_order<M_FIRST>, hf32_mode = #tla.hf32_mode<HF32_DISABLE>} : !l0c, !l0a, !l0b, i1, i64
100	}) : () -> ()	100	}) : () -> ()
101	"tla.vector"() ({	101	"tla.vector"() ({
102	tla.copy %t0, %gm_ub : !ub, !gm_ub	102	tla.copy %t0, %gm_ub : !ub, !gm_ub

Mpython/tla_dsl/tests/lit/tla-compile/auto-mutex-instruction-pipes.mlir+2-2

	@@ -31,7 +31,7 @@ module {
31	"tla.cube"() ({	31	"tla.cube"() ({
32	tla.copy %l1, %gm : !l1, !gm	32	tla.copy %l1, %gm : !l1, !gm
33	tla.copy %l0a, %l1 : !l0a, !l1	33	tla.copy %l0a, %l1 : !l0a, !l1
34	- tla.mmad %l0c, %l0a, %l0b, %init, %unit {compute_order = #tla.compute_order<M_FIRST>} : !l0c, !l0a, !l0b, i1, i64	34	+ tla.mmad %l0c, %l0a, %l0b, %init, %unit {compute_order = #tla.compute_order<M_FIRST>, hf32_mode = #tla.hf32_mode<HF32_DISABLE>} : !l0c, !l0a, !l0b, i1, i64
35	"tla.copy"(%ub, %l0c, %params) : (!ub, !l0c, !tla.copy_l0c2dst_params) -> ()	35	"tla.copy"(%ub, %l0c, %params) : (!ub, !l0c, !tla.copy_l0c2dst_params) -> ()
36	}) : () -> ()	36	}) : () -> ()
37	"tla.vector"() ({	37	"tla.vector"() ({
	@@ -112,7 +112,7 @@ module {
112	%unit64 = arith.extsi %unit32 : i32 to i64	112	%unit64 = arith.extsi %unit32 : i32 to i64
113	%params = "tla.CopyL0C2DstParams"() <{unit_flag = 3 : i64, relu_enable = false, quant_mode = #tla.quant_mode<NO_QUANT>, l0c2ub_mode = #tla.l0c2ub_mode<NO_SPLIT_VEC_0>}> : () -> !tla.copy_l0c2dst_params	113	%params = "tla.CopyL0C2DstParams"() <{unit_flag = 3 : i64, relu_enable = false, quant_mode = #tla.quant_mode<NO_QUANT>, l0c2ub_mode = #tla.l0c2ub_mode<NO_SPLIT_VEC_0>}> : () -> !tla.copy_l0c2dst_params
114	"tla.cube"() ({	114	"tla.cube"() ({
115	- tla.mmad %c, %a, %b, %init, %unit64 {compute_order = #tla.compute_order<M_FIRST>} : !unit_l0c, !unit_l0a, !unit_l0b, i1, i64	115	+ tla.mmad %c, %a, %b, %init, %unit64 {compute_order = #tla.compute_order<M_FIRST>, hf32_mode = #tla.hf32_mode<HF32_DISABLE>} : !unit_l0c, !unit_l0a, !unit_l0b, i1, i64
116	"tla.copy"(%ub, %c, %params) : (!unit_ub, !unit_l0c, !tla.copy_l0c2dst_params) -> ()	116	"tla.copy"(%ub, %c, %params) : (!unit_ub, !unit_l0c, !tla.copy_l0c2dst_params) -> ()
117	}) : () -> ()	117	}) : () -> ()
118	tla.return	118	tla.return

Mpython/tla_dsl/tests/lit/tla-compile/mmad-compute-order.mlir+6-6

	@@ -27,7 +27,7 @@ module {
27	tla.cube {	27	tla.cube {
28	%true = arith.constant true	28	%true = arith.constant true
29	%c3_i64 = arith.constant 3 : i64	29	%c3_i64 = arith.constant 3 : i64
30	- tla.mmad %11, %7, %9, %true, %c3_i64 {compute_order = #tla.compute_order<N_FIRST>} : !tla.tensor<!tla.layout<!tla.shape<(16,2),(16,2)>, !tla.stride<(16,256),(1,512)>, !tla.shape<32,32>, L0Clayout>, !tla.coord<0,0>, !tla.ptr<f32, l0c, 512>>, !tla.tensor<!tla.layout<!tla.shape<(16,2),(8,4)>, !tla.stride<(8,128),(1,256)>, !tla.shape<32,32>, zN>, !tla.coord<0,0>, !tla.ptr<f32, l0a, 512>>, !tla.tensor<!tla.layout<!tla.shape<(8,4),(16,2)>, !tla.stride<(1,256),(8,128)>, !tla.shape<32,32>, nZ>, !tla.coord<0,0>, !tla.ptr<f32, l0b, 512>>, i1, i64	30	+ tla.mmad %11, %7, %9, %true, %c3_i64 {compute_order = #tla.compute_order<N_FIRST>, hf32_mode = #tla.hf32_mode<HF32_DISABLE>} : !tla.tensor<!tla.layout<!tla.shape<(16,2),(16,2)>, !tla.stride<(16,256),(1,512)>, !tla.shape<32,32>, L0Clayout>, !tla.coord<0,0>, !tla.ptr<f32, l0c, 512>>, !tla.tensor<!tla.layout<!tla.shape<(16,2),(8,4)>, !tla.stride<(8,128),(1,256)>, !tla.shape<32,32>, zN>, !tla.coord<0,0>, !tla.ptr<f32, l0a, 512>>, !tla.tensor<!tla.layout<!tla.shape<(8,4),(16,2)>, !tla.stride<(1,256),(8,128)>, !tla.shape<32,32>, nZ>, !tla.coord<0,0>, !tla.ptr<f32, l0b, 512>>, i1, i64
31	}	31	}
32	tla.return	32	tla.return
33	}	33	}
	@@ -47,7 +47,7 @@ module {
47	tla.cube {	47	tla.cube {
48	%true = arith.constant true	48	%true = arith.constant true
49	%c3_i64 = arith.constant 3 : i64	49	%c3_i64 = arith.constant 3 : i64
50	- tla.mmad %11, %7, %9, %true, %c3_i64 {compute_order = #tla.compute_order<M_FIRST>} : !tla.tensor<!tla.layout<!tla.shape<(16,2),(16,2)>, !tla.stride<(16,256),(1,512)>, !tla.shape<32,32>, L0Clayout>, !tla.coord<0,0>, !tla.ptr<f32, l0c, 512>>, !tla.tensor<!tla.layout<!tla.shape<(16,2),(8,4)>, !tla.stride<(8,128),(1,256)>, !tla.shape<32,32>, zN>, !tla.coord<0,0>, !tla.ptr<f32, l0a, 512>>, !tla.tensor<!tla.layout<!tla.shape<(8,4),(16,2)>, !tla.stride<(1,256),(8,128)>, !tla.shape<32,32>, nZ>, !tla.coord<0,0>, !tla.ptr<f32, l0b, 512>>, i1, i64	50	+ tla.mmad %11, %7, %9, %true, %c3_i64 {compute_order = #tla.compute_order<M_FIRST>, hf32_mode = #tla.hf32_mode<HF32_DISABLE>} : !tla.tensor<!tla.layout<!tla.shape<(16,2),(16,2)>, !tla.stride<(16,256),(1,512)>, !tla.shape<32,32>, L0Clayout>, !tla.coord<0,0>, !tla.ptr<f32, l0c, 512>>, !tla.tensor<!tla.layout<!tla.shape<(16,2),(8,4)>, !tla.stride<(8,128),(1,256)>, !tla.shape<32,32>, zN>, !tla.coord<0,0>, !tla.ptr<f32, l0a, 512>>, !tla.tensor<!tla.layout<!tla.shape<(8,4),(16,2)>, !tla.stride<(1,256),(8,128)>, !tla.shape<32,32>, nZ>, !tla.coord<0,0>, !tla.ptr<f32, l0b, 512>>, i1, i64
51	}	51	}
52	tla.return	52	tla.return
53	}	53	}
	@@ -68,8 +68,8 @@ module {
68	%true = arith.constant true	68	%true = arith.constant true
69	%false = arith.constant false	69	%false = arith.constant false
70	%c3_i64 = arith.constant 3 : i64	70	%c3_i64 = arith.constant 3 : i64
71	- tla.mmad %11, %7, %9, %true, %c3_i64 {compute_order = #tla.compute_order<N_FIRST>} : !tla.tensor<!tla.layout<!tla.shape<(16,2),(16,2)>, !tla.stride<(16,256),(1,512)>, !tla.shape<32,32>, L0Clayout>, !tla.coord<0,0>, !tla.ptr<f32, l0c, 512>>, !tla.tensor<!tla.layout<!tla.shape<(16,2),(8,4)>, !tla.stride<(8,128),(1,256)>, !tla.shape<32,32>, zN>, !tla.coord<0,0>, !tla.ptr<f32, l0a, 512>>, !tla.tensor<!tla.layout<!tla.shape<(8,4),(16,2)>, !tla.stride<(1,256),(8,128)>, !tla.shape<32,32>, nZ>, !tla.coord<0,0>, !tla.ptr<f32, l0b, 512>>, i1, i64	71	+ tla.mmad %11, %7, %9, %true, %c3_i64 {compute_order = #tla.compute_order<N_FIRST>, hf32_mode = #tla.hf32_mode<HF32_DISABLE>} : !tla.tensor<!tla.layout<!tla.shape<(16,2),(16,2)>, !tla.stride<(16,256),(1,512)>, !tla.shape<32,32>, L0Clayout>, !tla.coord<0,0>, !tla.ptr<f32, l0c, 512>>, !tla.tensor<!tla.layout<!tla.shape<(16,2),(8,4)>, !tla.stride<(8,128),(1,256)>, !tla.shape<32,32>, zN>, !tla.coord<0,0>, !tla.ptr<f32, l0a, 512>>, !tla.tensor<!tla.layout<!tla.shape<(8,4),(16,2)>, !tla.stride<(1,256),(8,128)>, !tla.shape<32,32>, nZ>, !tla.coord<0,0>, !tla.ptr<f32, l0b, 512>>, i1, i64
72	- tla.mmad %11, %7, %9, %false, %c3_i64 {compute_order = #tla.compute_order<N_FIRST>} : !tla.tensor<!tla.layout<!tla.shape<(16,2),(16,2)>, !tla.stride<(16,256),(1,512)>, !tla.shape<32,32>, L0Clayout>, !tla.coord<0,0>, !tla.ptr<f32, l0c, 512>>, !tla.tensor<!tla.layout<!tla.shape<(16,2),(8,4)>, !tla.stride<(8,128),(1,256)>, !tla.shape<32,32>, zN>, !tla.coord<0,0>, !tla.ptr<f32, l0a, 512>>, !tla.tensor<!tla.layout<!tla.shape<(8,4),(16,2)>, !tla.stride<(1,256),(8,128)>, !tla.shape<32,32>, nZ>, !tla.coord<0,0>, !tla.ptr<f32, l0b, 512>>, i1, i64	72	+ tla.mmad %11, %7, %9, %false, %c3_i64 {compute_order = #tla.compute_order<N_FIRST>, hf32_mode = #tla.hf32_mode<HF32_DISABLE>} : !tla.tensor<!tla.layout<!tla.shape<(16,2),(16,2)>, !tla.stride<(16,256),(1,512)>, !tla.shape<32,32>, L0Clayout>, !tla.coord<0,0>, !tla.ptr<f32, l0c, 512>>, !tla.tensor<!tla.layout<!tla.shape<(16,2),(8,4)>, !tla.stride<(8,128),(1,256)>, !tla.shape<32,32>, zN>, !tla.coord<0,0>, !tla.ptr<f32, l0a, 512>>, !tla.tensor<!tla.layout<!tla.shape<(8,4),(16,2)>, !tla.stride<(1,256),(8,128)>, !tla.shape<32,32>, nZ>, !tla.coord<0,0>, !tla.ptr<f32, l0b, 512>>, i1, i64
73	}	73	}
74	tla.return	74	tla.return
75	}	75	}
	@@ -90,8 +90,8 @@ module {
90	%true = arith.constant true	90	%true = arith.constant true
91	%false = arith.constant false	91	%false = arith.constant false
92	%c3_i64 = arith.constant 3 : i64	92	%c3_i64 = arith.constant 3 : i64
93	- tla.mmad %11, %7, %9, %true, %c3_i64 {compute_order = #tla.compute_order<N_FIRST>} : !tla.tensor<!tla.layout<!tla.shape<(16,2),(16,2)>, !tla.stride<(16,256),(1,512)>, !tla.shape<32,32>, L0Clayout>, !tla.coord<0,0>, !tla.ptr<f32, l0c, 512>>, !tla.tensor<!tla.layout<!tla.shape<(16,2),(8,4)>, !tla.stride<(8,128),(1,256)>, !tla.shape<32,32>, zN>, !tla.coord<0,0>, !tla.ptr<f32, l0a, 512>>, !tla.tensor<!tla.layout<!tla.shape<(8,4),(16,2)>, !tla.stride<(1,256),(8,128)>, !tla.shape<32,32>, nZ>, !tla.coord<0,0>, !tla.ptr<f32, l0b, 512>>, i1, i64	93	+ tla.mmad %11, %7, %9, %true, %c3_i64 {compute_order = #tla.compute_order<N_FIRST>, hf32_mode = #tla.hf32_mode<HF32_DISABLE>} : !tla.tensor<!tla.layout<!tla.shape<(16,2),(16,2)>, !tla.stride<(16,256),(1,512)>, !tla.shape<32,32>, L0Clayout>, !tla.coord<0,0>, !tla.ptr<f32, l0c, 512>>, !tla.tensor<!tla.layout<!tla.shape<(16,2),(8,4)>, !tla.stride<(8,128),(1,256)>, !tla.shape<32,32>, zN>, !tla.coord<0,0>, !tla.ptr<f32, l0a, 512>>, !tla.tensor<!tla.layout<!tla.shape<(8,4),(16,2)>, !tla.stride<(1,256),(8,128)>, !tla.shape<32,32>, nZ>, !tla.coord<0,0>, !tla.ptr<f32, l0b, 512>>, i1, i64
94	- tla.mmad %11, %7, %9, %false, %c3_i64 {compute_order = #tla.compute_order<M_FIRST>} : !tla.tensor<!tla.layout<!tla.shape<(16,2),(16,2)>, !tla.stride<(16,256),(1,512)>, !tla.shape<32,32>, L0Clayout>, !tla.coord<0,0>, !tla.ptr<f32, l0c, 512>>, !tla.tensor<!tla.layout<!tla.shape<(16,2),(8,4)>, !tla.stride<(8,128),(1,256)>, !tla.shape<32,32>, zN>, !tla.coord<0,0>, !tla.ptr<f32, l0a, 512>>, !tla.tensor<!tla.layout<!tla.shape<(8,4),(16,2)>, !tla.stride<(1,256),(8,128)>, !tla.shape<32,32>, nZ>, !tla.coord<0,0>, !tla.ptr<f32, l0b, 512>>, i1, i64	94	+ tla.mmad %11, %7, %9, %false, %c3_i64 {compute_order = #tla.compute_order<M_FIRST>, hf32_mode = #tla.hf32_mode<HF32_DISABLE>} : !tla.tensor<!tla.layout<!tla.shape<(16,2),(16,2)>, !tla.stride<(16,256),(1,512)>, !tla.shape<32,32>, L0Clayout>, !tla.coord<0,0>, !tla.ptr<f32, l0c, 512>>, !tla.tensor<!tla.layout<!tla.shape<(16,2),(8,4)>, !tla.stride<(8,128),(1,256)>, !tla.shape<32,32>, zN>, !tla.coord<0,0>, !tla.ptr<f32, l0a, 512>>, !tla.tensor<!tla.layout<!tla.shape<(8,4),(16,2)>, !tla.stride<(1,256),(8,128)>, !tla.shape<32,32>, nZ>, !tla.coord<0,0>, !tla.ptr<f32, l0b, 512>>, i1, i64
95	}	95	}
96	tla.return	96	tla.return
97	}	97	}

Mpython/tla_dsl/tests/lit/tla-compile/mmad-end_to_end.mlir+1-1

	@@ -56,7 +56,7 @@ module attributes {tla.module_exec_units = "cube"} {
56	%38 = arith.constant true	56	%38 = arith.constant true
57	%39 = arith.constant 0 : i64	57	%39 = arith.constant 0 : i64
58	"tla.cube"() ({	58	"tla.cube"() ({
59	- "tla.mmad"(%37, %35, %36, %38, %39) {compute_order = #tla.compute_order<M_FIRST>} : (!tla.tensor<!tla.layout<!tla.shape<(16,2),(16,2)>, !tla.stride<(16,256),(1,512)>, !tla.shape<32,32>, L0Clayout>, !tla.coord<0,0>, !tla.ptr<f32, l0c, 512>>, !tla.tensor<!tla.layout<!tla.shape<(16,2),(8,4)>, !tla.stride<(8,128),(1,256)>, !tla.shape<32,32>, zN>, !tla.coord<0,0>, !tla.ptr<f32, l0a, 512>>, !tla.tensor<!tla.layout<!tla.shape<(8,4),(16,2)>, !tla.stride<(1,256),(8,128)>, !tla.shape<32,32>, nZ>, !tla.coord<0,0>, !tla.ptr<f32, l0b, 512>>, i1, i64) -> ()	59	+ "tla.mmad"(%37, %35, %36, %38, %39) {compute_order = #tla.compute_order<M_FIRST>, hf32_mode = #tla.hf32_mode<HF32_DISABLE>} : (!tla.tensor<!tla.layout<!tla.shape<(16,2),(16,2)>, !tla.stride<(16,256),(1,512)>, !tla.shape<32,32>, L0Clayout>, !tla.coord<0,0>, !tla.ptr<f32, l0c, 512>>, !tla.tensor<!tla.layout<!tla.shape<(16,2),(8,4)>, !tla.stride<(8,128),(1,256)>, !tla.shape<32,32>, zN>, !tla.coord<0,0>, !tla.ptr<f32, l0a, 512>>, !tla.tensor<!tla.layout<!tla.shape<(8,4),(16,2)>, !tla.stride<(1,256),(8,128)>, !tla.shape<32,32>, nZ>, !tla.coord<0,0>, !tla.ptr<f32, l0b, 512>>, i1, i64) -> ()
60	"tla.set_flag"(%2) : (!tla.flag) -> ()	60	"tla.set_flag"(%2) : (!tla.flag) -> ()
61	"tla.wait_flag"(%2) : (!tla.flag) -> ()	61	"tla.wait_flag"(%2) : (!tla.flag) -> ()
62	}) : () -> ()	62	}) : () -> ()

Apython/tla_dsl/tests/lit/tla-compile/mmad-hf32-mode-verify.test+60-0

@@ -0,0 +1,60 @@
	1	+// Reject a non-f32 mmad source operand when hf32_mode is not HF32_DISABLE:
	2	+// HF32 rounding only applies to f32 L0A/L0B operands.
	3	+// Each case is a separate compile so coverage does not depend on MLIR verifying
	4	+// multiple functions in a single pass.
	5	+
	6	+// RUN: sed -n '/BEGIN-F16-LHS$/,/END-F16-LHS$/p' %s \| not %tla_compile - -o - 2>&1 \| %filecheck %s --check-prefix=F16-LHS
	7	+// RUN: sed -n '/BEGIN-F16-RHS$/,/END-F16-RHS$/p' %s \| not %tla_compile - -o - 2>&1 \| %filecheck %s --check-prefix=F16-RHS
	8	+
	9	+// F16-LHS: 'tla.mmad' op hf32_mode HF32_NEAREST_ZERO requires f32 source operands, but lhs operand has element type 'f16'
	10	+// F16-RHS: 'tla.mmad' op hf32_mode HF32_NEAREST_ZERO requires f32 source operands, but rhs operand has element type 'f16'
	11	+
	12	+// BEGIN-F16-LHS
	13	+module {
	14	+ tla.func @mmad_hf32_f16_lhs() {
	15	+ %0 = tla.alloc_ptr{size_bytes = 4096} -> !tla.ptr<f32, l1, 512>
	16	+ %1 = tla.make_shape -> !tla.shape<32,32>
	17	+ %2 = tla.make_stride -> !tla.stride<32,1>
	18	+ %3 = tla.make_layout %1, %2 : !tla.shape<32,32>, !tla.stride<32,1> -> !tla.layout<!tla.shape<32,32>, !tla.stride<32,1>, !tla.shape<32,32>, row_major>
	19	+ %4 = tla.make_coord -> !tla.coord<0,0>
	20	+ %5 = tla.make_tensor %0, %3, %4 : !tla.ptr<f32, l1, 512>, !tla.layout<!tla.shape<32,32>, !tla.stride<32,1>, !tla.shape<32,32>, row_major>, !tla.coord<0,0> -> !tla.tensor<!tla.layout<!tla.shape<32,32>, !tla.stride<32,1>, !tla.shape<32,32>, row_major>, !tla.coord<0,0>, !tla.ptr<f32, l1, 512>>
	21	+ %6 = tla.alloc_ptr{size_bytes = 4096} -> !tla.ptr<f16, l0a, 512>
	22	+ %7 = tla.make_tensor_like %6 like %5 layoutTag("zN") : !tla.ptr<f16, l0a, 512>, !tla.tensor<!tla.layout<!tla.shape<32,32>, !tla.stride<32,1>, !tla.shape<32,32>, row_major>, !tla.coord<0,0>, !tla.ptr<f32, l1, 512>> -> !tla.tensor<!tla.layout<!tla.shape<(16,2),(8,4)>, !tla.stride<(8,128),(1,256)>, !tla.shape<32,32>, zN>, !tla.coord<0,0>, !tla.ptr<f16, l0a, 512>>
	23	+ %8 = tla.alloc_ptr{size_bytes = 4096} -> !tla.ptr<f32, l0b, 512>
	24	+ %9 = tla.make_tensor_like %8 like %5 layoutTag("nZ") : !tla.ptr<f32, l0b, 512>, !tla.tensor<!tla.layout<!tla.shape<32,32>, !tla.stride<32,1>, !tla.shape<32,32>, row_major>, !tla.coord<0,0>, !tla.ptr<f32, l1, 512>> -> !tla.tensor<!tla.layout<!tla.shape<(8,4),(16,2)>, !tla.stride<(1,256),(8,128)>, !tla.shape<32,32>, nZ>, !tla.coord<0,0>, !tla.ptr<f32, l0b, 512>>
	25	+ %10 = tla.alloc_ptr{size_bytes = 4096} -> !tla.ptr<f32, l0c, 512>
	26	+ %11 = tla.make_tensor_like %10 like %5 layoutTag("L0Clayout") : !tla.ptr<f32, l0c, 512>, !tla.tensor<!tla.layout<!tla.shape<32,32>, !tla.stride<32,1>, !tla.shape<32,32>, row_major>, !tla.coord<0,0>, !tla.ptr<f32, l1, 512>> -> !tla.tensor<!tla.layout<!tla.shape<(16,2),(16,2)>, !tla.stride<(16,256),(1,512)>, !tla.shape<32,32>, L0Clayout>, !tla.coord<0,0>, !tla.ptr<f32, l0c, 512>>
	27	+ tla.cube {
	28	+ %true = arith.constant true
	29	+ %c3_i64 = arith.constant 3 : i64
	30	+ tla.mmad %11, %7, %9, %true, %c3_i64 {compute_order = #tla.compute_order<M_FIRST>, hf32_mode = #tla.hf32_mode<HF32_NEAREST_ZERO>} : !tla.tensor<!tla.layout<!tla.shape<(16,2),(16,2)>, !tla.stride<(16,256),(1,512)>, !tla.shape<32,32>, L0Clayout>, !tla.coord<0,0>, !tla.ptr<f32, l0c, 512>>, !tla.tensor<!tla.layout<!tla.shape<(16,2),(8,4)>, !tla.stride<(8,128),(1,256)>, !tla.shape<32,32>, zN>, !tla.coord<0,0>, !tla.ptr<f16, l0a, 512>>, !tla.tensor<!tla.layout<!tla.shape<(8,4),(16,2)>, !tla.stride<(1,256),(8,128)>, !tla.shape<32,32>, nZ>, !tla.coord<0,0>, !tla.ptr<f32, l0b, 512>>, i1, i64
	31	+ }
	32	+ tla.return
	33	+ }
	34	+}
	35	+// END-F16-LHS
	36	+
	37	+// BEGIN-F16-RHS
	38	+module {
	39	+ tla.func @mmad_hf32_f16_rhs() {
	40	+ %0 = tla.alloc_ptr{size_bytes = 4096} -> !tla.ptr<f32, l1, 512>
	41	+ %1 = tla.make_shape -> !tla.shape<32,32>
	42	+ %2 = tla.make_stride -> !tla.stride<32,1>
	43	+ %3 = tla.make_layout %1, %2 : !tla.shape<32,32>, !tla.stride<32,1> -> !tla.layout<!tla.shape<32,32>, !tla.stride<32,1>, !tla.shape<32,32>, row_major>
	44	+ %4 = tla.make_coord -> !tla.coord<0,0>
	45	+ %5 = tla.make_tensor %0, %3, %4 : !tla.ptr<f32, l1, 512>, !tla.layout<!tla.shape<32,32>, !tla.stride<32,1>, !tla.shape<32,32>, row_major>, !tla.coord<0,0> -> !tla.tensor<!tla.layout<!tla.shape<32,32>, !tla.stride<32,1>, !tla.shape<32,32>, row_major>, !tla.coord<0,0>, !tla.ptr<f32, l1, 512>>
	46	+ %6 = tla.alloc_ptr{size_bytes = 4096} -> !tla.ptr<f32, l0a, 512>
	47	+ %7 = tla.make_tensor_like %6 like %5 layoutTag("zN") : !tla.ptr<f32, l0a, 512>, !tla.tensor<!tla.layout<!tla.shape<32,32>, !tla.stride<32,1>, !tla.shape<32,32>, row_major>, !tla.coord<0,0>, !tla.ptr<f32, l1, 512>> -> !tla.tensor<!tla.layout<!tla.shape<(16,2),(8,4)>, !tla.stride<(8,128),(1,256)>, !tla.shape<32,32>, zN>, !tla.coord<0,0>, !tla.ptr<f32, l0a, 512>>
	48	+ %8 = tla.alloc_ptr{size_bytes = 4096} -> !tla.ptr<f16, l0b, 512>
	49	+ %9 = tla.make_tensor_like %8 like %5 layoutTag("nZ") : !tla.ptr<f16, l0b, 512>, !tla.tensor<!tla.layout<!tla.shape<32,32>, !tla.stride<32,1>, !tla.shape<32,32>, row_major>, !tla.coord<0,0>, !tla.ptr<f32, l1, 512>> -> !tla.tensor<!tla.layout<!tla.shape<(8,4),(16,2)>, !tla.stride<(1,256),(8,128)>, !tla.shape<32,32>, nZ>, !tla.coord<0,0>, !tla.ptr<f16, l0b, 512>>
	50	+ %10 = tla.alloc_ptr{size_bytes = 4096} -> !tla.ptr<f32, l0c, 512>
	51	+ %11 = tla.make_tensor_like %10 like %5 layoutTag("L0Clayout") : !tla.ptr<f32, l0c, 512>, !tla.tensor<!tla.layout<!tla.shape<32,32>, !tla.stride<32,1>, !tla.shape<32,32>, row_major>, !tla.coord<0,0>, !tla.ptr<f32, l1, 512>> -> !tla.tensor<!tla.layout<!tla.shape<(16,2),(16,2)>, !tla.stride<(16,256),(1,512)>, !tla.shape<32,32>, L0Clayout>, !tla.coord<0,0>, !tla.ptr<f32, l0c, 512>>
	52	+ tla.cube {
	53	+ %true = arith.constant true
	54	+ %c3_i64 = arith.constant 3 : i64
	55	+ tla.mmad %11, %7, %9, %true, %c3_i64 {compute_order = #tla.compute_order<M_FIRST>, hf32_mode = #tla.hf32_mode<HF32_NEAREST_ZERO>} : !tla.tensor<!tla.layout<!tla.shape<(16,2),(16,2)>, !tla.stride<(16,256),(1,512)>, !tla.shape<32,32>, L0Clayout>, !tla.coord<0,0>, !tla.ptr<f32, l0c, 512>>, !tla.tensor<!tla.layout<!tla.shape<(16,2),(8,4)>, !tla.stride<(8,128),(1,256)>, !tla.shape<32,32>, zN>, !tla.coord<0,0>, !tla.ptr<f32, l0a, 512>>, !tla.tensor<!tla.layout<!tla.shape<(8,4),(16,2)>, !tla.stride<(1,256),(8,128)>, !tla.shape<32,32>, nZ>, !tla.coord<0,0>, !tla.ptr<f16, l0b, 512>>, i1, i64
	56	+ }
	57	+ tla.return
	58	+ }
	59	+}
	60	+// END-F16-RHS

Apython/tla_dsl/tests/lit/tla-compile/mmad-hf32-mode.mlir+137-0

@@ -0,0 +1,137 @@
	1	+// Verify the hf32_mode attribute on tla.mmad lowers to SPR.CTRL[46] (enable)
	2	+// and SPR.CTRL[47] (rounding select), the DSL equivalents of AscendC::SetHF32Mode
	3	+// and AscendC::SetHF32TransMode.
	4	+//
	5	+// HF32_NEAREST_ZERO -> set_ctrl true at ctrl[46], set_ctrl true at ctrl[47] (func entry).
	6	+// HF32_NEAREST_EVEN -> set_ctrl true at ctrl[46], set_ctrl false at ctrl[47] (func entry).
	7	+// HF32_DISABLE -> set_ctrl false at ctrl[46], set_ctrl false at ctrl[47] (func entry).
	8	+// mixed (NEAREST_ZERO + NEAREST_EVEN) -> per-mmad set_ctrl at ctrl[46]/ctrl[47].
	9	+//
	10	+// RUN: %tla_compile %s -o - \| %filecheck %s --check-prefix=NZ
	11	+// RUN: %tla_compile %s -o - \| %filecheck %s --check-prefix=NE
	12	+// RUN: %tla_compile %s -o - \| %filecheck %s --check-prefix=DIS
	13	+// RUN: %tla_compile %s -o - \| %filecheck %s --check-prefix=MIXED
	14	+
	15	+module {
	16	+ tla.func @mmad_hf32_nearest_zero() {
	17	+ %0 = tla.alloc_ptr{size_bytes = 4096} -> !tla.ptr<f32, l1, 512>
	18	+ %1 = tla.make_shape -> !tla.shape<32,32>
	19	+ %2 = tla.make_stride -> !tla.stride<32,1>
	20	+ %3 = tla.make_layout %1, %2 : !tla.shape<32,32>, !tla.stride<32,1> -> !tla.layout<!tla.shape<32,32>, !tla.stride<32,1>, !tla.shape<32,32>, row_major>
	21	+ %4 = tla.make_coord -> !tla.coord<0,0>
	22	+ %5 = tla.make_tensor %0, %3, %4 : !tla.ptr<f32, l1, 512>, !tla.layout<!tla.shape<32,32>, !tla.stride<32,1>, !tla.shape<32,32>, row_major>, !tla.coord<0,0> -> !tla.tensor<!tla.layout<!tla.shape<32,32>, !tla.stride<32,1>, !tla.shape<32,32>, row_major>, !tla.coord<0,0>, !tla.ptr<f32, l1, 512>>
	23	+ %6 = tla.alloc_ptr{size_bytes = 4096} -> !tla.ptr<f32, l0a, 512>
	24	+ %7 = tla.make_tensor_like %6 like %5 layoutTag("zN") : !tla.ptr<f32, l0a, 512>, !tla.tensor<!tla.layout<!tla.shape<32,32>, !tla.stride<32,1>, !tla.shape<32,32>, row_major>, !tla.coord<0,0>, !tla.ptr<f32, l1, 512>> -> !tla.tensor<!tla.layout<!tla.shape<(16,2),(8,4)>, !tla.stride<(8,128),(1,256)>, !tla.shape<32,32>, zN>, !tla.coord<0,0>, !tla.ptr<f32, l0a, 512>>
	25	+ %8 = tla.alloc_ptr{size_bytes = 4096} -> !tla.ptr<f32, l0b, 512>
	26	+ %9 = tla.make_tensor_like %8 like %5 layoutTag("nZ") : !tla.ptr<f32, l0b, 512>, !tla.tensor<!tla.layout<!tla.shape<32,32>, !tla.stride<32,1>, !tla.shape<32,32>, row_major>, !tla.coord<0,0>, !tla.ptr<f32, l1, 512>> -> !tla.tensor<!tla.layout<!tla.shape<(8,4),(16,2)>, !tla.stride<(1,256),(8,128)>, !tla.shape<32,32>, nZ>, !tla.coord<0,0>, !tla.ptr<f32, l0b, 512>>
	27	+ %10 = tla.alloc_ptr{size_bytes = 4096} -> !tla.ptr<f32, l0c, 512>
	28	+ %11 = tla.make_tensor_like %10 like %5 layoutTag("L0Clayout") : !tla.ptr<f32, l0c, 512>, !tla.tensor<!tla.layout<!tla.shape<32,32>, !tla.stride<32,1>, !tla.shape<32,32>, row_major>, !tla.coord<0,0>, !tla.ptr<f32, l1, 512>> -> !tla.tensor<!tla.layout<!tla.shape<(16,2),(16,2)>, !tla.stride<(16,256),(1,512)>, !tla.shape<32,32>, L0Clayout>, !tla.coord<0,0>, !tla.ptr<f32, l0c, 512>>
	29	+ tla.cube {
	30	+ %true = arith.constant true
	31	+ %c3_i64 = arith.constant 3 : i64
	32	+ tla.mmad %11, %7, %9, %true, %c3_i64 {compute_order = #tla.compute_order<M_FIRST>, hf32_mode = #tla.hf32_mode<HF32_NEAREST_ZERO>} : !tla.tensor<!tla.layout<!tla.shape<(16,2),(16,2)>, !tla.stride<(16,256),(1,512)>, !tla.shape<32,32>, L0Clayout>, !tla.coord<0,0>, !tla.ptr<f32, l0c, 512>>, !tla.tensor<!tla.layout<!tla.shape<(16,2),(8,4)>, !tla.stride<(8,128),(1,256)>, !tla.shape<32,32>, zN>, !tla.coord<0,0>, !tla.ptr<f32, l0a, 512>>, !tla.tensor<!tla.layout<!tla.shape<(8,4),(16,2)>, !tla.stride<(1,256),(8,128)>, !tla.shape<32,32>, nZ>, !tla.coord<0,0>, !tla.ptr<f32, l0b, 512>>, i1, i64
	33	+ }
	34	+ tla.return
	35	+ }
	36	+
	37	+ tla.func @mmad_hf32_nearest_even() {
	38	+ %0 = tla.alloc_ptr{size_bytes = 4096} -> !tla.ptr<f32, l1, 512>
	39	+ %1 = tla.make_shape -> !tla.shape<32,32>
	40	+ %2 = tla.make_stride -> !tla.stride<32,1>
	41	+ %3 = tla.make_layout %1, %2 : !tla.shape<32,32>, !tla.stride<32,1> -> !tla.layout<!tla.shape<32,32>, !tla.stride<32,1>, !tla.shape<32,32>, row_major>
	42	+ %4 = tla.make_coord -> !tla.coord<0,0>
	43	+ %5 = tla.make_tensor %0, %3, %4 : !tla.ptr<f32, l1, 512>, !tla.layout<!tla.shape<32,32>, !tla.stride<32,1>, !tla.shape<32,32>, row_major>, !tla.coord<0,0> -> !tla.tensor<!tla.layout<!tla.shape<32,32>, !tla.stride<32,1>, !tla.shape<32,32>, row_major>, !tla.coord<0,0>, !tla.ptr<f32, l1, 512>>
	44	+ %6 = tla.alloc_ptr{size_bytes = 4096} -> !tla.ptr<f32, l0a, 512>
	45	+ %7 = tla.make_tensor_like %6 like %5 layoutTag("zN") : !tla.ptr<f32, l0a, 512>, !tla.tensor<!tla.layout<!tla.shape<32,32>, !tla.stride<32,1>, !tla.shape<32,32>, row_major>, !tla.coord<0,0>, !tla.ptr<f32, l1, 512>> -> !tla.tensor<!tla.layout<!tla.shape<(16,2),(8,4)>, !tla.stride<(8,128),(1,256)>, !tla.shape<32,32>, zN>, !tla.coord<0,0>, !tla.ptr<f32, l0a, 512>>
	46	+ %8 = tla.alloc_ptr{size_bytes = 4096} -> !tla.ptr<f32, l0b, 512>
	47	+ %9 = tla.make_tensor_like %8 like %5 layoutTag("nZ") : !tla.ptr<f32, l0b, 512>, !tla.tensor<!tla.layout<!tla.shape<32,32>, !tla.stride<32,1>, !tla.shape<32,32>, row_major>, !tla.coord<0,0>, !tla.ptr<f32, l1, 512>> -> !tla.tensor<!tla.layout<!tla.shape<(8,4),(16,2)>, !tla.stride<(1,256),(8,128)>, !tla.shape<32,32>, nZ>, !tla.coord<0,0>, !tla.ptr<f32, l0b, 512>>
	48	+ %10 = tla.alloc_ptr{size_bytes = 4096} -> !tla.ptr<f32, l0c, 512>
	49	+ %11 = tla.make_tensor_like %10 like %5 layoutTag("L0Clayout") : !tla.ptr<f32, l0c, 512>, !tla.tensor<!tla.layout<!tla.shape<32,32>, !tla.stride<32,1>, !tla.shape<32,32>, row_major>, !tla.coord<0,0>, !tla.ptr<f32, l1, 512>> -> !tla.tensor<!tla.layout<!tla.shape<(16,2),(16,2)>, !tla.stride<(16,256),(1,512)>, !tla.shape<32,32>, L0Clayout>, !tla.coord<0,0>, !tla.ptr<f32, l0c, 512>>
	50	+ tla.cube {
	51	+ %true = arith.constant true
	52	+ %c3_i64 = arith.constant 3 : i64
	53	+ tla.mmad %11, %7, %9, %true, %c3_i64 {compute_order = #tla.compute_order<M_FIRST>, hf32_mode = #tla.hf32_mode<HF32_NEAREST_EVEN>} : !tla.tensor<!tla.layout<!tla.shape<(16,2),(16,2)>, !tla.stride<(16,256),(1,512)>, !tla.shape<32,32>, L0Clayout>, !tla.coord<0,0>, !tla.ptr<f32, l0c, 512>>, !tla.tensor<!tla.layout<!tla.shape<(16,2),(8,4)>, !tla.stride<(8,128),(1,256)>, !tla.shape<32,32>, zN>, !tla.coord<0,0>, !tla.ptr<f32, l0a, 512>>, !tla.tensor<!tla.layout<!tla.shape<(8,4),(16,2)>, !tla.stride<(1,256),(8,128)>, !tla.shape<32,32>, nZ>, !tla.coord<0,0>, !tla.ptr<f32, l0b, 512>>, i1, i64
	54	+ }
	55	+ tla.return
	56	+ }
	57	+
	58	+ tla.func @mmad_hf32_disable() {
	59	+ %0 = tla.alloc_ptr{size_bytes = 4096} -> !tla.ptr<f32, l1, 512>
	60	+ %1 = tla.make_shape -> !tla.shape<32,32>
	61	+ %2 = tla.make_stride -> !tla.stride<32,1>
	62	+ %3 = tla.make_layout %1, %2 : !tla.shape<32,32>, !tla.stride<32,1> -> !tla.layout<!tla.shape<32,32>, !tla.stride<32,1>, !tla.shape<32,32>, row_major>
	63	+ %4 = tla.make_coord -> !tla.coord<0,0>
	64	+ %5 = tla.make_tensor %0, %3, %4 : !tla.ptr<f32, l1, 512>, !tla.layout<!tla.shape<32,32>, !tla.stride<32,1>, !tla.shape<32,32>, row_major>, !tla.coord<0,0> -> !tla.tensor<!tla.layout<!tla.shape<32,32>, !tla.stride<32,1>, !tla.shape<32,32>, row_major>, !tla.coord<0,0>, !tla.ptr<f32, l1, 512>>
	65	+ %6 = tla.alloc_ptr{size_bytes = 4096} -> !tla.ptr<f32, l0a, 512>
	66	+ %7 = tla.make_tensor_like %6 like %5 layoutTag("zN") : !tla.ptr<f32, l0a, 512>, !tla.tensor<!tla.layout<!tla.shape<32,32>, !tla.stride<32,1>, !tla.shape<32,32>, row_major>, !tla.coord<0,0>, !tla.ptr<f32, l1, 512>> -> !tla.tensor<!tla.layout<!tla.shape<(16,2),(8,4)>, !tla.stride<(8,128),(1,256)>, !tla.shape<32,32>, zN>, !tla.coord<0,0>, !tla.ptr<f32, l0a, 512>>
	67	+ %8 = tla.alloc_ptr{size_bytes = 4096} -> !tla.ptr<f32, l0b, 512>
	68	+ %9 = tla.make_tensor_like %8 like %5 layoutTag("nZ") : !tla.ptr<f32, l0b, 512>, !tla.tensor<!tla.layout<!tla.shape<32,32>, !tla.stride<32,1>, !tla.shape<32,32>, row_major>, !tla.coord<0,0>, !tla.ptr<f32, l1, 512>> -> !tla.tensor<!tla.layout<!tla.shape<(8,4),(16,2)>, !tla.stride<(1,256),(8,128)>, !tla.shape<32,32>, nZ>, !tla.coord<0,0>, !tla.ptr<f32, l0b, 512>>
	69	+ %10 = tla.alloc_ptr{size_bytes = 4096} -> !tla.ptr<f32, l0c, 512>
	70	+ %11 = tla.make_tensor_like %10 like %5 layoutTag("L0Clayout") : !tla.ptr<f32, l0c, 512>, !tla.tensor<!tla.layout<!tla.shape<32,32>, !tla.stride<32,1>, !tla.shape<32,32>, row_major>, !tla.coord<0,0>, !tla.ptr<f32, l1, 512>> -> !tla.tensor<!tla.layout<!tla.shape<(16,2),(16,2)>, !tla.stride<(16,256),(1,512)>, !tla.shape<32,32>, L0Clayout>, !tla.coord<0,0>, !tla.ptr<f32, l0c, 512>>
	71	+ tla.cube {
	72	+ %true = arith.constant true
	73	+ %c3_i64 = arith.constant 3 : i64
	74	+ tla.mmad %11, %7, %9, %true, %c3_i64 {compute_order = #tla.compute_order<M_FIRST>, hf32_mode = #tla.hf32_mode<HF32_DISABLE>} : !tla.tensor<!tla.layout<!tla.shape<(16,2),(16,2)>, !tla.stride<(16,256),(1,512)>, !tla.shape<32,32>, L0Clayout>, !tla.coord<0,0>, !tla.ptr<f32, l0c, 512>>, !tla.tensor<!tla.layout<!tla.shape<(16,2),(8,4)>, !tla.stride<(8,128),(1,256)>, !tla.shape<32,32>, zN>, !tla.coord<0,0>, !tla.ptr<f32, l0a, 512>>, !tla.tensor<!tla.layout<!tla.shape<(8,4),(16,2)>, !tla.stride<(1,256),(8,128)>, !tla.shape<32,32>, nZ>, !tla.coord<0,0>, !tla.ptr<f32, l0b, 512>>, i1, i64
	75	+ }
	76	+ tla.return
	77	+ }
	78	+
	79	+ tla.func @mmad_hf32_mixed() {
	80	+ %0 = tla.alloc_ptr{size_bytes = 4096} -> !tla.ptr<f32, l1, 512>
	81	+ %1 = tla.make_shape -> !tla.shape<32,32>
	82	+ %2 = tla.make_stride -> !tla.stride<32,1>
	83	+ %3 = tla.make_layout %1, %2 : !tla.shape<32,32>, !tla.stride<32,1> -> !tla.layout<!tla.shape<32,32>, !tla.stride<32,1>, !tla.shape<32,32>, row_major>
	84	+ %4 = tla.make_coord -> !tla.coord<0,0>
	85	+ %5 = tla.make_tensor %0, %3, %4 : !tla.ptr<f32, l1, 512>, !tla.layout<!tla.shape<32,32>, !tla.stride<32,1>, !tla.shape<32,32>, row_major>, !tla.coord<0,0> -> !tla.tensor<!tla.layout<!tla.shape<32,32>, !tla.stride<32,1>, !tla.shape<32,32>, row_major>, !tla.coord<0,0>, !tla.ptr<f32, l1, 512>>
	86	+ %6 = tla.alloc_ptr{size_bytes = 4096} -> !tla.ptr<f32, l0a, 512>
	87	+ %7 = tla.make_tensor_like %6 like %5 layoutTag("zN") : !tla.ptr<f32, l0a, 512>, !tla.tensor<!tla.layout<!tla.shape<32,32>, !tla.stride<32,1>, !tla.shape<32,32>, row_major>, !tla.coord<0,0>, !tla.ptr<f32, l1, 512>> -> !tla.tensor<!tla.layout<!tla.shape<(16,2),(8,4)>, !tla.stride<(8,128),(1,256)>, !tla.shape<32,32>, zN>, !tla.coord<0,0>, !tla.ptr<f32, l0a, 512>>
	88	+ %8 = tla.alloc_ptr{size_bytes = 4096} -> !tla.ptr<f32, l0b, 512>
	89	+ %9 = tla.make_tensor_like %8 like %5 layoutTag("nZ") : !tla.ptr<f32, l0b, 512>, !tla.tensor<!tla.layout<!tla.shape<32,32>, !tla.stride<32,1>, !tla.shape<32,32>, row_major>, !tla.coord<0,0>, !tla.ptr<f32, l1, 512>> -> !tla.tensor<!tla.layout<!tla.shape<(8,4),(16,2)>, !tla.stride<(1,256),(8,128)>, !tla.shape<32,32>, nZ>, !tla.coord<0,0>, !tla.ptr<f32, l0b, 512>>
	90	+ %10 = tla.alloc_ptr{size_bytes = 4096} -> !tla.ptr<f32, l0c, 512>
	91	+ %11 = tla.make_tensor_like %10 like %5 layoutTag("L0Clayout") : !tla.ptr<f32, l0c, 512>, !tla.tensor<!tla.layout<!tla.shape<32,32>, !tla.stride<32,1>, !tla.shape<32,32>, row_major>, !tla.coord<0,0>, !tla.ptr<f32, l1, 512>> -> !tla.tensor<!tla.layout<!tla.shape<(16,2),(16,2)>, !tla.stride<(16,256),(1,512)>, !tla.shape<32,32>, L0Clayout>, !tla.coord<0,0>, !tla.ptr<f32, l0c, 512>>
	92	+ tla.cube {
	93	+ %true = arith.constant true
	94	+ %false = arith.constant false
	95	+ %c3_i64 = arith.constant 3 : i64
	96	+ tla.mmad %11, %7, %9, %true, %c3_i64 {compute_order = #tla.compute_order<M_FIRST>, hf32_mode = #tla.hf32_mode<HF32_NEAREST_ZERO>} : !tla.tensor<!tla.layout<!tla.shape<(16,2),(16,2)>, !tla.stride<(16,256),(1,512)>, !tla.shape<32,32>, L0Clayout>, !tla.coord<0,0>, !tla.ptr<f32, l0c, 512>>, !tla.tensor<!tla.layout<!tla.shape<(16,2),(8,4)>, !tla.stride<(8,128),(1,256)>, !tla.shape<32,32>, zN>, !tla.coord<0,0>, !tla.ptr<f32, l0a, 512>>, !tla.tensor<!tla.layout<!tla.shape<(8,4),(16,2)>, !tla.stride<(1,256),(8,128)>, !tla.shape<32,32>, nZ>, !tla.coord<0,0>, !tla.ptr<f32, l0b, 512>>, i1, i64
	97	+ tla.mmad %11, %7, %9, %false, %c3_i64 {compute_order = #tla.compute_order<M_FIRST>, hf32_mode = #tla.hf32_mode<HF32_NEAREST_EVEN>} : !tla.tensor<!tla.layout<!tla.shape<(16,2),(16,2)>, !tla.stride<(16,256),(1,512)>, !tla.shape<32,32>, L0Clayout>, !tla.coord<0,0>, !tla.ptr<f32, l0c, 512>>, !tla.tensor<!tla.layout<!tla.shape<(16,2),(8,4)>, !tla.stride<(8,128),(1,256)>, !tla.shape<32,32>, zN>, !tla.coord<0,0>, !tla.ptr<f32, l0a, 512>>, !tla.tensor<!tla.layout<!tla.shape<(8,4),(16,2)>, !tla.stride<(1,256),(8,128)>, !tla.shape<32,32>, nZ>, !tla.coord<0,0>, !tla.ptr<f32, l0b, 512>>, i1, i64
	98	+ }
	99	+ tla.return
	100	+ }
	101	+
	102	+}
	103	+
	104	+// NZ-LABEL: func.func @mmad_hf32_nearest_zero
	105	+// NZ: hivm.hir.set_ctrl true at ctrl[46]
	106	+// NZ: hivm.hir.set_ctrl true at ctrl[47]
	107	+// NZ: call @mmad_float_float_float
	108	+// NZ: hivm.hir.set_ctrl false at ctrl[46]
	109	+// NZ: hivm.hir.set_ctrl false at ctrl[47]
	110	+// NZ-NOT: tla.mmad
	111	+
	112	+// NE-LABEL: func.func @mmad_hf32_nearest_even
	113	+// NE: hivm.hir.set_ctrl true at ctrl[46]
	114	+// NE: hivm.hir.set_ctrl false at ctrl[47]
	115	+// NE: call @mmad_float_float_float
	116	+// NE: hivm.hir.set_ctrl false at ctrl[46]
	117	+// NE: hivm.hir.set_ctrl false at ctrl[47]
	118	+// NE-NOT: tla.mmad
	119	+
	120	+// DIS-LABEL: func.func @mmad_hf32_disable
	121	+// DIS: hivm.hir.set_ctrl false at ctrl[46]
	122	+// DIS: hivm.hir.set_ctrl false at ctrl[47]
	123	+// DIS: call @mmad_float_float_float
	124	+// DIS: hivm.hir.set_ctrl false at ctrl[46]
	125	+// DIS: hivm.hir.set_ctrl false at ctrl[47]
	126	+// DIS-NOT: tla.mmad
	127	+
	128	+// MIXED-LABEL: func.func @mmad_hf32_mixed
	129	+// MIXED: hivm.hir.set_ctrl true at ctrl[46]
	130	+// MIXED: hivm.hir.set_ctrl true at ctrl[47]
	131	+// MIXED: call @mmad_float_float_float
	132	+// MIXED: hivm.hir.set_ctrl true at ctrl[46]
	133	+// MIXED: hivm.hir.set_ctrl false at ctrl[47]
	134	+// MIXED: call @mmad_float_float_float
	135	+// MIXED: hivm.hir.set_ctrl false at ctrl[46]
	136	+// MIXED: hivm.hir.set_ctrl false at ctrl[47]
	137	+// MIXED-NOT: tla.mmad

Mpython/tla_dsl/tests/test_frontend_lowering.py+72-0

	@@ -200,6 +200,54 @@ def mmad_bad_compute_order_kernel(
200	with tla.cube():	200	with tla.cube():
201	tla.mmad(acc, lhs, rhs, init_c=True, compute_order="N_FIRST")	201	tla.mmad(acc, lhs, rhs, init_c=True, compute_order="N_FIRST")
202		202
		203	+@tla.kernel
		204	+def mmad_hf32_mode_kernel(
		205	+ lhs: tla.Tensor, rhs: tla.Tensor, acc: tla.Tensor
		206	+) -> None:
		207	+ with tla.cube():
		208	+ tla.mmad(
		209	+ acc,
		210	+ lhs,
		211	+ rhs,
		212	+ init_c=True,
		213	+ hf32_mode=tla.params.HF32Mode.HF32_NEAREST_ZERO,
		214	+ )
		215	+
		216	+@tla.kernel
		217	+def mmad_hf32_mode_nearest_even_kernel(
		218	+ lhs: tla.Tensor, rhs: tla.Tensor, acc: tla.Tensor
		219	+) -> None:
		220	+ with tla.cube():
		221	+ tla.mmad(
		222	+ acc,
		223	+ lhs,
		224	+ rhs,
		225	+ init_c=True,
		226	+ hf32_mode=tla.params.HF32Mode.HF32_NEAREST_EVEN,
		227	+ )
		228	+
		229	+@tla.kernel
		230	+def mmad_default_hf32_mode_kernel(
		231	+ lhs: tla.Tensor, rhs: tla.Tensor, acc: tla.Tensor
		232	+) -> None:
		233	+ with tla.cube():
		234	+ tla.mmad(acc, lhs, rhs, init_c=True)
		235	+
		236	+@tla.kernel
		237	+def mmad_default_hf32_mode_kernel_v2(
		238	+ lhs: tla.Tensor, rhs: tla.Tensor, acc: tla.Tensor
		239	+) -> None:
		240	+ with tla.cube():
		241	+ tla.mmad(acc, lhs, rhs, init_c=True, hf32_mode=tla.params.HF32Mode.HF32_DISABLE)
		242	+
		243	+
		244	+@tla.kernel
		245	+def mmad_bad_hf32_mode_kernel(
		246	+ lhs: tla.Tensor, rhs: tla.Tensor, acc: tla.Tensor
		247	+) -> None:
		248	+ with tla.cube():
		249	+ tla.mmad(acc, lhs, rhs, init_c=True, hf32_mode="nearest_zero")
		250	+
203	@tla.kernel	251	@tla.kernel
204	def make_shape_index_arg_ok(dim: "index") -> None:	252	def make_shape_index_arg_ok(dim: "index") -> None:
205	tla.make_shape(dim, 16)	253	tla.make_shape(dim, 16)
	@@ -457,6 +505,30 @@ def test_mmad_without_region_lowers() -> None:
457	raise	505	raise
458	assert "tla.mmad" in mlir	506	assert "tla.mmad" in mlir
459		507
		508	+def test_hf32_mode_emits_attr() -> None:
		509	+ ta, tb, tc = _mmad_tensor_args()
		510	+ mlir = mmad_hf32_mode_kernel.dump_mlir(type_args=(ta, tb, tc))
		511	+ assert "tla.mmad" in mlir
		512	+ assert "hf32_mode = #tla.hf32_mode<HF32_NEAREST_ZERO>" in mlir
		513	+
		514	+def test_hf32_mode_nearest_even_emits_attr() -> None:
		515	+ ta, tb, tc = _mmad_tensor_args()
		516	+ mlir = mmad_hf32_mode_nearest_even_kernel.dump_mlir(type_args=(ta, tb, tc))
		517	+ assert "tla.mmad" in mlir
		518	+ assert "hf32_mode = #tla.hf32_mode<HF32_NEAREST_EVEN>" in mlir
		519	+
		520	+def test_default_hf32_mode_is_disable() -> None:
		521	+ ta, tb, tc = _mmad_tensor_args()
		522	+ for _kernel in (mmad_default_hf32_mode_kernel, mmad_default_hf32_mode_kernel_v2):
		523	+ mlir = _kernel.dump_mlir(type_args=(ta, tb, tc))
		524	+ assert "tla.mmad" in mlir
		525	+ assert "hf32_mode = #tla.hf32_mode<HF32_DISABLE>" in mlir
		526	+
		527	+def test_rejects_invalid_hf32_mode() -> None:
		528	+ ta, tb, tc = _mmad_tensor_args()
		529	+ with pytest.raises(TlaLoweringError):
		530	+ _ = mmad_bad_hf32_mode_kernel.dump_mlir(type_args=(ta, tb, tc))
		531	+
460	def test_make_shape_accepts_index_typed_components() -> None:	532	def test_make_shape_accepts_index_typed_components() -> None:
461	"""Lowering emits tla.make_shape with a dynamic dim when the component is a kernel parameter."""	533	"""Lowering emits tla.make_shape with a dynamic dim when the component is a kernel parameter."""
462	mlir = make_shape_index_arg_ok.dump_mlir(type_args=(4,))	534	mlir = make_shape_index_arg_ok.dump_mlir(type_args=(4,))