Diff 346153

mlir/include/mlir/Dialect/Linalg/IR/LinalgNamedStructuredOps.yaml

--- !LinalgOpConfig		--- !LinalgOpConfig
metadata: !LinalgOpMetadata		metadata: !LinalgOpMetadata
name: matmul		name: matmul
cpp_op_name: MatmulOp		cpp_class_name: MatmulOp
doc: \|-		doc: \|-
Performs a matrix multiplication of two 2D inputs.		Performs a matrix multiplication of two 2D inputs.

Numeric casting is performed on the operands to the inner multiply, promoting		Numeric casting is performed on the operands to the inner multiply, promoting
them to the same data type as the accumulator/output.		them to the same data type as the accumulator/output.
implements:		implements:
- LinalgContractionOpInterface		- LinalgContractionOpInterface
structured_op: !LinalgStructuredOpConfig		structured_op: !LinalgStructuredOpConfig
▲ Show 20 Lines • Show All 45 Lines • ▼ Show 20 Lines	value: !ScalarExpression
symbolic_cast:		symbolic_cast:
type_var: U		type_var: U
operands:		operands:
- !ScalarExpression		- !ScalarExpression
scalar_arg: B		scalar_arg: B
--- !LinalgOpConfig		--- !LinalgOpConfig
metadata: !LinalgOpMetadata		metadata: !LinalgOpMetadata
name: batch_matmul		name: batch_matmul
cpp_op_name: BatchMatmulOp		cpp_class_name: BatchMatmulOp
doc: \|-		doc: \|-
Performs a batched matrix multiplication of two 3D inputs.		Performs a batched matrix multiplication of two 3D inputs.

Numeric casting is performed on the operands to the inner multiply, promoting		Numeric casting is performed on the operands to the inner multiply, promoting
them to the same data type as the accumulator/output.		them to the same data type as the accumulator/output.
implements:		implements:
- LinalgContractionOpInterface		- LinalgContractionOpInterface
structured_op: !LinalgStructuredOpConfig		structured_op: !LinalgStructuredOpConfig
▲ Show 20 Lines • Show All 46 Lines • ▼ Show 20 Lines	value: !ScalarExpression
symbolic_cast:		symbolic_cast:
type_var: U		type_var: U
operands:		operands:
- !ScalarExpression		- !ScalarExpression
scalar_arg: B		scalar_arg: B
--- !LinalgOpConfig		--- !LinalgOpConfig
metadata: !LinalgOpMetadata		metadata: !LinalgOpMetadata
name: matvec		name: matvec
cpp_op_name: MatvecOp		cpp_class_name: MatvecOp
doc: \|-		doc: \|-
Performs a matrix-vector multiplication.		Performs a matrix-vector multiplication.

Numeric casting is performed on the operands to the inner multiply, promoting		Numeric casting is performed on the operands to the inner multiply, promoting
them to the same data type as the accumulator/output.		them to the same data type as the accumulator/output.
implements:		implements:
- LinalgContractionOpInterface		- LinalgContractionOpInterface
structured_op: !LinalgStructuredOpConfig		structured_op: !LinalgStructuredOpConfig
▲ Show 20 Lines • Show All 44 Lines • ▼ Show 20 Lines	value: !ScalarExpression
symbolic_cast:		symbolic_cast:
type_var: U		type_var: U
operands:		operands:
- !ScalarExpression		- !ScalarExpression
scalar_arg: y		scalar_arg: y
--- !LinalgOpConfig		--- !LinalgOpConfig
metadata: !LinalgOpMetadata		metadata: !LinalgOpMetadata
name: vecmat		name: vecmat
cpp_op_name: VecmatOp		cpp_class_name: VecmatOp
doc: \|-		doc: \|-
Performs a vector-matrix multiplication.		Performs a vector-matrix multiplication.

Numeric casting is performed on the operands to the inner multiply, promoting		Numeric casting is performed on the operands to the inner multiply, promoting
them to the same data type as the accumulator/output.		them to the same data type as the accumulator/output.
implements:		implements:
- LinalgContractionOpInterface		- LinalgContractionOpInterface
structured_op: !LinalgStructuredOpConfig		structured_op: !LinalgStructuredOpConfig
▲ Show 20 Lines • Show All 44 Lines • ▼ Show 20 Lines	value: !ScalarExpression
symbolic_cast:		symbolic_cast:
type_var: U		type_var: U
operands:		operands:
- !ScalarExpression		- !ScalarExpression
scalar_arg: A		scalar_arg: A
--- !LinalgOpConfig		--- !LinalgOpConfig
metadata: !LinalgOpMetadata		metadata: !LinalgOpMetadata
name: dot		name: dot
cpp_op_name: DotOp		cpp_class_name: DotOp
doc: \|-		doc: \|-
Performs a dot product of two vectors to a scalar result.		Performs a dot product of two vectors to a scalar result.

Numeric casting is performed on the operands to the inner multiply, promoting		Numeric casting is performed on the operands to the inner multiply, promoting
them to the same data type as the accumulator/output.		them to the same data type as the accumulator/output.
implements:		implements:
- LinalgContractionOpInterface		- LinalgContractionOpInterface
structured_op: !LinalgStructuredOpConfig		structured_op: !LinalgStructuredOpConfig
Show All 40 Lines	value: !ScalarExpression
- !ScalarExpression		- !ScalarExpression
scalar_arg: A		scalar_arg: A
- !ScalarExpression		- !ScalarExpression
symbolic_cast:		symbolic_cast:
type_var: U		type_var: U
operands:		operands:
- !ScalarExpression		- !ScalarExpression
scalar_arg: B		scalar_arg: B
		--- !LinalgOpConfig
		metadata: !LinalgOpMetadata
		name: fill_rng_2d
		cpp_class_name: FillRng2DOp
		doc: \|-
		Fills the output tensor with pseudo random numbers.

		The operation generations pseudo random numbers using a linear congruential
		generator. It provides no guarantees regarding the distribution of the
		generated random numbers. Instead of generating the random numbers
		sequentially, it instantiates one random number generator per data element
		and runs them in parallel. The seed operand and the indices of the data
		element seed the random number generation. The min and max operands limit
		the range of the generated random numbers.

		Note: The captures are hard-coded till there is capture support on the C++
		side.
		structured_op: !LinalgStructuredOpConfig
		args:
		- !<LinalgTensorDef>
		name: O
		usage: output
		shape: affine_map<()[s0, s1] -> (s0, s1)>
		element_type_var: T
		indexing_maps: !LinalgIndexingMapsConfig
		static_indexing_maps:
		- affine_map<(d0, d1)[s0, s1] -> (d0, d1)>
		iterator_types:
		- parallel
		- parallel
		assignments:
		- !ScalarAssign
		arg: O
		value: !ScalarExpression
		symbolic_cast:
		type_var: T
		operands:
		- !ScalarExpression
		scalar_apply:
		fn_name: add
		operands:
		- !ScalarExpression
		scalar_apply:
		fn_name: mul
		operands:
		- !ScalarExpression
		scalar_apply:
		fn_name: add
		operands:
		- !ScalarExpression
		symbolic_cast:
		type_var: F64
		operands:
		- !ScalarExpression
		scalar_const: '2147483647 : i64'
		- !ScalarExpression
		symbolic_cast:
		type_var: F64
		operands:
		- !ScalarExpression
		scalar_apply:
		fn_name: add
		operands:
		- !ScalarExpression
		scalar_apply:
		fn_name: mul
		operands:
		- !ScalarExpression
		scalar_apply:
		fn_name: add
		operands:
		- !ScalarExpression
		symbolic_cast:
		type_var: I32
		operands:
		- !ScalarExpression
		scalar_index: 1
		- !ScalarExpression
		scalar_apply:
		fn_name: add
		operands:
		- !ScalarExpression
		scalar_apply:
		fn_name: mul
		operands:
		- !ScalarExpression
		scalar_apply:
		fn_name: add
		operands:
		- !ScalarExpression
		symbolic_cast:
		type_var: I32
		operands:
		- !ScalarExpression
		scalar_index: 0
		- !ScalarExpression
		symbolic_cast:
		type_var: I32
		operands:
		- !ScalarExpression
		scalar_const: '42 : i64'
		- !ScalarExpression
		symbolic_cast:
		type_var: I32
		operands:
		- !ScalarExpression
		scalar_const: '1103515245 : i64'
		- !ScalarExpression
		symbolic_cast:
		type_var: I32
		operands:
		- !ScalarExpression
		scalar_const: '12345 : i64'
		- !ScalarExpression
		symbolic_cast:
		type_var: I32
		operands:
		- !ScalarExpression
		scalar_const: '1103515245 : i64'
		- !ScalarExpression
		symbolic_cast:
		type_var: I32
		operands:
		- !ScalarExpression
		scalar_const: '12345 : i64'
		- !ScalarExpression
		scalar_apply:
		fn_name: mul
		operands:
		- !ScalarExpression
		scalar_apply:
		fn_name: sub
		operands:
		- !ScalarExpression
		symbolic_cast:
		type_var: F64
		operands:
		- !ScalarExpression
		scalar_const: '1000 : i64'
		- !ScalarExpression
		symbolic_cast:
		type_var: F64
		operands:
		- !ScalarExpression
		scalar_const: '-1000 : i64'
		- !ScalarExpression
		symbolic_cast:
		type_var: F64
		operands:
		- !ScalarExpression
		scalar_const: '2.3283063999999999E-10 : f64'
		- !ScalarExpression
		symbolic_cast:
		type_var: F64
		operands:
		- !ScalarExpression
		scalar_const: '-1000 : i64'
		nicolasvasilacheUnsubmitted Done Reply Inline Actions nit: NL nicolasvasilache: nit: NL

mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp

Show First 20 Lines • Show All 214 Lines • ▼ Show 20 Lines
// language to make such cases explicit. In the mean-time, violating this will		// language to make such cases explicit. In the mean-time, violating this will
// fail verification, which is deemed acceptable.		// fail verification, which is deemed acceptable.
//===----------------------------------------------------------------------===//		//===----------------------------------------------------------------------===//

namespace {		namespace {

class RegionBuilderHelper {		class RegionBuilderHelper {
public:		public:
RegionBuilderHelper(Block &block) : block(block) {}		RegionBuilderHelper(MLIRContext *context, Block &block)
		: context(context), block(block) {}

// Generates operations to cast the given operand to a specified type.		// Generates operations to cast the given operand to a specified type.
// If the cast cannot be performed, a warning will be issued and the		// If the cast cannot be performed, a warning will be issued and the
// operand returned as-is (which will presumably yield a verification		// operand returned as-is (which will presumably yield a verification
// issue downstream).		// issue downstream).
Value cast(Type toType, Value operand) {		Value cast(Type toType, Value operand) {
OpBuilder builder = getBuilder(operand);		OpBuilder builder = getBuilder();
auto loc = operand.getLoc();		auto loc = operand.getLoc();

if (operand.getType() == toType)		if (operand.getType() == toType)
return operand;		return operand;
if (auto toIntType = toType.dyn_cast<IntegerType>()) {		if (auto toIntType = toType.dyn_cast<IntegerType>()) {
// If operand is floating point, cast directly to the int type.		// If operand is floating point, cast directly to the int type.
if (operand.getType().isa<FloatType>())		if (operand.getType().isa<FloatType>())
return builder.create<FPToSIOp>(loc, toType, operand);		return builder.create<FPToSIOp>(loc, toType, operand);
		// Cast index operands directly to the int type.
		if (operand.getType().isIndex())
		return builder.create<IndexCastOp>(loc, toType, operand);
if (auto fromIntType = operand.getType().dyn_cast<IntegerType>()) {		if (auto fromIntType = operand.getType().dyn_cast<IntegerType>()) {
// Either sign extend or truncate.		// Either sign extend or truncate.
if (toIntType.getWidth() > fromIntType.getWidth())		if (toIntType.getWidth() > fromIntType.getWidth())
return builder.create<SignExtendIOp>(loc, toType, operand);		return builder.create<SignExtendIOp>(loc, toType, operand);
else if (toIntType.getWidth() < fromIntType.getWidth())		if (toIntType.getWidth() < fromIntType.getWidth())
return builder.create<TruncateIOp>(loc, toType, operand);		return builder.create<TruncateIOp>(loc, toType, operand);
}		}
} else if (auto toFloatType = toType.dyn_cast<FloatType>()) {		} else if (auto toFloatType = toType.dyn_cast<FloatType>()) {
// If operand is integer, cast directly to the float type.		// If operand is integer, cast directly to the float type.
// Note that it is unclear how to cast from BF16<->FP16.		// Note that it is unclear how to cast from BF16<->FP16.
if (operand.getType().isa<IntegerType>())		if (operand.getType().isa<IntegerType>())
return builder.create<SIToFPOp>(loc, toFloatType, operand);		return builder.create<SIToFPOp>(loc, toFloatType, operand);
if (auto fromFloatType = operand.getType().dyn_cast<FloatType>()) {		if (auto fromFloatType = operand.getType().dyn_cast<FloatType>()) {
if (toFloatType.getWidth() > fromFloatType.getWidth())		if (toFloatType.getWidth() > fromFloatType.getWidth())
return builder.create<FPExtOp>(loc, toFloatType, operand);		return builder.create<FPExtOp>(loc, toFloatType, operand);
else if (toFloatType.getWidth() < fromFloatType.getWidth())		if (toFloatType.getWidth() < fromFloatType.getWidth())
return builder.create<FPTruncOp>(loc, toFloatType, operand);		return builder.create<FPTruncOp>(loc, toFloatType, operand);
}		}
}		}

emitWarning(operand.getLoc()) << "could not cast operand of type "		emitWarning(operand.getLoc()) << "could not cast operand of type "
<< operand.getType() << " to " << toType;		<< operand.getType() << " to " << toType;
return operand;		return operand;
}		}

Value applyfn__add(Value lhs, Value rhs) {		Value applyfn__add(Value lhs, Value rhs) {
OpBuilder builder = getBuilder(lhs);		OpBuilder builder = getBuilder();
if (isFloatingPoint(lhs))		if (isFloatingPoint(lhs))
return builder.create<AddFOp>(lhs.getLoc(), lhs, rhs);		return builder.create<AddFOp>(lhs.getLoc(), lhs, rhs);
else if (isInteger(lhs))		if (isInteger(lhs))
return builder.create<AddIOp>(lhs.getLoc(), lhs, rhs);		return builder.create<AddIOp>(lhs.getLoc(), lhs, rhs);
llvm_unreachable("unsupported non numeric type");		llvm_unreachable("unsupported non numeric type");
}		}

		Value applyfn__sub(Value lhs, Value rhs) {
		OpBuilder builder = getBuilder();
		if (isFloatingPoint(lhs))
		return builder.create<SubFOp>(lhs.getLoc(), lhs, rhs);
		if (isInteger(lhs))
		return builder.create<SubIOp>(lhs.getLoc(), lhs, rhs);
		llvm_unreachable("unsupported non numeric type");
		}

Value applyfn__mul(Value lhs, Value rhs) {		Value applyfn__mul(Value lhs, Value rhs) {
OpBuilder builder = getBuilder(lhs);		OpBuilder builder = getBuilder();
if (isFloatingPoint(lhs))		if (isFloatingPoint(lhs))
return builder.create<MulFOp>(lhs.getLoc(), lhs, rhs);		return builder.create<MulFOp>(lhs.getLoc(), lhs, rhs);
else if (isInteger(lhs))		if (isInteger(lhs))
return builder.create<MulIOp>(lhs.getLoc(), lhs, rhs);		return builder.create<MulIOp>(lhs.getLoc(), lhs, rhs);
llvm_unreachable("unsupported non numeric type");		llvm_unreachable("unsupported non numeric type");
}		}

void yieldOutputs(ValueRange values) {		void yieldOutputs(ValueRange values) {
assert(!values.empty() && "linalg ops must yield outputs");		assert(!values.empty() && "linalg ops must yield outputs");
if (values.empty())		if (values.empty())
return;		return;
Value first = values.front();		Value first = values.front();
OpBuilder builder = getBuilder(first);		OpBuilder builder = getBuilder();
builder.create<YieldOp>(first.getLoc(), values);		builder.create<YieldOp>(first.getLoc(), values);
}		}

		Value constant(std::string value) {
		OpBuilder builder = getBuilder();
		Location loc = builder.getUnknownLoc();
		Attribute valueAttr = parseAttribute(value, builder.getContext());
		return builder.create<ConstantOp>(loc, valueAttr.getType(), valueAttr);
		}

		Value index(int64_t dim) {
		OpBuilder builder = getBuilder();
		return builder.create<IndexOp>(builder.getUnknownLoc(), dim);
		}

		Type getIntegerType(unsigned width) {
		return IntegerType::get(context, width);
		}

		Type getFloat32Type() { return Float32Type::get(context); }

		Type getFloat64Type() { return Float64Type::get(context); }

private:		private:
		MLIRContext *context;
Block &block;		Block &block;

bool isFloatingPoint(Value value) { return value.getType().isa<FloatType>(); }		bool isFloatingPoint(Value value) { return value.getType().isa<FloatType>(); }
bool isInteger(Value value) { return value.getType().isa<IntegerType>(); }		bool isInteger(Value value) { return value.getType().isa<IntegerType>(); }

OpBuilder getBuilder(Value value) {		OpBuilder getBuilder() {
OpBuilder builder(value.getContext());		OpBuilder builder(context);
builder.setInsertionPointToEnd(&block);		builder.setInsertionPointToEnd(&block);
return builder;		return builder;
}		}
};		};

} // namespace		} // namespace

//===----------------------------------------------------------------------===//		//===----------------------------------------------------------------------===//
▲ Show 20 Lines • Show All 1,164 Lines • ▼ Show 20 Lines	else
intStrides.push_back(ShapedType::kDynamicStrideOrOffset);		intStrides.push_back(ShapedType::kDynamicStrideOrOffset);
}		}
auto layout =		auto layout =
makeStridedLinearLayoutMap(intStrides, intOffset, type.getContext());		makeStridedLinearLayoutMap(intStrides, intOffset, type.getContext());
return canonicalizeStridedLayout(		return canonicalizeStridedLayout(
MemRefType::Builder(type).setShape(newSizes).setAffineMaps({layout}));		MemRefType::Builder(type).setShape(newSizes).setAffineMaps({layout}));
}		}


template <typename AffineExprTy>		template <typename AffineExprTy>
unsigned getMaxPosOfType(ArrayRef<ReassociationExprs> exprArrays) {		unsigned getMaxPosOfType(ArrayRef<ReassociationExprs> exprArrays) {
unsigned pos = 0;		unsigned pos = 0;
for (const auto &exprs : exprArrays) {		for (const auto &exprs : exprArrays) {
for (auto expr : exprs) {		for (auto expr : exprs) {
expr.walk([&pos](AffineExpr e) {		expr.walk([&pos](AffineExpr e) {
if (auto d = e.dyn_cast<AffineExprTy>())		if (auto d = e.dyn_cast<AffineExprTy>())
pos = std::max(pos, d.getPosition());		pos = std::max(pos, d.getPosition());
▲ Show 20 Lines • Show All 1,673 Lines • Show Last 20 Lines

mlir/python/mlir/dialects/linalg/opdsl/lang/comprehension.py

# Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.		# Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
# See https://llvm.org/LICENSE.txt for license information.		# See https://llvm.org/LICENSE.txt for license information.
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception		# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
"""Model classes representing a tensor comprehension.		"""Model classes representing a tensor comprehension.

These classes model the language more at an AST level as evaluated. Reasoning		These classes model the language more at an AST level as evaluated. Reasoning
about it typically involves processing this form into config objects that		about it typically involves processing this form into config objects that
represent actual op definitions (i.e. YAML).		represent actual op definitions (i.e. YAML).
"""		"""

from typing import Any, Callable, Dict, List, Optional, Sequence, Set, Tuple, Union		from typing import Any, Dict, List, Optional, Sequence, Set, Tuple, Union

from mlir import ir as _ir		from mlir import ir as _ir

from .affine import *		from .affine import *
from .scalar_expr import *		from .scalar_expr import *
from .types import *		from .types import *
from .yaml_helper import *		from .yaml_helper import *

Show All 11 Lines	def visit_tensor_exprs(self, callback):
"""Visits all tensor expression reachable by the expression."""		"""Visits all tensor expression reachable by the expression."""
callback(self)		callback(self)

def _get_all_dim_defs(self) -> Set[DimDef]:		def _get_all_dim_defs(self) -> Set[DimDef]:
"""Recursively gets all DimDef affine expressions that are referenced."""		"""Recursively gets all DimDef affine expressions that are referenced."""
results = set()		results = set()

def visit_dim_def(dim_def):		def visit_dim_def(dim_def):
if isinstance(dim_def, DimDef):		if isinstance(dim_def, DimDef):
results.add(dim_def)		results.add(dim_def)

def visit_affine_exprs(expr):		def visit_affine_exprs(expr):
if isinstance(expr, TensorUse):		if isinstance(expr, TensorUse):
for ind in expr.indices:		for ind in expr.indices:
ind.visit_affine_exprs(visit_dim_def)		ind.visit_affine_exprs(visit_dim_def)
if isinstance(expr, ReduceApply):		if isinstance(expr, ReduceApply):
for ind in expr.reduce.reduce_dims:		for ind in expr.reduce.reduce_dims:
ind.visit_affine_exprs(visit_dim_def)		ind.visit_affine_exprs(visit_dim_def)

self.visit_tensor_exprs(visit_affine_exprs)		self.visit_tensor_exprs(visit_affine_exprs)
return results		return results

def collect_uses(self, uses: Set["TensorUse"]):		def collect_uses(self, uses: Set["TensorUse"]):
"""Collects all TensorUses reachable through this expression."""		"""Collects all TensorUses reachable through this expression."""

def visit_tensor_use(expr):		def visit_tensor_use(expr):
if isinstance(expr, TensorUse):		if isinstance(expr, TensorUse):
uses.add(expr)		uses.add(expr)

self.visit_tensor_exprs(visit_tensor_use)		self.visit_tensor_exprs(visit_tensor_use)

def collect_indices(self, indices: Set["index"]):		def collect_indices(self, indices: Set["index"]):
"""Collects all index accesses reachable through this expression."""		"""Collects all index accesses reachable through this expression."""

def visit_index(expr):		def visit_index(expr):
if isinstance(expr, index):		if isinstance(expr, index):
indices.add(expr)		indices.add(expr)

self.visit_tensor_exprs(visit_index)		self.visit_tensor_exprs(visit_index)

def collect_captures(self, captures: Set["CaptureDef"]):		def collect_captures(self, captures: Set["CaptureDef"]):
"""Collects all CaptureDefs reachable through this expression."""		"""Collects all CaptureDefs reachable through this expression."""

def visit_capture_def(expr):		def visit_capture_def(expr):
if isinstance(expr, CaptureDef):		if isinstance(expr, CaptureDef):
captures.add(expr)		captures.add(expr)

self.visit_tensor_exprs(visit_capture_def)		self.visit_tensor_exprs(visit_capture_def)

def __add__(self, rhs: "TensorExpression") -> "TensorExpression":		def __add__(self, rhs: "TensorExpression") -> "TensorExpression":
return PrimFn.add(self, rhs)		return PrimFn.add(self, rhs)

def __mul__(self, rhs) -> "TensorExpression":		def __mul__(self, rhs) -> "TensorExpression":
return PrimFn.mul(self, rhs)		return PrimFn.mul(self, rhs)

▲ Show 20 Lines • Show All 74 Lines • ▼ Show 20 Lines	def attach(self, index: int, tensor_name: str, owner: "LinalgOpDef"):
if self.owner:		if self.owner:
raise ValueError(f"TensorDef already registered with op: {self}")		raise ValueError(f"TensorDef already registered with op: {self}")
self.registered_index = index		self.registered_index = index
self.tensor_name = tensor_name		self.tensor_name = tensor_name
self.owner = owner		self.owner = owner

def __getitem__(self, dims) -> TensorUse:		def __getitem__(self, dims) -> TensorUse:
assert self.owner, "TensorDef is not attached to an op"		assert self.owner, "TensorDef is not attached to an op"
state = AffineBuildState(global_state=self.owner._affine_state,		state = AffineBuildState(
allow_new_symbols=False)		global_state=self.owner._affine_state, allow_new_symbols=False)
if not isinstance(dims, tuple):		if not isinstance(dims, tuple):
dims = (dims,) # Handle single subscript case.		dims = (dims,) # Handle single subscript case.
# Special case: (None) is a 0d-scalar use.		# Special case: (None) is a 0d-scalar use.
if dims == (None,):		if dims == (None,):
dims = ()		dims = ()

exprs = []		exprs = []
for expr_def in dims:		for expr_def in dims:
Show All 19 Lines	class TensorDef:
def __hash__(self):		def __hash__(self):
return hash(id(self))		return hash(id(self))

def __repr__(self):		def __repr__(self):
output = "OUTPUT " if self.output else ""		output = "OUTPUT " if self.output else ""
return (f"{self.tensor_name}:TensorDef({output}{repr(self.type_var)}, "		return (f"{self.tensor_name}:TensorDef({output}{repr(self.type_var)}, "
f"shape={self.shape})")		f"shape={self.shape})")


class CaptureDef(TensorExpression):		class CaptureDef(TensorExpression):
"""Defines an SSA value captured by the operation.		"""Defines an SSA value captured by the operation.

The captured SSA values are not indexed by the indexing_maps of the		The captured SSA values are not indexed by the indexing_maps of the
structured op (as opposed to memrefs and tensors). A unique name		structured op (as opposed to memrefs and tensors). A unique name
identifies the captures and an index determines their position the		identifies the captures and an index determines their position the
operation's parameter list.		operation's parameter list.
"""		"""
Show All 14 Lines	def attach(self, index: int, capture_name: str, owner: "LinalgOpDef"):
self.owner = owner		self.owner = owner

def to_scalar_expression(self) -> ScalarExpression:		def to_scalar_expression(self) -> ScalarExpression:
return ScalarCapture(self.capture_name).expr()		return ScalarCapture(self.capture_name).expr()

def __repr__(self):		def __repr__(self):
return (f"{self.capture_name}:CaptureDef({repr(self.type_var)})")		return (f"{self.capture_name}:CaptureDef({repr(self.type_var)})")


class Comprehension:		class Comprehension:
"""Represents a single comprehension."""		"""Represents a single comprehension."""

def __init__(self, *bindings: Tuple[TensorUse, TensorExpression]):		def __init__(self, *bindings: Tuple[TensorUse, TensorExpression]):
self.definitions = list() # List[TensorUse]		self.definitions = list() # List[TensorUse]
self.values = list() # List[TensorExpression]		self.values = list() # List[TensorExpression]

# Find the lhs to reduction rhs.		# Find the lhs to reduction rhs.
▲ Show 20 Lines • Show All 92 Lines • ▼ Show 20 Lines	class PrimApply(TensorExpression):
def visit_tensor_exprs(self, callback):		def visit_tensor_exprs(self, callback):
super().visit_tensor_exprs(callback)		super().visit_tensor_exprs(callback)
for arg in self.args:		for arg in self.args:
arg.visit_tensor_exprs(callback)		arg.visit_tensor_exprs(callback)

def __repr__(self):		def __repr__(self):
return f"{repr(self.prim)}({', '.join(repr(a) for a in self.args)})"		return f"{repr(self.prim)}({', '.join(repr(a) for a in self.args)})"


class const(TensorExpression):		class const(TensorExpression):
"""Returns the given constant floating point or integer value."""		"""Returns the given constant floating point or integer value."""

def __init__(self, type_var: TypeVar, value: Any):		def __init__(self, value: Any):
if not isinstance(type_var, TypeVar):		with _ir.Context():
raise ValueError(f"const requires a TypeVar. Got: {repr(type_var)}")		if isinstance(value, float):
if not (isinstance(value, float) or isinstance(value, int)):		self.value = str(_ir.FloatAttr.get_f64(float(value)))
		elif isinstance(value, int):
		self.value = str(
		_ir.IntegerAttr.get(_ir.IntegerType.get_signless(64), int(value)))
		else:
raise ValueError(f"const requires int or float. Got: {type(value)}")		raise ValueError(f"const requires int or float. Got: {type(value)}")
self.type_var = type_var
self.value = value

def to_scalar_expression(self) -> ScalarExpression:		def to_scalar_expression(self) -> ScalarExpression:
return ScalarConst(self.type_var, self.value).expr()		return ScalarConst(self.value).expr()

def __repr__(self):		def __repr__(self):
return f"const({self.type_var}, {self.value})"		return f"const({self.type_var}, {self.value})"


class index(TensorExpression):		class index(TensorExpression):
"""Returns the iteration index for a given dimension name.		"""Returns the iteration index for a given dimension name.

Resolves the given dimension name to obtain its position in the iteration		Resolves the given dimension name to obtain its position in the iteration
domain of the operation.		domain of the operation.
"""		"""

def __init__(self, dim : DimDef):		def __init__(self, dim: DimDef):
self.dim_def = dim		self.dim_def = dim
self.dim = -1		self.dim = -1

def resolve_dimension_name(self, affine_state: AffineBuildState):		def resolve_dimension_name(self, affine_state: AffineBuildState):
self.dim = affine_state.get_dim(self.dim_def.dimname)		self.dim = affine_state.get_dim(self.dim_def.dimname)

def to_scalar_expression(self) -> ScalarExpression:		def to_scalar_expression(self) -> ScalarExpression:
assert self.dim != -1, "Dimension name not resolved"		assert self.dim != -1, "Dimension name not resolved"
▲ Show 20 Lines • Show All 58 Lines • ▼ Show 20 Lines

ContractionOpInterface = OpInterfaceDef("LinalgContractionOpInterface")		ContractionOpInterface = OpInterfaceDef("LinalgContractionOpInterface")


class OpMetadataDef(YAMLObject):		class OpMetadataDef(YAMLObject):
"""Metadata about the op (generally not behavior impacting)."""		"""Metadata about the op (generally not behavior impacting)."""
yaml_tag = "!LinalgOpMetadata"		yaml_tag = "!LinalgOpMetadata"

def __init__(self, name: str, cpp_class_name: Optional[str], doc: Optional[str]):		def __init__(self, name: str, cpp_class_name: Optional[str],
		doc: Optional[str]):
self.name = name		self.name = name
self.cpp_class_name = cpp_class_name if cpp_class_name is not None else name		self.cpp_class_name = cpp_class_name if cpp_class_name is not None else name
self.doc = doc		self.doc = doc
self.implements = [] # type: List[OpInterfaceDef]		self.implements = [] # type: List[OpInterfaceDef]

def to_yaml_custom_dict(self):		def to_yaml_custom_dict(self):
d = dict(		d = dict(
name=self.name,		name=self.name,
cpp_class_name=self.cpp_class_name,		cpp_class_name=self.cpp_class_name,
doc=self.doc,		doc=self.doc,
)		)
if self.implements:		if self.implements:
d["implements"] = [intr.cpp_name for intr in self.implements]		d["implements"] = [intr.cpp_name for intr in self.implements]
return d		return d


class LinalgOpDef:		class LinalgOpDef:
"""Definition of a linalg op."""		"""Definition of a linalg op."""

def __init__(self,		def __init__(self,
name: str,		name: str,
cpp_class_name: Optional[str] = None,		cpp_class_name: Optional[str] = None,
doc: Optional[str] = None):		doc: Optional[str] = None):
self.metadata = OpMetadataDef(name=name, cpp_class_name=cpp_class_name, doc=doc)		self.metadata = OpMetadataDef(
		name=name, cpp_class_name=cpp_class_name, doc=doc)
self.registered_tensors = dict() # type: Dict[str, TensorDef]		self.registered_tensors = dict() # type: Dict[str, TensorDef]
self.registered_captures = dict() # type: Dict[str, CaptureDef]		self.registered_captures = dict() # type: Dict[str, CaptureDef]
self.comprehensions = list() # type: List[Comprehension]		self.comprehensions = list() # type: List[Comprehension]
self._affine_state = AffineBuildState()		self._affine_state = AffineBuildState()

@property		@property
def inputs(self) -> Sequence[TensorDef]:		def inputs(self) -> Sequence[TensorDef]:
return [t for t in self.registered_tensors.values() if not t.output]		return [t for t in self.registered_tensors.values() if not t.output]
Show All 35 Lines

mlir/python/mlir/dialects/linalg/opdsl/lang/config.py

# Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.		# Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
# See https://llvm.org/LICENSE.txt for license information.		# See https://llvm.org/LICENSE.txt for license information.
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception		# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
"""Represents configured ops as emitted for code generation.		"""Represents configured ops as emitted for code generation.

Classes in this module generally are directly serializable to YAML for use		Classes in this module generally are directly serializable to YAML for use
by the code generator.		by the code generator.

TODO: These should just be dumb containers or serialization code but they		TODO: These should just be dumb containers or serialization code but they
currently encode too many details of how the language is interpreted. Move this		currently encode too many details of how the language is interpreted. Move this
to helpers on the comprehension objects themselves.		to helpers on the comprehension objects themselves.
"""		"""

from typing import Any, Dict, Optional		from typing import Dict, Optional

from mlir import ir as _ir		from mlir import ir as _ir

from .comprehension import *		from .comprehension import *
from .yaml_helper import *		from .yaml_helper import *

__all__ = [		__all__ = [
"LinalgStructuredOpConfig",		"LinalgStructuredOpConfig",
▲ Show 20 Lines • Show All 42 Lines • ▼ Show 20 Lines	return dict(
usage=self.usage,		usage=self.usage,
shape=_serialize_affine_map(self.shape_map),		shape=_serialize_affine_map(self.shape_map),
element_type_var=self.tensor_def.type_var.name,		element_type_var=self.tensor_def.type_var.name,
)		)

def __repr__(self):		def __repr__(self):
return f"Def({self.tensor_def}, shape_map={self.shape_map}, indexing_map={self.indexing_map})"		return f"Def({self.tensor_def}, shape_map={self.shape_map}, indexing_map={self.indexing_map})"


class CaptureDefConfig(YAMLObject):		class CaptureDefConfig(YAMLObject):
"""Wrapper around a CaptureDef."""		"""Wrapper around a CaptureDef."""
yaml_tag = "LinalgCaptureDef"		yaml_tag = "LinalgCaptureDef"

def __init__(self, capture_def: CaptureDef):		def __init__(self, capture_def: CaptureDef):
self.capture_def = capture_def		self.capture_def = capture_def

def to_yaml_custom_dict(self):		def to_yaml_custom_dict(self):
Show All 27 Lines	if self.static_indexing_maps is not None:
_serialize_affine_map(m) for m in self.static_indexing_maps		_serialize_affine_map(m) for m in self.static_indexing_maps
])		])
raise ValueError(		raise ValueError(
f"LinalgIndexingMapsConfig must have one type of indexing map"		f"LinalgIndexingMapsConfig must have one type of indexing map"
f"(got none)")		f"(got none)")


class LinalgStructuredOpConfig(YAMLObject):		class LinalgStructuredOpConfig(YAMLObject):
"""Configuration for metadata sufficient to construct a linalg single		"""Configuration for metadata sufficient to construct a linalg named op."""
contraction named op."""

yaml_tag = "!LinalgStructuredOpConfig"		yaml_tag = "!LinalgStructuredOpConfig"

def __init__(self,		def __init__(self,
comprehension: Comprehension,		comprehension: Comprehension,
context: Optional[_ir.Context] = None):		context: Optional[_ir.Context] = None):
self.context = context if context is not None else _ir.Context()		self.context = context if context is not None else _ir.Context()
self.affine_state = AffineBuildState()		self.affine_state = AffineBuildState()
Show All 25 Lines	def __init__(self,
for use in collected_uses:		for use in collected_uses:
self.add_use(use)		self.add_use(use)

# Now normalize all defs and uses indexing maps now that full count of		# Now normalize all defs and uses indexing maps now that full count of
# dims and symbols are known.		# dims and symbols are known.
for cuse in self.uses.values():		for cuse in self.uses.values():
cuse.indexing_map = self._normalize_affine_map(cuse.indexing_map)		cuse.indexing_map = self._normalize_affine_map(cuse.indexing_map)
for cdef in self.tensor_args.values():		for cdef in self.tensor_args.values():
cdef.shape_map = self._normalize_affine_map(cdef.shape_map,		cdef.shape_map = self._normalize_affine_map(
with_dims=False)		cdef.shape_map, with_dims=False)

# Now for each write use, propagate the indexing maps from the use to the		# Now for each write use, propagate the indexing maps from the use to the
# tensor, ensuring that there are not conflicts.		# tensor, ensuring that there are not conflicts.
for write_use, _ in self.writes:		for write_use, _ in self.writes:
write_tensor_def = self.tensor_args[write_use.tensor_def]		write_tensor_def = self.tensor_args[write_use.tensor_def]
if write_tensor_def.indexing_map:		if write_tensor_def.indexing_map:
raise ValueError(		raise ValueError(
f"Unexpected multi-write to a single tensor: {write_tensor_def}")		f"Unexpected multi-write to a single tensor: {write_tensor_def}")
Show All 24 Lines	if len(all_reduction_dims) != 1:
f"All writes within a generic must have the same reduction "		f"All writes within a generic must have the same reduction "
f"dims. Got: {all_reduction_dims}")		f"dims. Got: {all_reduction_dims}")
self.reduction_dims = next(iter(all_reduction_dims))		self.reduction_dims = next(iter(all_reduction_dims))

# Check the index dimension exists and resolve		# Check the index dimension exists and resolve
for index in collected_indices:		for index in collected_indices:
if index.dim_def.dimname not in self.affine_state.all_dims:		if index.dim_def.dimname not in self.affine_state.all_dims:
raise ValueError(		raise ValueError(
f"The dimension {index.dim.dimname} is not part of the iteration "		f"The dimension {index.dim.dimname} is not part of the iteration "
f"domain {self.affine_state.all_dims}")		f"domain {self.affine_state.all_dims}")
index.resolve_dimension_name(self.affine_state)		index.resolve_dimension_name(self.affine_state)

# Generate the scalar assignments (used to build a body).		# Generate the scalar assignments (used to build a body).
self.assignments = [		self.assignments = [
ScalarAssign(write_use.tensor_name, read_expr.to_scalar_expression())		ScalarAssign(write_use.tensor_name, read_expr.to_scalar_expression())
for write_use, read_expr in self.writes		for write_use, read_expr in self.writes
]		]

@property		@property
def ordered_tensor_args(self) -> Sequence[TensorDefConfig]:		def ordered_tensor_args(self) -> Sequence[TensorDefConfig]:
return sorted(self.tensor_args.values(),		return sorted(
		self.tensor_args.values(),
key=lambda tdc: tdc.tensor_def.registered_index)		key=lambda tdc: tdc.tensor_def.registered_index)

@property		@property
def ordered_tensor_uses(self) -> Sequence[TensorUseConfig]:		def ordered_tensor_uses(self) -> Sequence[TensorUseConfig]:
return sorted(self.uses.values(),		return sorted(
		self.uses.values(),
key=lambda tuc: tuc.tensor_use.tensor_def.registered_index)		key=lambda tuc: tuc.tensor_use.tensor_def.registered_index)

@property		@property
def ordered_capture_args(self) -> Sequence[CaptureDefConfig]:		def ordered_capture_args(self) -> Sequence[CaptureDefConfig]:
return sorted(self.capture_args.values(),		return sorted(
		self.capture_args.values(),
key=lambda cdc: cdc.capture_def.registered_index)		key=lambda cdc: cdc.capture_def.registered_index)

@property		@property
def ordered_dims(self) -> Sequence[Tuple[str, int]]:		def ordered_dims(self) -> Sequence[Tuple[str, int]]:
"""Gets the ordered list of dim bindings (symbolic name, position).		"""Gets the ordered list of dim bindings (symbolic name, position).

TODO: The original parser relies on parse ordering to arrive at the		TODO: The original parser relies on parse ordering to arrive at the
iterator types, but that ordering is not defined on the Python side, so		iterator types, but that ordering is not defined on the Python side, so
this may be ambiguous.		this may be ambiguous.
Show All 14 Lines	def get_type(symbolic_name, position):
return "parallel"		return "parallel"

return [get_type(*dim) for dim in self.ordered_dims]		return [get_type(*dim) for dim in self.ordered_dims]

def add_tensor_arg(self, tensor_def: TensorDef):		def add_tensor_arg(self, tensor_def: TensorDef):
if tensor_def in self.tensor_args:		if tensor_def in self.tensor_args:
return		return
with self.context:		with self.context:
local_state = AffineBuildState(global_state=self.affine_state,		local_state = AffineBuildState(
allow_new_dims=False)		global_state=self.affine_state, allow_new_dims=False)
exprs = []		exprs = []
for expr in tensor_def.shape:		for expr in tensor_def.shape:
exprs.append(expr.build(state=local_state))		exprs.append(expr.build(state=local_state))
assert local_state.local_dim_count == 0		assert local_state.local_dim_count == 0
indexing_map = _ir.AffineMap.get(dim_count=0,		indexing_map = _ir.AffineMap.get(
symbol_count=local_state.symbol_count,		dim_count=0, symbol_count=local_state.symbol_count, exprs=exprs)
exprs=exprs)

def_config = TensorDefConfig(tensor_def, indexing_map)		def_config = TensorDefConfig(tensor_def, indexing_map)
self.tensor_args[tensor_def] = def_config		self.tensor_args[tensor_def] = def_config

def add_use(self, tensor_use: TensorUse):		def add_use(self, tensor_use: TensorUse):
if tensor_use in self.uses:		if tensor_use in self.uses:
return		return
with self.context:		with self.context:
local_state = AffineBuildState(global_state=self.affine_state,		local_state = AffineBuildState(
allow_new_symbols=False)		global_state=self.affine_state, allow_new_symbols=False)
exprs = []		exprs = []
for expr in tensor_use.indices:		for expr in tensor_use.indices:
exprs.append(expr.build(state=local_state))		exprs.append(expr.build(state=local_state))
assert local_state.local_symbol_count == 0		assert local_state.local_symbol_count == 0
indexing_map = _ir.AffineMap.get(dim_count=local_state.dim_count,		indexing_map = _ir.AffineMap.get(
		dim_count=local_state.dim_count,
symbol_count=local_state.symbol_count,		symbol_count=local_state.symbol_count,
exprs=exprs)		exprs=exprs)

use_config = TensorUseConfig(tensor_use, indexing_map)		use_config = TensorUseConfig(tensor_use, indexing_map)
self.uses[tensor_use] = use_config		self.uses[tensor_use] = use_config

def add_capture_arg(self, capture_def: CaptureDef):		def add_capture_arg(self, capture_def: CaptureDef):
if capture_def in self.capture_args:		if capture_def in self.capture_args:
return		return
def_config = CaptureDefConfig(capture_def)		def_config = CaptureDefConfig(capture_def)
self.capture_args[capture_def] = def_config		self.capture_args[capture_def] = def_config

def _normalize_affine_map(self,		def _normalize_affine_map(self,
affine_map: _ir.AffineMap,		affine_map: _ir.AffineMap,
with_dims: bool = True) -> _ir.AffineMap:		with_dims: bool = True) -> _ir.AffineMap:
"""Normalizes an indexing map to have the max known symbols and dims."""		"""Normalizes an indexing map to have the max known symbols and dims."""
with self.context:		with self.context:
return _ir.AffineMap.get(		return _ir.AffineMap.get(
dim_count=self.affine_state.dim_count if with_dims else 0,		dim_count=self.affine_state.dim_count if with_dims else 0,
symbol_count=self.affine_state.symbol_count,		symbol_count=self.affine_state.symbol_count,
exprs=list(affine_map.results))		exprs=list(affine_map.results))

def to_yaml_custom_dict(self):		def to_yaml_custom_dict(self):
self_dict = dict(		self_dict = dict(args=self.ordered_tensor_args)
args=self.ordered_tensor_args,		if self.ordered_capture_args:
		stellaraccidentUnsubmitted Done Reply Inline Actions In Python, one generally just says `if self.ordered_capture_args` and let the default empty-container-to-False rule apply. stellaraccident: In Python, one generally just says `if self.ordered_capture_args` and let the default empty…
captures=self.ordered_capture_args,		self_dict["captures"] = self.ordered_capture_args
# TODO: Refactor the hierarchy internally when supporting more		# TODO: Refactor the hierarchy internally when supporting more
# than static (preserving this serialized form).		# than static (preserving this serialized form).
indexing_maps=LinalgIndexingMapsConfig(		self_dict["indexing_maps"] = LinalgIndexingMapsConfig(
static_indexing_maps=self.indexing_maps),		static_indexing_maps=self.indexing_maps)
iterator_types=self.iterator_types,		self_dict["iterator_types"] = self.iterator_types
assignments=self.assignments,		self_dict["assignments"] = self.assignments
)
return self_dict		return self_dict

def __repr__(self):		def __repr__(self):
lines = [f"LinalgGenericOpConfig(reduction_dims={self.reduction_dims},"]		lines = [f"LinalgGenericOpConfig(reduction_dims={self.reduction_dims},"]
lines.append("tensor_args=[")		lines.append("tensor_args=[")
for def_config in self.ordered_tensor_args:		for def_config in self.ordered_tensor_args:
lines.append(f" {repr(def_config)}")		lines.append(f" {repr(def_config)}")
lines.append("], capture_args=[")		lines.append("], capture_args=[")
Show All 34 Lines	class LinalgOpConfig(YAMLObject):
def from_linalg_op_def(		def from_linalg_op_def(
tc_op_def: LinalgOpDef,		tc_op_def: LinalgOpDef,
context: Optional[_ir.Context] = None) -> Sequence["LinalgOpConfig"]:		context: Optional[_ir.Context] = None) -> Sequence["LinalgOpConfig"]:
"""Expands a LinalgOpDef into corresponding Linalg configured ops."""		"""Expands a LinalgOpDef into corresponding Linalg configured ops."""
# TODO: Many LinalgOpDef patterns need to expand to multiple generics.		# TODO: Many LinalgOpDef patterns need to expand to multiple generics.
assert len(		assert len(
tc_op_def.comprehensions) == 1, "Only one comprehension supported"		tc_op_def.comprehensions) == 1, "Only one comprehension supported"
return [		return [
LinalgOpConfig(tc_op_def.metadata,		LinalgOpConfig(
structured_op=LinalgStructuredOpConfig(		tc_op_def.metadata,
tc_op_def.comprehensions[0], context)),		structured_op=LinalgStructuredOpConfig(tc_op_def.comprehensions[0],
		context)),
]		]

def __repr__(self):		def __repr__(self):
return (f"LinalgOpConfig(metadata={self.metadata},\n"		return (f"LinalgOpConfig(metadata={self.metadata},\n"
f"structured_op={self.structured_op})")		f"structured_op={self.structured_op})")

mlir/python/mlir/dialects/linalg/opdsl/lang/emitter.py

# Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.		# Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
# See https://llvm.org/LICENSE.txt for license information.		# See https://llvm.org/LICENSE.txt for license information.
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception		# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception

from typing import Any, Dict, Sequence		from typing import Dict, Sequence

from mlir.ir import *		from mlir.ir import *
from mlir.dialects import linalg		from mlir.dialects import linalg
from mlir.dialects import std		from mlir.dialects import std
# TODO: resolve name collision for Linalg functionality that is injected inside		# TODO: resolve name collision for Linalg functionality that is injected inside
# the _mlir.dialects.linalg directly via pybind.		# the _mlir.dialects.linalg directly via pybind.
from _mlir.dialects.linalg import fill_builtin_region		from _mlir.dialects.linalg import fill_builtin_region

from .scalar_expr import *		from .scalar_expr import *
from .config import *		from .config import *

__all__ = [		__all__ = [
"emit_generic_structured_op",		"emit_generic_structured_op",
"emit_named_structured_op",		"emit_named_structured_op",
]		]


def isa(cls : Type, ty : Type):		def isa(cls: Type, ty: Type):
try:		try:
cls(ty)		cls(ty)
return True		return True
except ValueError:		except ValueError:
return False		return False


def prepare_common_structured_op(op_config: LinalgStructuredOpConfig,		def prepare_common_structured_op(op_config: LinalgStructuredOpConfig,
*ins: Value,		*ins: Value, outs: Sequence[Value],
outs: Sequence[Value],
captures: Sequence[Value]):		captures: Sequence[Value]):
all_arg_defs = op_config.ordered_tensor_args		all_arg_defs = op_config.ordered_tensor_args
in_arg_defs = [arg for arg in all_arg_defs if arg.usage == "input"]		in_arg_defs = [arg for arg in all_arg_defs if arg.usage == "input"]
out_arg_defs = [arg for arg in all_arg_defs if arg.usage == "output"]		out_arg_defs = [arg for arg in all_arg_defs if arg.usage == "output"]
capture_arg_defs = op_config.ordered_capture_args		capture_arg_defs = op_config.ordered_capture_args

# Verify outs and captures are sequences.		# Verify outs and captures are sequences.
if not isinstance(outs, Sequence):		if not isinstance(outs, Sequence):
Show All 37 Lines	def prepare_common_structured_op(op_config: LinalgStructuredOpConfig,
capture_arg_mapping = dict() # type: Dict[str, Value]		capture_arg_mapping = dict() # type: Dict[str, Value]
for arg_def, capture_value in zip(capture_arg_defs, captures):		for arg_def, capture_value in zip(capture_arg_defs, captures):
_add_type_mapping(arg_def.capture_def.type_var.name, capture_value.type,		_add_type_mapping(arg_def.capture_def.type_var.name, capture_value.type,
type_mapping)		type_mapping)
capture_arg_mapping[arg_def.capture_def.capture_name] = capture_value		capture_arg_mapping[arg_def.capture_def.capture_name] = capture_value

# Emit the generic op.		# Emit the generic op.
# TODO: Support emission of pure memref form.		# TODO: Support emission of pure memref form.
indexing_maps_attr = ArrayAttr.get(		indexing_maps_attr = ArrayAttr.get([
[AffineMapAttr.get(am)		AffineMapAttr.get(am)
# TODO: linalg verification does not currently allow symbols.		# TODO: linalg verification does not currently allow symbols.
# Compress them for now.		# Compress them for now.
for am in AffineMap.compress_unused_symbols(op_config.indexing_maps, Context.current)])		for am in AffineMap.compress_unused_symbols(op_config.indexing_maps,
		Context.current)
		])
iterator_types_attr = ArrayAttr.get(		iterator_types_attr = ArrayAttr.get(
[StringAttr.get(s) for s in op_config.iterator_types])		[StringAttr.get(s) for s in op_config.iterator_types])

return (all_arg_defs, in_arg_defs, out_arg_defs, outs, result_types,		return (all_arg_defs, in_arg_defs, out_arg_defs, outs, result_types,
type_mapping, capture_arg_mapping, indexing_maps_attr,		type_mapping, capture_arg_mapping, indexing_maps_attr,
iterator_types_attr)		iterator_types_attr)


▲ Show 20 Lines • Show All 41 Lines • ▼ Show 20 Lines	def emit_named_structured_op(op_config: LinalgStructuredOpConfig,
captures: Sequence[Value] = ()):		captures: Sequence[Value] = ()):
all_arg_defs, in_arg_defs, out_arg_defs, outs, result_types, type_mapping, \		all_arg_defs, in_arg_defs, out_arg_defs, outs, result_types, type_mapping, \
capture_arg_mapping, indexing_maps_attr, iterator_types_attr = \		capture_arg_mapping, indexing_maps_attr, iterator_types_attr = \
prepare_common_structured_op(op_config, *ins, outs = outs,		prepare_common_structured_op(op_config, *ins, outs = outs,
captures = captures)		captures = captures)

# If we get here, there must exist a builtin class `op_class_name`.		# If we get here, there must exist a builtin class `op_class_name`.
ctx = Context.current		ctx = Context.current
fully_qualified_name = 'linalg.' + op_name		fully_qualified_name = "linalg." + op_name
if (not ctx.is_registered_operation(fully_qualified_name) or		if (not ctx.is_registered_operation(fully_qualified_name) or
not op_class_name in linalg.__dict__.keys()):		not op_class_name in linalg.__dict__.keys()):
raise NotImplementedError(		raise NotImplementedError(
f"Unknown named op_name / op_class_name: {op_name} / {op_class_name}")		f"Unknown named op_name / op_class_name: {op_name} / {op_class_name}")

named_op = getattr(linalg, op_class_name)(ins, outs, result_types)		named_op = getattr(linalg, op_class_name)(ins, outs, result_types)
linalgDialect = ctx.get_dialect_descriptor("linalg")		linalgDialect = ctx.get_dialect_descriptor("linalg")
fill_builtin_region(linalgDialect, named_op.operation)		fill_builtin_region(linalgDialect, named_op.operation)
# Note: mlir-linalg-ods-yaml-gen.cpp uses a special linalg.memoized_indexing_maps		# Note: mlir-linalg-ods-yaml-gen.cpp uses a special linalg.memoized_indexing_maps
# attribute that the non-yaml path does not. The non-yaml path hardcodes the		# attribute that the non-yaml path does not. The non-yaml path hardcodes the
# indexing_maps in C++ directly.		# indexing_maps in C++ directly.
named_op.operation.attributes["linalg.memoized_indexing_maps"] = indexing_maps_attr		named_op.operation.attributes[
		"linalg.memoized_indexing_maps"] = indexing_maps_attr
# iterator_types are hardcoded in C++ both in the yaml and non-yaml path.		# iterator_types are hardcoded in C++ both in the yaml and non-yaml path.

if len(result_types) == 1:		if len(result_types) == 1:
return named_op.result		return named_op.result
else:		else:
return named_op.results		return named_op.results


class _BodyBuilder:		class _BodyBuilder:
"""Constructs a structured op body by evaluating assignments."""		"""Constructs a structured op body by evaluating assignments."""

def __init__(self,		def __init__(self, type_mapping: Dict[str, Type],
type_mapping: Dict[str, Type],
block_arg_mapping: Dict[str, Value],		block_arg_mapping: Dict[str, Value],
capture_arg_mapping: Dict[str, Value]):		capture_arg_mapping: Dict[str, Value]):
self.type_mapping = type_mapping		self.type_mapping = type_mapping
self.block_arg_mapping = block_arg_mapping		self.block_arg_mapping = block_arg_mapping
self.capture_arg_mapping = capture_arg_mapping		self.capture_arg_mapping = capture_arg_mapping
self.yield_mapping = dict() # type: Dict[str, Value]		self.yield_mapping = dict() # type: Dict[str, Value]

def assign(self, assignment: ScalarAssign):		def assign(self, assignment: ScalarAssign):
Show All 9 Lines	if expr.scalar_arg:
return self.block_arg_mapping[expr.scalar_arg.arg]		return self.block_arg_mapping[expr.scalar_arg.arg]
except KeyError:		except KeyError:
raise ValueError(f"Argument {expr.scalar_arg.arg} is not bound for "		raise ValueError(f"Argument {expr.scalar_arg.arg} is not bound for "
f"this structured op.")		f"this structured op.")
elif expr.scalar_capture:		elif expr.scalar_capture:
try:		try:
return self.capture_arg_mapping[expr.scalar_capture.capture]		return self.capture_arg_mapping[expr.scalar_capture.capture]
except KeyError:		except KeyError:
raise ValueError(f"Capture {expr.scalar_capture.capture} is not bound for "		raise ValueError(
		f"Capture {expr.scalar_capture.capture} is not bound for "
f"this structured op.")		f"this structured op.")
elif expr.scalar_const:		elif expr.scalar_const:
return self.constant(expr.scalar_const.type_var.name, expr.scalar_const.value)		value_attr = Attribute.parse(expr.scalar_const.value)
		return std.ConstantOp(value_attr.type, value_attr).result
elif expr.scalar_index:		elif expr.scalar_index:
return self.index(expr.scalar_index.dim)		dim_attr = IntegerAttr.get(
		IntegerType.get_signless(64), expr.scalar_index.dim)
		return linalg.IndexOp(IndexType.get(), dim_attr).result
elif expr.scalar_apply:		elif expr.scalar_apply:
try:		try:
fn = getattr(self, f"_eval_{expr.scalar_apply.fn_name}")		fn = getattr(self, f"_eval_{expr.scalar_apply.fn_name}")
except AttributeError:		except AttributeError:
raise ValueError(		raise ValueError(
f"Function '{expr.scalar_apply.fn_name}' is not a known "		f"Function '{expr.scalar_apply.fn_name}' is not a known "
"scalar body function")		"scalar body function")
operand_values = [		operand_values = [
self.expression(operand) for operand in expr.scalar_apply.operands		self.expression(operand) for operand in expr.scalar_apply.operands
]		]
return fn(*operand_values)		return fn(*operand_values)
elif expr.symbolic_cast:		elif expr.symbolic_cast:
operand_value = self.expression(expr.symbolic_cast.operand)		operand_value = self.expression(expr.symbolic_cast.operand)
return self.cast(expr.symbolic_cast.to_type.name, operand_value)		return self.cast(expr.symbolic_cast.to_type.name, operand_value)
raise NotImplementedError(f"Unimplemented scalar body expression: {expr}")		raise NotImplementedError(f"Unimplemented scalar body expression: {expr}")

def constant(self, type_var_name: str, value: Any) -> Value:
try:
type = self.type_mapping[type_var_name]
except KeyError:
raise ValueError(f"Unbound type variable '{type_var_name}' ("
f"expected one of {self.type_mappings.keys()}")
try:
if(_is_floating_point_type(type)):
return std.ConstantOp(type, FloatAttr.get(type, float(value))).result
elif(_is_integer_type(type)):
return std.ConstantOp(type, IntegerAttr.get(type, int(value))).result
except ValueError:
raise ValueError(f"Unable to cast value {value} to type {type}")
raise NotImplementedError(f"Unimplemented constant type {type}")

def index(self, dim: int) -> Value:
dim_attr = IntegerAttr.get(IntegerType.get_signless(64), dim)
return linalg.IndexOp(IndexType.get(), dim_attr).result

def cast(self, type_var_name: str, operand: Value) -> Value:		def cast(self, type_var_name: str, operand: Value) -> Value:
try:		try:
to_type = self.type_mapping[type_var_name]		to_type = self.type_mapping[type_var_name]
except KeyError:		except KeyError:
raise ValueError(f"Unbound type variable '{type_var_name}' ("		raise ValueError(f"Unbound type variable '{type_var_name}' ("
f"expected one of {self.type_mappings.keys()}")		f"expected one of {self.type_mappings.keys()}")
if operand.type == to_type:		if operand.type == to_type:
return operand		return operand
if _is_integer_type(to_type):		if _is_integer_type(to_type):
return self._cast_to_integer(to_type, operand)		return self._cast_to_integer(to_type, operand)
elif _is_floating_point_type(to_type):		elif _is_floating_point_type(to_type):
return self._cast_to_floating_point(to_type, operand)		return self._cast_to_floating_point(to_type, operand)

def _cast_to_integer(self, to_type: Type, operand: Value) -> Value:		def _cast_to_integer(self, to_type: Type, operand: Value) -> Value:
to_width = IntegerType(to_type).width		to_width = IntegerType(to_type).width
operand_type = operand.type		operand_type = operand.type
if _is_floating_point_type(operand_type):		if _is_floating_point_type(operand_type):
return std.FPToSIOp(to_type, operand).result		return std.FPToSIOp(to_type, operand).result
if _is_index_type(operand_type):		if _is_index_type(operand_type):
return std.IndexCastOp(to_type, operand).result		return std.IndexCastOp(to_type, operand).result
# Assume integer.		# Assume integer.
▲ Show 20 Lines • Show All 81 Lines • ▼ Show 20 Lines	for v in values:
types.append(t.element_type)		types.append(t.element_type)
return types		return types


def _get_tensor_def_names(		def _get_tensor_def_names(
*tensor_def_configs: TensorDefConfig) -> Sequence[str]:		*tensor_def_configs: TensorDefConfig) -> Sequence[str]:
return [tdc.tensor_def.tensor_name for tdc in tensor_def_configs]		return [tdc.tensor_def.tensor_name for tdc in tensor_def_configs]


def _add_type_mapping(name: str, type: Type, type_mapping: Dict[str, Type]):		def _add_type_mapping(name: str, type: Type, type_mapping: Dict[str, Type]):
if name in type_mapping:		if name in type_mapping:
if type_mapping[name] != type:		if type_mapping[name] != type:
raise ValueError(f"Cannot overwrite type mapping {name} = "		raise ValueError(f"Cannot overwrite type mapping {name} = "
f"{type_mapping[name]} by type {type}")		f"{type_mapping[name]} by type {type}")
type_mapping[name] = type		type_mapping[name] = type


def _is_floating_point_type(t: Type) -> bool:		def _is_floating_point_type(t: Type) -> bool:
# TODO: Create a FloatType in the Python API and implement the switch		# TODO: Create a FloatType in the Python API and implement the switch
# there.		# there.
return (F64Type.isinstance(t) or F32Type.isinstance(t) or		return (F64Type.isinstance(t) or F32Type.isinstance(t) or
F16Type.isinstance(t) or BF16Type.isinstance(t))		F16Type.isinstance(t) or BF16Type.isinstance(t))


def _is_integer_type(t: Type) -> bool:		def _is_integer_type(t: Type) -> bool:
return IntegerType.isinstance(t)		return IntegerType.isinstance(t)


def _is_index_type(t: Type) -> bool:		def _is_index_type(t: Type) -> bool:
return IndexType.isinstance(t)		return IndexType.isinstance(t)


def _get_floating_point_width(t: Type) -> int:		def _get_floating_point_width(t: Type) -> int:
# TODO: Create a FloatType in the Python API and implement the switch		# TODO: Create a FloatType in the Python API and implement the switch
# there.		# there.
if F64Type.isinstance(t):		if F64Type.isinstance(t):
return 64		return 64
if F32Type.isinstance(t):		if F32Type.isinstance(t):
return 32		return 32
if F16Type.isinstance(t):		if F16Type.isinstance(t):
return 16		return 16
if BF16Type.isinstance(t):		if BF16Type.isinstance(t):
return 16		return 16
raise NotImplementedError(f"Unhandled floating point type switch {t}")		raise NotImplementedError(f"Unhandled floating point type switch {t}")

mlir/python/mlir/dialects/linalg/opdsl/lang/scalar_expr.py

# Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.		# Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
# See https://llvm.org/LICENSE.txt for license information.		# See https://llvm.org/LICENSE.txt for license information.
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception		# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
"""Models DAGs of scalar math expressions.		"""Models DAGs of scalar math expressions.

Used for generating region bodies at the "math" level where they are still type		Used for generating region bodies at the "math" level where they are still type
polymorphic. This is modeled to be polymorphic by attribute name for interop		polymorphic. This is modeled to be polymorphic by attribute name for interop
with serialization schemes that are just plain-old-dicts.		with serialization schemes that are just plain-old-dicts.

These classes are typically not user accessed and are created as a by-product		These classes are typically not user accessed and are created as a by-product
of interpreting a comprehension DSL and model the operations to perform in the		of interpreting a comprehension DSL and model the operations to perform in the
op body. The class hierarchy is laid out to map well to a form of YAML that		op body. The class hierarchy is laid out to map well to a form of YAML that
can be easily consumed from the C++ side, not necessarily for ergonomics.		can be easily consumed from the C++ side, not necessarily for ergonomics.
"""		"""

from typing import Any, Optional, Sequence		from typing import Optional, Sequence

from .yaml_helper import *		from .yaml_helper import *
from .types import *		from .types import *

__all__ = [		__all__ = [
"ScalarAssign",		"ScalarAssign",
"ScalarApplyFn",		"ScalarApplyFn",
"ScalarArg",		"ScalarArg",
Show All 26 Lines	def __init__(self, arg: str):
self.arg = arg		self.arg = arg

def expr(self) -> "ScalarExpression":		def expr(self) -> "ScalarExpression":
return ScalarExpression(scalar_arg=self)		return ScalarExpression(scalar_arg=self)

def __repr__(self):		def __repr__(self):
return f"(ScalarArg({self.arg})"		return f"(ScalarArg({self.arg})"


class ScalarCapture:		class ScalarCapture:
"""A type of ScalarExpression that references a named capture."""		"""A type of ScalarExpression that references a named capture."""

def __init__(self, capture: str):		def __init__(self, capture: str):
self.capture = capture		self.capture = capture

def expr(self) -> "ScalarExpression":		def expr(self) -> "ScalarExpression":
return ScalarExpression(scalar_capture=self)		return ScalarExpression(scalar_capture=self)

def __repr__(self):		def __repr__(self):
return f"(ScalarCapture({self.capture})"		return f"(ScalarCapture({self.capture})"


class ScalarConst:		class ScalarConst:
"""A type of ScalarExpression representing a constant."""		"""A type of ScalarExpression representing a constant."""

def __init__(self, type_var: TypeVar, value: Any):		def __init__(self, value: str):
self.type_var = type_var
self.value = value		self.value = value

def expr(self) -> "ScalarExpression":		def expr(self) -> "ScalarExpression":
return ScalarExpression(scalar_const=self)		return ScalarExpression(scalar_const=self)

def __repr__(self):		def __repr__(self):
return f"(ScalarConst({self.type_var}, {self.value})"		return f"(ScalarConst({self.value})"


class ScalarIndex:		class ScalarIndex:
"""A type of ScalarExpression accessing an iteration index."""		"""A type of ScalarExpression accessing an iteration index."""

def __init__(self, dim : int):		def __init__(self, dim: int):
self.dim = dim		self.dim = dim

def expr(self) -> "ScalarExpression":		def expr(self) -> "ScalarExpression":
return ScalarExpression(scalar_index=self)		return ScalarExpression(scalar_index=self)

def __repr__(self):		def __repr__(self):
return f"(ScalarIndex({self.dim})"		return f"(ScalarIndex({self.dim})"


class ScalarSymbolicCast:		class ScalarSymbolicCast:
"""A type of ScalarExpression that symbolically casts an operand to a TypeVar.		"""A type of ScalarExpression that symbolically casts an operand to a TypeVar."""
"""

def __init__(self, to_type: TypeVar, operand: "ScalarExpression"):		def __init__(self, to_type: TypeVar, operand: "ScalarExpression"):
self.to_type = to_type		self.to_type = to_type
self.operand = operand		self.operand = operand

def expr(self) -> "ScalarExpression":		def expr(self) -> "ScalarExpression":
return ScalarExpression(symbolic_cast=self)		return ScalarExpression(symbolic_cast=self)

Show All 30 Lines	def __init__(self,
self.scalar_arg = scalar_arg		self.scalar_arg = scalar_arg
self.scalar_capture = scalar_capture		self.scalar_capture = scalar_capture
self.scalar_const = scalar_const		self.scalar_const = scalar_const
self.scalar_index = scalar_index		self.scalar_index = scalar_index
self.symbolic_cast = symbolic_cast		self.symbolic_cast = symbolic_cast

def to_yaml_custom_dict(self):		def to_yaml_custom_dict(self):
if self.scalar_apply:		if self.scalar_apply:
return dict(scalar_apply=dict(		return dict(
		scalar_apply=dict(
fn_name=self.scalar_apply.fn_name,		fn_name=self.scalar_apply.fn_name,
operands=list(self.scalar_apply.operands),		operands=list(self.scalar_apply.operands),
))		))
elif self.scalar_arg:		elif self.scalar_arg:
return dict(scalar_arg=self.scalar_arg.arg)		return dict(scalar_arg=self.scalar_arg.arg)
elif self.scalar_capture:		elif self.scalar_capture:
return dict(scalar_capture=self.scalar_capture.capture)		return dict(scalar_capture=self.scalar_capture.capture)
elif self.scalar_const:		elif self.scalar_const:
return dict(scalar_const=dict(type_var=self.scalar_const.type_var.name,		return dict(scalar_const=self.scalar_const.value)
attributes=[self.scalar_const.value]))
elif self.scalar_index:		elif self.scalar_index:
return dict(scalar_index=self.scalar_index.dim)		return dict(scalar_index=self.scalar_index.dim)
elif self.symbolic_cast:		elif self.symbolic_cast:
# Note that even though operands must be arity 1, we write it the		# Note that even though operands must be arity 1, we write it the
# same way as for apply because it allows handling code to be more		# same way as for apply because it allows handling code to be more
# generic vs having a special form.		# generic vs having a special form.
return dict(symbolic_cast=dict(type_var=self.symbolic_cast.to_type.name,		return dict(
		symbolic_cast=dict(
		type_var=self.symbolic_cast.to_type.name,
operands=[self.symbolic_cast.operand]))		operands=[self.symbolic_cast.operand]))
else:		else:
raise ValueError(f"Unexpected ScalarExpression type: {self}")		raise ValueError(f"Unexpected ScalarExpression type: {self}")


class ScalarAssign(YAMLObject):		class ScalarAssign(YAMLObject):
"""An assignment to a named argument (LHS of a comprehension)."""		"""An assignment to a named argument (LHS of a comprehension)."""
yaml_tag = "!ScalarAssign"		yaml_tag = "!ScalarAssign"

Show All 9 Lines

mlir/python/mlir/dialects/linalg/opdsl/ops/core_named_ops.py

	from ..lang import *			from ..lang import *

	T1 = TV.T1			T1 = TV.T1
	T2 = TV.T2			T2 = TV.T2

	Batch = S.Batch			Batch = S.Batch


	@linalg_structured_op			@linalg_structured_op
	def matmul(A=TensorDef(T1, S.M, S.K),			def matmul(
				A=TensorDef(T1, S.M, S.K),
	B=TensorDef(T2, S.K, S.N),			B=TensorDef(T2, S.K, S.N),
	C=TensorDef(U, S.M, S.N, output=True)):			C=TensorDef(U, S.M, S.N, output=True)):
	"""Performs a matrix multiplication of two 2D inputs.			"""Performs a matrix multiplication of two 2D inputs.

	Numeric casting is performed on the operands to the inner multiply, promoting			Numeric casting is performed on the operands to the inner multiply, promoting
	them to the same data type as the accumulator/output.			them to the same data type as the accumulator/output.
	"""			"""
	implements(ContractionOpInterface)			implements(ContractionOpInterface)
	C[D.m, D.n] += cast(U, A[D.m, D.k]) * cast(U, B[D.k, D.n])			C[D.m, D.n] += cast(U, A[D.m, D.k]) * cast(U, B[D.k, D.n])


	@linalg_structured_op			@linalg_structured_op
	def batch_matmul(A=TensorDef(T1, Batch, S.M, S.K),			def batch_matmul(
				A=TensorDef(T1, Batch, S.M, S.K),
	B=TensorDef(T2, Batch, S.K, S.N),			B=TensorDef(T2, Batch, S.K, S.N),
	C=TensorDef(U, Batch, S.M, S.N, output=True)):			C=TensorDef(U, Batch, S.M, S.N, output=True)):
	"""Performs a batched matrix multiplication of two 3D inputs.			"""Performs a batched matrix multiplication of two 3D inputs.

	Numeric casting is performed on the operands to the inner multiply, promoting			Numeric casting is performed on the operands to the inner multiply, promoting
	them to the same data type as the accumulator/output.			them to the same data type as the accumulator/output.
	"""			"""
	implements(ContractionOpInterface)			implements(ContractionOpInterface)
	C[D.b, D.m, D.n] += cast(U, A[D.b, D.m, D.k]) * cast(U, B[D.b, D.k, D.n])			C[D.b, D.m, D.n] += cast(U, A[D.b, D.m, D.k]) * cast(U, B[D.b, D.k, D.n])


	@linalg_structured_op			@linalg_structured_op
	def matvec(A=TensorDef(T1, S.M, S.N),			def matvec(
				A=TensorDef(T1, S.M, S.N),
	y=TensorDef(T2, S.N),			y=TensorDef(T2, S.N),
	x=TensorDef(U, S.M, output=True)):			x=TensorDef(U, S.M, output=True)):
	"""Performs a matrix-vector multiplication.			"""Performs a matrix-vector multiplication.

	Numeric casting is performed on the operands to the inner multiply, promoting			Numeric casting is performed on the operands to the inner multiply, promoting
	them to the same data type as the accumulator/output.			them to the same data type as the accumulator/output.
	"""			"""
	implements(ContractionOpInterface)			implements(ContractionOpInterface)
	x[D.m] += cast(U, A[D.m, D.n]) * cast(U, y[D.n])			x[D.m] += cast(U, A[D.m, D.n]) * cast(U, y[D.n])


	@linalg_structured_op			@linalg_structured_op
	def vecmat(y=TensorDef(T1, S.M),			def vecmat(
				y=TensorDef(T1, S.M),
	A=TensorDef(T2, S.M, S.N),			A=TensorDef(T2, S.M, S.N),
	x=TensorDef(U, S.N, output=True)):			x=TensorDef(U, S.N, output=True)):
	"""Performs a vector-matrix multiplication.			"""Performs a vector-matrix multiplication.

	Numeric casting is performed on the operands to the inner multiply, promoting			Numeric casting is performed on the operands to the inner multiply, promoting
	them to the same data type as the accumulator/output.			them to the same data type as the accumulator/output.
	"""			"""
	implements(ContractionOpInterface)			implements(ContractionOpInterface)
	x[D.n] += cast(U, y[D.m]) * cast(U, A[D.m, D.n])			x[D.n] += cast(U, y[D.m]) * cast(U, A[D.m, D.n])


	@linalg_structured_op			@linalg_structured_op
	def dot(A=TensorDef(T1, S.M), B=TensorDef(T2, S.M), C=TensorDef(U,			def dot(
	output=True)):			A=TensorDef(T1, S.M), B=TensorDef(T2, S.M), C=TensorDef(U, output=True)):
	"""Performs a dot product of two vectors to a scalar result.			"""Performs a dot product of two vectors to a scalar result.

	Numeric casting is performed on the operands to the inner multiply, promoting			Numeric casting is performed on the operands to the inner multiply, promoting
	them to the same data type as the accumulator/output.			them to the same data type as the accumulator/output.
	"""			"""
	implements(ContractionOpInterface)			implements(ContractionOpInterface)
	C[None] += cast(U, A[D.m]) * cast(U, B[D.m])			C[None] += cast(U, A[D.m]) * cast(U, B[D.m])


				@linalg_structured_op
				def fill_rng_2d(O=TensorDef(T, S.M, S.N, output=True)):
				"""Fills the output tensor with pseudo random numbers.

				The operation generations pseudo random numbers using a linear congruential
				generator. It provides no guarantees regarding the distribution of the
				generated random numbers. Instead of generating the random numbers
				sequentially, it instantiates one random number generator per data element
				and runs them in parallel. The seed operand and the indices of the data
				element seed the random number generation. The min and max operands limit
				the range of the generated random numbers.

				Note: The captures are hard-coded till there is capture support on the C++
				side.
				"""
				min = cast(F64, const(-1000))
				max = cast(F64, const(+1000))
				seed = cast(I32, const(42))
				multiplier = cast(I32, const(1103515245))
				increment = cast(I32, const(12345))
				rand1 = (cast(I32, index(D.m)) + seed) * multiplier + increment
				rand2 = (cast(I32, index(D.n)) + rand1) * multiplier + increment
				inv_range = cast(F64, const(2.3283064e-10))
				offset = cast(F64, const(2147483647))
				scaling = (max - min) * inv_range
				O[D.m, D.n] = cast(T, (offset + cast(F64, rand2)) * scaling + min)

mlir/test/CMakeLists.txt

Show First 20 Lines • Show All 59 Lines • ▼ Show 20 Lines	set(MLIR_TEST_DEPENDS
FileCheck count not		FileCheck count not
mlir-capi-execution-engine-test		mlir-capi-execution-engine-test
mlir-capi-ir-test		mlir-capi-ir-test
mlir-capi-pass-test		mlir-capi-pass-test
mlir-capi-sparse-tensor-test		mlir-capi-sparse-tensor-test
mlir-cpu-runner		mlir-cpu-runner
mlir-edsc-builder-api-test		mlir-edsc-builder-api-test
mlir-linalg-ods-gen		mlir-linalg-ods-gen
		mlir-linalg-ods-yaml-gen
mlir-lsp-server		mlir-lsp-server
mlir-opt		mlir-opt
mlir-reduce		mlir-reduce
mlir-sdbm-api-test		mlir-sdbm-api-test
mlir-tblgen		mlir-tblgen
mlir-translate		mlir-translate
mlir_runner_utils		mlir_runner_utils
mlir_c_runner_utils		mlir_c_runner_utils
▲ Show 20 Lines • Show All 58 Lines • Show Last 20 Lines

mlir/test/Dialect/Linalg/generalize-named-polymorphic-ops.mlir

	Show All 23 Lines
	// CHECK-LABEL: @generalize_matmul_tensor_i32			// CHECK-LABEL: @generalize_matmul_tensor_i32
	// CHECK: ^{{.*}}(%[[A_ARG:.+]]: i32, %[[B_ARG:.+]]: i32, %[[C_ARG:.+]]: i32)			// CHECK: ^{{.*}}(%[[A_ARG:.+]]: i32, %[[B_ARG:.+]]: i32, %[[C_ARG:.+]]: i32)
	// CHECK-NEXT: %[[MUL:.+]] = muli %[[A_ARG]], %[[B_ARG]] : i32			// CHECK-NEXT: %[[MUL:.+]] = muli %[[A_ARG]], %[[B_ARG]] : i32
	// CHECK-NEXT: %[[ADD:.+]] = addi %[[C_ARG]], %[[MUL]] : i32			// CHECK-NEXT: %[[ADD:.+]] = addi %[[C_ARG]], %[[MUL]] : i32
	// CHECK-NEXT: linalg.yield %[[ADD]] : i32			// CHECK-NEXT: linalg.yield %[[ADD]] : i32
	// CHECK-NEXT: -> tensor<16x32xi32>			// CHECK-NEXT: -> tensor<16x32xi32>

	// -----			// -----

				func @generalize_fill_rng_2d_f32(%O: tensor<16x32xf32>) -> tensor<16x32xf32> {
				%0 = linalg.fill_rng_2d outs(%O : tensor<16x32xf32>) -> tensor<16x32xf32>
				return %0: tensor<16x32xf32>
				}

				// CHECK-LABEL: @generalize_fill_rng_2d_f32
				// CHECK-SAME: (%[[O:.+]]: tensor<16x32xf32>)
				// CHECK-DAG: %[[MIN:.+]] = constant -1000 : i64
				// CHECK-DAG: %[[MAX:.+]] = constant 1000 : i64
				// CHECK-DAG: %[[SEED:.+]] = constant 42 : i32
				// CHECK-DAG: %[[IDX0:.+]] = linalg.index 0 : index
				// CHECK-DAG: %[[IDX1:.+]] = linalg.index 1 : index
				// CHECK-DAG: %[[IDX0_CAST:.+]] = index_cast %[[IDX0]] : index to i32
				// CHECK-DAG: %[[IDX1_CAST:.+]] = index_cast %[[IDX1]] : index to i32
				// CHECK-DAG: %[[VAL0:.+]] = addi %[[IDX0_CAST]], %[[SEED]] : i32
				// CHECK-DAG: %[[CST0:.+]] = constant 1103515245 : i32
				// CHECK-DAG: %[[CST1:.+]] = constant 12345 : i32
				// CHECK-DAG: %[[VAL1:.+]] = muli %[[VAL0]], %[[CST0]] : i32
				// CHECK-DAG: %[[VAL2:.+]] = addi %[[VAL1]], %[[CST1]] : i32
				// Skip random number computation for the second index.
				// CHECK-DAG: %[[MIN_CAST1:.+]] = sitofp %[[MIN]] : i64 to f64
				// CHECK-DAG: %[[MAX_CAST:.+]] = sitofp %[[MAX]] : i64 to f64
				// CHECK-DAG: %[[DIFF:.+]] = subf %[[MAX_CAST]], %[[MIN_CAST1]] : f64
				// CHECK-DAG: %[[CST2:.+]] = constant 2.3283063999999999E-10 : f64
				// CHECK-DAG: %[[FACT:.+]] = mulf %[[DIFF]], %[[CST2]] : f64
				// CHECK-DAG: %[[VAL4:.+]] = mulf %{{.+}}, %[[FACT]] : f64
				// CHECK-DAG: %[[MIN_CAST2:.+]] = sitofp %[[MIN]] : i64 to f64
				// CHECK-DAG: %[[VAL5:.+]] = addf %[[VAL4]], %[[MIN_CAST2]] : f64
				// CHECK-DAG: %[[VAL6:.+]] = fptrunc %[[VAL5]] : f64 to f32
				// CHECK-NEXT: linalg.yield %[[VAL6]] : f32
				// CHECK-NEXT: -> tensor<16x32xf32>

				// -----

				func @generalize_fill_rng_2d_i32(%O: tensor<16x32xi32>) -> tensor<16x32xi32> {
				%0 = linalg.fill_rng_2d outs(%O : tensor<16x32xi32>) -> tensor<16x32xi32>
				return %0: tensor<16x32xi32>
				}

				// CHECK-LABEL: @generalize_fill_rng_2d_i32
				// CHECK-SAME: (%[[O:.+]]: tensor<16x32xi32>)
				// Verifies floating point to integer cast.
				// CHECK: %[[VAL6:.+]] = fptosi %{{.+}} : f64 to i32
				// CHECK-NEXT: linalg.yield %[[VAL6]] : i32
				// CHECK-NEXT: -> tensor<16x32xi32>

				// -----
	// Verifies floating point to integer cast.			// Verifies floating point to integer cast.
	func @generalize_matmul_tensor_f32_f32_i16(%A : tensor<16x8xf32>, %B: tensor<8x32xf32>, %C: tensor<16x32xi16>) -> tensor<16x32xi16> {			func @generalize_matmul_tensor_f32_f32_i16(%A : tensor<16x8xf32>, %B: tensor<8x32xf32>, %C: tensor<16x32xi16>) -> tensor<16x32xi16> {
	%0 = linalg.matmul ins(%A, %B: tensor<16x8xf32>, tensor<8x32xf32>)			%0 = linalg.matmul ins(%A, %B: tensor<16x8xf32>, tensor<8x32xf32>)
	outs(%C: tensor<16x32xi16>) -> tensor<16x32xi16>			outs(%C: tensor<16x32xi16>) -> tensor<16x32xi16>
	return %0: tensor<16x32xi16>			return %0: tensor<16x32xi16>
	}			}

	// CHECK-LABEL: @generalize_matmul_tensor_f32_f32_i16			// CHECK-LABEL: @generalize_matmul_tensor_f32_f32_i16
	▲ Show 20 Lines • Show All 109 Lines • Show Last 20 Lines

mlir/test/lit.cfg.py

	Show All 15 Lines
	# Configuration file for the 'lit' test runner.			# Configuration file for the 'lit' test runner.

	# name: The name of this test suite.			# name: The name of this test suite.
	config.name = 'MLIR'			config.name = 'MLIR'

	config.test_format = lit.formats.ShTest(not llvm_config.use_lit_shell)			config.test_format = lit.formats.ShTest(not llvm_config.use_lit_shell)

	# suffixes: A list of file extensions to treat as test files.			# suffixes: A list of file extensions to treat as test files.
	config.suffixes = ['.td', '.mlir', '.toy', '.ll', '.tc', '.py', '.test']			config.suffixes = ['.td', '.mlir', '.toy', '.ll', '.tc', '.py', '.yaml', '.test']

	# test_source_root: The root path where tests are located.			# test_source_root: The root path where tests are located.
	config.test_source_root = os.path.dirname(__file__)			config.test_source_root = os.path.dirname(__file__)

	# test_exec_root: The root path where tests should be run.			# test_exec_root: The root path where tests should be run.
	config.test_exec_root = os.path.join(config.mlir_obj_root, 'test')			config.test_exec_root = os.path.join(config.mlir_obj_root, 'test')

	config.substitutions.append(('%PATH%', config.environment['PATH']))			config.substitutions.append(('%PATH%', config.environment['PATH']))
	▲ Show 20 Lines • Show All 82 Lines • Show Last 20 Lines

mlir/test/mlir-linalg-ods-gen/test-linalg-ods-yaml-gen.yaml

This file was added.

				# RUN: mlir-linalg-ods-yaml-gen %s --o-ods-decl=- \| FileCheck %s --check-prefix=ODS
				# RUN: mlir-linalg-ods-yaml-gen %s --o-impl=- \| FileCheck %s --check-prefix=IMPL

				gysitAuthorUnsubmitted Done Reply Inline Actions The dash (-) is used here to write the result to the standard output instead of writing to a file (hope this trick works on all platforms but mlir-linalg-ods-gen uses (-) as default file name so should be fine) gysit: The dash (-) is used here to write the result to the standard output instead of writing to a…
				# @linalg_structured_op
				# def test1(O=TensorDef(T, S.M, S.N, output=True)):
				# """Title.

				# Detailed description.
				# """
				# O[D.m, D.n] = cast(T, const(42)) + cast(T, index(D.n))

				--- !LinalgOpConfig
				metadata: !LinalgOpMetadata
				name: test1
				cpp_class_name: Test1Op
				doc: \|-
				Title.

				Detailed description.
				structured_op: !LinalgStructuredOpConfig
				args:
				- !<LinalgTensorDef>
				name: O
				usage: output
				shape: affine_map<()[s0, s1] -> (s0, s1)>
				element_type_var: T
				indexing_maps: !LinalgIndexingMapsConfig
				static_indexing_maps:
				- affine_map<(d0, d1)[s0, s1] -> (d0, d1)>
				iterator_types:
				- parallel
				- parallel
				assignments:
				- !ScalarAssign
				arg: O
				value: !ScalarExpression
				scalar_apply:
				fn_name: add
				operands:
				- !ScalarExpression
				symbolic_cast:
				type_var: T
				operands:
				- !ScalarExpression
				scalar_const: '42 : i64'
				- !ScalarExpression
				symbolic_cast:
				type_var: T
				operands:
				- !ScalarExpression
				scalar_index: 1

				# ODS-LABEL: def Test1Op : LinalgStructuredBase_Op<"test1"

				# ODS: let summary = [{ Title. }];
				# ODS-NEXT: let description = [{
				# ODS-NEXT: Detailed description.
				# ODS-NEXT: }];

				# ODS: let arguments =
				# ODS-NEXT: Variadic<AnyShaped>:$inputs,
				# ODS-NEXT: Variadic<AnyShaped>:$outputs

				# ODS: let builders =
				# ODS: $_state.addOperands(inputs);
				# ODS-NEXT: $_state.addOperands(outputs);
				# ODS-NEXT: $_state.addAttribute(
				# ODS-NEXT: "operand_segment_sizes",
				# ODS-NEXT: $_builder.getI32VectorAttr({
				# ODS-NEXT: static_cast<int32_t>(inputs.size()),
				# ODS-NEXT: static_cast<int32_t>(outputs.size())}));
				# ODS-NEXT: createAndFillStructuredOpRegion<Test1Op>(
				# ODS-NEXT: $_builder,
				# ODS-NEXT: $_state,
				# ODS-NEXT: TypeRange(inputs),
				# ODS-NEXT: TypeRange(outputs)

				# IMPL-LABEL: void Test1Op::regionBuilder
				# IMPL-SAME: (Block &block, ValueRange captures)
				# IMPL: Value [[VAL0:[a-z0-9]+]] = helper.constant("42 : i64");
				# IMPL-DAG: Value [[VAL1:[a-z0-9]+]] = helper.cast(block.getArgument(0).getType(), [[VAL0]]);
				# IMPL-DAG: Value [[VAL2:[a-z0-9]+]] = helper.index(1);
				# IMPL-DAG: Value [[VAL3:[a-z0-9]+]] = helper.cast(block.getArgument(0).getType(), [[VAL2]]);
				# IMPL-DAG: Value [[VAL4:[a-z0-9]+]] = helper.applyfn__add([[VAL1]], [[VAL3]]);


				# @linalg_structured_op
				# def test2(I=TensorDef(T, S.M, S.N),
				# O=TensorDef(T, S.M, S.N, output=True)):
				# """Title.

				# Detailed description.
				# """
				# O[D.m, D.n] = I[D.n, D.m]

				--- !LinalgOpConfig
				metadata: !LinalgOpMetadata
				name: test2
				cpp_class_name: Test2Op
				doc: \|-
				Title.

				Detailed description.
				structured_op: !LinalgStructuredOpConfig
				args:
				- !<LinalgTensorDef>
				name: I
				usage: input
				shape: affine_map<()[s0, s1] -> (s0, s1)>
				element_type_var: T
				- !<LinalgTensorDef>
				name: O
				usage: output
				shape: affine_map<()[s0, s1] -> (s0, s1)>
				element_type_var: T
				indexing_maps: !LinalgIndexingMapsConfig
				static_indexing_maps:
				- affine_map<(d0, d1)[s0, s1] -> (d1, d0)>
				- affine_map<(d0, d1)[s0, s1] -> (d0, d1)>
				iterator_types:
				- parallel
				- parallel
				assignments:
				- !ScalarAssign
				arg: O
				value: !ScalarExpression
				scalar_arg: I

				# IMPL-LABEL: Test2Op::iterator_types()
				# IMPL-NEXT: { getParallelIteratorTypeName(), getParallelIteratorTypeName() }

				# IMPL: Test2Op::indexing_maps()
				# IMPL: "affine_map<(d0, d1)[s0, s1] -> (d1, d0)>"
				# IMPL: "affine_map<(d0, d1)[s0, s1] -> (d0, d1)>"

				# IMPL: void Test2Op::regionBuilder(Block &block, ValueRange captures)
				# IMPL: yields.push_back(block.getArgument(0));

mlir/test/python/dialects/linalg/opdsl/arguments.py

This file was added.

				# RUN: %PYTHON -m mlir.dialects.linalg.opdsl.dump_oplib --file %s \| FileCheck %s

				from mlir.dialects.linalg.opdsl.lang import *


				# CHECK: ---
				# CHECK-LABEL: matmul
				# CHECK: args:
				# CHECK: name: A
				# CHECK: usage: input
				# CHECK: shape: affine_map<()[s0, s1, s2] -> (s0, s2)>
				# CHECK: element_type_var: T
				# CHECK: name: B
				# CHECK: usage: input
				# CHECK: shape: affine_map<()[s0, s1, s2] -> (s2, s1)>
				# CHECK: element_type_var: T
				# CHECK: name: C
				# CHECK: usage: output
				# CHECK: shape: affine_map<()[s0, s1, s2] -> (s0, s1)>
				# CHECK: element_type_var: U
				@linalg_structured_op
				def matmul(
				A=TensorDef(T, S.M, S.K),
				B=TensorDef(T, S.K, S.N),
				C=TensorDef(U, S.M, S.N, output=True)):
				C[D.m, D.n] += cast(U, A[D.m, D.k]) * cast(U, B[D.k, D.n])


				# CHECK: ---
				# CHECK-LABEL: fill
				# CHECK: captures:
				# CHECK: - !<LinalgCaptureDef>
				# CHECK: name: value
				# CHECK: type_var: T
				@linalg_structured_op
				def fill(O=TensorDef(T, S.M, S.K, output=True), value=CaptureDef(T)):
				O[D.m, D.n] = value

mlir/test/python/dialects/linalg/opdsl/assignments.py

	# RUN: %PYTHON -m mlir.dialects.linalg.opdsl.dump_oplib --file %s \| FileCheck %s			# RUN: %PYTHON -m mlir.dialects.linalg.opdsl.dump_oplib --file %s \| FileCheck %s

	from mlir.dialects.linalg.opdsl.lang import *			from mlir.dialects.linalg.opdsl.lang import *


	# CHECK: ---			# CHECK: ---
	# CHECK-LABEL: matmul			# CHECK-LABEL: matmul
	# CHECK: assignments:			# CHECK: assignments:
	# CHECK: -			# CHECK: -
	# CHECK: arg: C			# CHECK: arg: C
	# CHECK: value:			# CHECK: value:
	# CHECK: scalar_apply:			# CHECK: scalar_apply:
	# CHECK: fn_name: add			# CHECK: fn_name: add
	# CHECK: operands:			# CHECK: operands:
	# CHECK: scalar_apply:			# CHECK: scalar_apply:
	# CHECK: fn_name: mul			# CHECK: fn_name: mul
	# CHECK: operands:			# CHECK: operands:
	# CHECK: symbolic_cast:			# CHECK: symbolic_cast:
	# CHECK: type_var: U			# CHECK: type_var: U
	# CHECK: operands:			# CHECK: operands:
	# CHECK: scalar_arg: A			# CHECK: scalar_arg: A
	# CHECK: symbolic_cast:			# CHECK: symbolic_cast:
	# CHECK: type_var: U			# CHECK: type_var: U
	# CHECK: operands:			# CHECK: operands:
	# CHECK: scalar_arg: B			# CHECK: scalar_arg: B
	@linalg_structured_op			@linalg_structured_op
	def matmul(A=TensorDef(T, S.M, S.K),			def matmul(
				A=TensorDef(T, S.M, S.K),
	B=TensorDef(T, S.K, S.N),			B=TensorDef(T, S.K, S.N),
	C=TensorDef(U, S.M, S.N, output=True)):			C=TensorDef(U, S.M, S.N, output=True)):
	C[D.m, D.n] += cast(U, A[D.m, D.k]) * cast(U, B[D.k, D.n])			C[D.m, D.n] += cast(U, A[D.m, D.k]) * cast(U, B[D.k, D.n])


				# CHECK: ---
				# CHECK-LABEL: constants
				# CHECK: assignments:
				# CHECK: -
				# CHECK: arg: O
				# CHECK: scalar_apply:
				# CHECK: fn_name: sub
				# CHECK: operands:
				# CHECK: scalar_apply:
				# CHECK: fn_name: add
				# CHECK: operands:
				# CHECK: symbolic_cast:
				# CHECK: type_var: T
				# CHECK: operands:
				# CHECK: scalar_const: '3.1415926535897931 : f64'
				# CHECK: symbolic_cast:
				# CHECK: type_var: T
				# CHECK: operands:
				# CHECK: scalar_const: '42 : i64'
				# CHECK: symbolic_cast:
				# CHECK: type_var: T
				# CHECK: operands:
				# CHECK: scalar_const: '1.{{[0]*}}e+03 : f64'
				@linalg_structured_op
				def constants(O=TensorDef(T, S.M, S.K, output=True)):
				pi = cast(T, const(3.1415926535897931))
				cst42 = cast(T, const(42))
				cst1000 = cast(T, const(1e+3))
				O[D.m, D.n] = pi + cst42 - cst1000


				# CHECK: ---
				# CHECK-LABEL: indices
				# CHECK: assignments:
				# CHECK: -
				# CHECK: arg: O
				# CHECK: scalar_apply:
				# CHECK: fn_name: add
				# CHECK: operands:
				# CHECK: scalar_index: 1
				# CHECK: scalar_index: 0
				@linalg_structured_op
				def indices(O=TensorDef(T, S.M, S.K, output=True)):
				O[D.m, D.n] = index(D.n) + index(D.m)


				# CHECK: ---
				# CHECK-LABEL: fill
				# CHECK: assignments:
				# CHECK: -
				# CHECK: arg: O
				# CHECK: scalar_capture: value
				@linalg_structured_op
				def fill(O=TensorDef(T, S.M, S.K, output=True), value=CaptureDef(T)):
				O[D.m, D.n] = value
				nicolasvasilacheUnsubmitted Done Reply Inline Actions nit: NL nicolasvasilache: nit: NL

mlir/test/python/dialects/linalg/opdsl/emit_structured_generic.py

# RUN: %PYTHON %s \| FileCheck %s		# RUN: %PYTHON %s \| FileCheck %s

from typing import Optional, Sequence

from mlir.ir import *		from mlir.ir import *
from mlir.dialects import builtin		from mlir.dialects import builtin
from mlir.dialects import linalg		from mlir.dialects import linalg
from mlir.dialects import std		from mlir.dialects import std

from mlir.dialects.linalg.opdsl.lang import *		from mlir.dialects.linalg.opdsl.lang import *


@linalg_structured_op		@linalg_structured_op
def matmul_mono(A=TensorDef(T, S.M, S.K),		def matmul_mono(
		A=TensorDef(T, S.M, S.K),
B=TensorDef(T, S.K, S.N),		B=TensorDef(T, S.K, S.N),
C=TensorDef(T, S.M, S.N, output=True)):		C=TensorDef(T, S.M, S.N, output=True)):
C[D.m, D.n] += A[D.m, D.k] * B[D.k, D.n]		C[D.m, D.n] += A[D.m, D.k] * B[D.k, D.n]


@linalg_structured_op		@linalg_structured_op
def matmul_poly(A=TensorDef(TV.T1, S.M, S.K),		def matmul_poly(
		A=TensorDef(TV.T1, S.M, S.K),
B=TensorDef(TV.T2, S.K, S.N),		B=TensorDef(TV.T2, S.K, S.N),
C=TensorDef(U, S.M, S.N, output=True)):		C=TensorDef(U, S.M, S.N, output=True)):
C[D.m, D.n] += cast(U, A[D.m, D.k]) * cast(U, B[D.k, D.n])		C[D.m, D.n] += cast(U, A[D.m, D.k]) * cast(U, B[D.k, D.n])


@linalg_structured_op		@linalg_structured_op
def fill_rng_2d(A=TensorDef(T, S.M, S.N, output=True),		def fill_rng(
		O=TensorDef(T, S.M, S.N, output=True),
min=CaptureDef(F64),		min=CaptureDef(F64),
max=CaptureDef(F64),		max=CaptureDef(F64),
seed=CaptureDef(I32)):		seed=CaptureDef(I32)):
multiplier = const(I32, 1103515245)		multiplier = cast(I32, const(1103515245))
increment = const(I32, 12345)		increment = cast(I32, const(12345))
temp1 = (cast(I32, index(D.m)) + seed) * multiplier + increment		rand1 = (cast(I32, index(D.m)) + seed) * multiplier + increment
temp2 = (cast(I32, index(D.n)) + temp1) * multiplier + increment		rand2 = (cast(I32, index(D.n)) + rand1) * multiplier + increment
inv_randmax = const(F64, 2.3283064e-10)		inv_range = cast(F64, const(2.3283064e-10))
scaling = (max - min) * inv_randmax		offset = cast(F64, const(2147483647))
A[D.m, D.n] = cast(T, cast(F64, temp2) * scaling + min)		scaling = (max - min) * inv_range
		O[D.m, D.n] = cast(T, (offset + cast(F64, rand2)) * scaling + min)


with Context() as ctx, Location.unknown():		with Context() as ctx, Location.unknown():
module = Module.create()		module = Module.create()
f16 = F16Type.get()		f16 = F16Type.get()
f32 = F32Type.get()		f32 = F32Type.get()
f64 = F64Type.get()		f64 = F64Type.get()
i8 = IntegerType.get_signless(8)		i8 = IntegerType.get_signless(8)
i16 = IntegerType.get_signless(16)		i16 = IntegerType.get_signless(16)
Show All 13 Lines	with InsertionPoint(module.body):

# CHECK: %[[INITC:.+]] = linalg.init_tensor [4, 8] : tensor<4x8xf32>		# CHECK: %[[INITC:.+]] = linalg.init_tensor [4, 8] : tensor<4x8xf32>
# CHECK: linalg.generic		# CHECK: linalg.generic
# CHECK-SAME: indexing_maps = [#[[$MAPA]], #[[$MAPB]], #[[$MAPC]]]		# CHECK-SAME: indexing_maps = [#[[$MAPA]], #[[$MAPB]], #[[$MAPC]]]
# CHECK-SAME: iterator_types = ["parallel", "parallel", "reduction"]		# CHECK-SAME: iterator_types = ["parallel", "parallel", "reduction"]
# CHECK-SAME: ins(%[[A]], %[[B]]		# CHECK-SAME: ins(%[[A]], %[[B]]
# CHECK-SAME: outs(%[[INITC]]		# CHECK-SAME: outs(%[[INITC]]

@builtin.FuncOp.from_py_func(RankedTensorType.get((4, 16), f32),		@builtin.FuncOp.from_py_func(
RankedTensorType.get((16, 8), f32))		RankedTensorType.get((4, 16), f32), RankedTensorType.get((16, 8), f32))
def test_matmul_mono(lhs, rhs):		def test_matmul_mono(lhs, rhs):
init_result = linalg.InitTensorOp([4, 8], f32)		init_result = linalg.InitTensorOp([4, 8], f32)
return matmul_mono(lhs, rhs, outs=[init_result.result])		return matmul_mono(lhs, rhs, outs=[init_result.result])

# CHECK-LABEL: @test_i8i8i32_matmul		# CHECK-LABEL: @test_i8i8i32_matmul
# CHECK: ^{{.*}}(%[[A_ARG:.+]]: i8, %[[B_ARG:.+]]: i8, %[[C_ARG:.+]]: i32)		# CHECK: ^{{.*}}(%[[A_ARG:.+]]: i8, %[[B_ARG:.+]]: i8, %[[C_ARG:.+]]: i32)
# CHECK-NEXT: %[[A_CAST:.+]] = sexti %[[A_ARG]] : i8 to i32		# CHECK-NEXT: %[[A_CAST:.+]] = sexti %[[A_ARG]] : i8 to i32
# CHECK-NEXT: %[[B_CAST:.+]] = sexti %[[B_ARG]] : i8 to i32		# CHECK-NEXT: %[[B_CAST:.+]] = sexti %[[B_ARG]] : i8 to i32
# CHECK-NEXT: %[[MUL:.+]] = muli %[[A_CAST]], %[[B_CAST]] : i32		# CHECK-NEXT: %[[MUL:.+]] = muli %[[A_CAST]], %[[B_CAST]] : i32
# CHECK-NEXT: %[[ADD:.+]] = addi %[[C_ARG]], %[[MUL]] : i32		# CHECK-NEXT: %[[ADD:.+]] = addi %[[C_ARG]], %[[MUL]] : i32
# CHECK-NEXT: linalg.yield %[[ADD]] : i32		# CHECK-NEXT: linalg.yield %[[ADD]] : i32
# CHECK-NEXT: -> tensor<4x8xi32>		# CHECK-NEXT: -> tensor<4x8xi32>
@builtin.FuncOp.from_py_func(RankedTensorType.get((4, 16), i8),		@builtin.FuncOp.from_py_func(
RankedTensorType.get((16, 8), i8),		RankedTensorType.get((4, 16), i8), RankedTensorType.get((16, 8), i8),
RankedTensorType.get((4, 8), i32))		RankedTensorType.get((4, 8), i32))
def test_i8i8i32_matmul(lhs, rhs, init_result):		def test_i8i8i32_matmul(lhs, rhs, init_result):
return matmul_poly(lhs, rhs, outs=[init_result])		return matmul_poly(lhs, rhs, outs=[init_result])

# CHECK-LABEL: @test_i8i16i32_matmul		# CHECK-LABEL: @test_i8i16i32_matmul
# CHECK: ^{{.*}}(%[[A_ARG:.+]]: i8, %[[B_ARG:.+]]: i16, %[[C_ARG:.+]]: i32)		# CHECK: ^{{.*}}(%[[A_ARG:.+]]: i8, %[[B_ARG:.+]]: i16, %[[C_ARG:.+]]: i32)
# CHECK-NEXT: %[[A_CAST:.+]] = sexti %[[A_ARG]] : i8 to i32		# CHECK-NEXT: %[[A_CAST:.+]] = sexti %[[A_ARG]] : i8 to i32
# CHECK-NEXT: %[[B_CAST:.+]] = sexti %[[B_ARG]] : i16 to i32		# CHECK-NEXT: %[[B_CAST:.+]] = sexti %[[B_ARG]] : i16 to i32
# CHECK-NEXT: %[[MUL:.+]] = muli %[[A_CAST]], %[[B_CAST]] : i32		# CHECK-NEXT: %[[MUL:.+]] = muli %[[A_CAST]], %[[B_CAST]] : i32
# CHECK-NEXT: %[[ADD:.+]] = addi %[[C_ARG]], %[[MUL]] : i32		# CHECK-NEXT: %[[ADD:.+]] = addi %[[C_ARG]], %[[MUL]] : i32
# CHECK-NEXT: linalg.yield %[[ADD]] : i32		# CHECK-NEXT: linalg.yield %[[ADD]] : i32
# CHECK-NEXT: -> tensor<4x8xi32>		# CHECK-NEXT: -> tensor<4x8xi32>
@builtin.FuncOp.from_py_func(RankedTensorType.get((4, 16), i8),		@builtin.FuncOp.from_py_func(
RankedTensorType.get((16, 8), i16),		RankedTensorType.get((4, 16), i8), RankedTensorType.get((16, 8), i16),
RankedTensorType.get((4, 8), i32))		RankedTensorType.get((4, 8), i32))
def test_i8i16i32_matmul(lhs, rhs, init_result):		def test_i8i16i32_matmul(lhs, rhs, init_result):
return matmul_poly(lhs, rhs, outs=[init_result])		return matmul_poly(lhs, rhs, outs=[init_result])

# CHECK-LABEL: @test_i32i32i16_matmul		# CHECK-LABEL: @test_i32i32i16_matmul
# CHECK: ^{{.*}}(%[[A_ARG:.+]]: i32, %[[B_ARG:.+]]: i32, %[[C_ARG:.+]]: i16)		# CHECK: ^{{.*}}(%[[A_ARG:.+]]: i32, %[[B_ARG:.+]]: i32, %[[C_ARG:.+]]: i16)
# CHECK-NEXT: %[[A_CAST:.+]] = trunci %[[A_ARG]] : i32 to i16		# CHECK-NEXT: %[[A_CAST:.+]] = trunci %[[A_ARG]] : i32 to i16
# CHECK-NEXT: %[[B_CAST:.+]] = trunci %[[B_ARG]] : i32 to i16		# CHECK-NEXT: %[[B_CAST:.+]] = trunci %[[B_ARG]] : i32 to i16
# CHECK-NEXT: %[[MUL:.+]] = muli %[[A_CAST]], %[[B_CAST]] : i16		# CHECK-NEXT: %[[MUL:.+]] = muli %[[A_CAST]], %[[B_CAST]] : i16
# CHECK-NEXT: %[[ADD:.+]] = addi %[[C_ARG]], %[[MUL]] : i16		# CHECK-NEXT: %[[ADD:.+]] = addi %[[C_ARG]], %[[MUL]] : i16
# CHECK-NEXT: linalg.yield %[[ADD]] : i16		# CHECK-NEXT: linalg.yield %[[ADD]] : i16
# CHECK-NEXT: -> tensor<4x8xi16>		# CHECK-NEXT: -> tensor<4x8xi16>
@builtin.FuncOp.from_py_func(RankedTensorType.get((4, 16), i32),		@builtin.FuncOp.from_py_func(
RankedTensorType.get((16, 8), i32),		RankedTensorType.get((4, 16), i32), RankedTensorType.get((16, 8), i32),
RankedTensorType.get((4, 8), i16))		RankedTensorType.get((4, 8), i16))
def test_i32i32i16_matmul(lhs, rhs, init_result):		def test_i32i32i16_matmul(lhs, rhs, init_result):
return matmul_poly(lhs, rhs, outs=[init_result])		return matmul_poly(lhs, rhs, outs=[init_result])

# CHECK-LABEL: @test_i8i8f32_matmul		# CHECK-LABEL: @test_i8i8f32_matmul
# CHECK: ^{{.*}}(%[[A_ARG:.+]]: i8, %[[B_ARG:.+]]: i8, %[[C_ARG:.+]]: f32)		# CHECK: ^{{.*}}(%[[A_ARG:.+]]: i8, %[[B_ARG:.+]]: i8, %[[C_ARG:.+]]: f32)
# CHECK-NEXT: %[[A_CAST:.+]] = sitofp %[[A_ARG]] : i8 to f32		# CHECK-NEXT: %[[A_CAST:.+]] = sitofp %[[A_ARG]] : i8 to f32
# CHECK-NEXT: %[[B_CAST:.+]] = sitofp %[[B_ARG]] : i8 to f32		# CHECK-NEXT: %[[B_CAST:.+]] = sitofp %[[B_ARG]] : i8 to f32
# CHECK-NEXT: %[[MUL:.+]] = mulf %[[A_CAST]], %[[B_CAST]] : f32		# CHECK-NEXT: %[[MUL:.+]] = mulf %[[A_CAST]], %[[B_CAST]] : f32
# CHECK-NEXT: %[[ADD:.+]] = addf %[[C_ARG]], %[[MUL]] : f32		# CHECK-NEXT: %[[ADD:.+]] = addf %[[C_ARG]], %[[MUL]] : f32
# CHECK-NEXT: linalg.yield %[[ADD]] : f32		# CHECK-NEXT: linalg.yield %[[ADD]] : f32
# CHECK-NEXT: -> tensor<4x8xf32>		# CHECK-NEXT: -> tensor<4x8xf32>
@builtin.FuncOp.from_py_func(RankedTensorType.get((4, 16), i8),		@builtin.FuncOp.from_py_func(
RankedTensorType.get((16, 8), i8),		RankedTensorType.get((4, 16), i8), RankedTensorType.get((16, 8), i8),
RankedTensorType.get((4, 8), f32))		RankedTensorType.get((4, 8), f32))
def test_i8i8f32_matmul(lhs, rhs, init_result):		def test_i8i8f32_matmul(lhs, rhs, init_result):
return matmul_poly(lhs, rhs, outs=[init_result])		return matmul_poly(lhs, rhs, outs=[init_result])

# CHECK-LABEL: @test_f16f16f32_matmul		# CHECK-LABEL: @test_f16f16f32_matmul
# CHECK: ^{{.*}}(%[[A_ARG:.+]]: f16, %[[B_ARG:.+]]: f16, %[[C_ARG:.+]]: f32)		# CHECK: ^{{.*}}(%[[A_ARG:.+]]: f16, %[[B_ARG:.+]]: f16, %[[C_ARG:.+]]: f32)
# CHECK-NEXT: %[[A_CAST:.+]] = fpext %[[A_ARG]] : f16 to f32		# CHECK-NEXT: %[[A_CAST:.+]] = fpext %[[A_ARG]] : f16 to f32
# CHECK-NEXT: %[[B_CAST:.+]] = fpext %[[B_ARG]] : f16 to f32		# CHECK-NEXT: %[[B_CAST:.+]] = fpext %[[B_ARG]] : f16 to f32
# CHECK-NEXT: %[[MUL:.+]] = mulf %[[A_CAST]], %[[B_CAST]] : f32		# CHECK-NEXT: %[[MUL:.+]] = mulf %[[A_CAST]], %[[B_CAST]] : f32
# CHECK-NEXT: %[[ADD:.+]] = addf %[[C_ARG]], %[[MUL]] : f32		# CHECK-NEXT: %[[ADD:.+]] = addf %[[C_ARG]], %[[MUL]] : f32
# CHECK-NEXT: linalg.yield %[[ADD]] : f32		# CHECK-NEXT: linalg.yield %[[ADD]] : f32
# CHECK-NEXT: -> tensor<4x8xf32>		# CHECK-NEXT: -> tensor<4x8xf32>
@builtin.FuncOp.from_py_func(RankedTensorType.get((4, 16), f16),		@builtin.FuncOp.from_py_func(
RankedTensorType.get((16, 8), f16),		RankedTensorType.get((4, 16), f16), RankedTensorType.get((16, 8), f16),
RankedTensorType.get((4, 8), f32))		RankedTensorType.get((4, 8), f32))
def test_f16f16f32_matmul(lhs, rhs, init_result):		def test_f16f16f32_matmul(lhs, rhs, init_result):
return matmul_poly(lhs, rhs, outs=[init_result])		return matmul_poly(lhs, rhs, outs=[init_result])

# CHECK-LABEL: @test_f64f64f32_matmul		# CHECK-LABEL: @test_f64f64f32_matmul
# CHECK: ^{{.*}}(%[[A_ARG:.+]]: f64, %[[B_ARG:.+]]: f64, %[[C_ARG:.+]]: f32)		# CHECK: ^{{.*}}(%[[A_ARG:.+]]: f64, %[[B_ARG:.+]]: f64, %[[C_ARG:.+]]: f32)
# CHECK-NEXT: %[[A_CAST:.+]] = fptrunc %[[A_ARG]] : f64 to f32		# CHECK-NEXT: %[[A_CAST:.+]] = fptrunc %[[A_ARG]] : f64 to f32
# CHECK-NEXT: %[[B_CAST:.+]] = fptrunc %[[B_ARG]] : f64 to f32		# CHECK-NEXT: %[[B_CAST:.+]] = fptrunc %[[B_ARG]] : f64 to f32
# CHECK-NEXT: %[[MUL:.+]] = mulf %[[A_CAST]], %[[B_CAST]] : f32		# CHECK-NEXT: %[[MUL:.+]] = mulf %[[A_CAST]], %[[B_CAST]] : f32
# CHECK-NEXT: %[[ADD:.+]] = addf %[[C_ARG]], %[[MUL]] : f32		# CHECK-NEXT: %[[ADD:.+]] = addf %[[C_ARG]], %[[MUL]] : f32
# CHECK-NEXT: linalg.yield %[[ADD]] : f32		# CHECK-NEXT: linalg.yield %[[ADD]] : f32
# CHECK-NEXT: -> tensor<4x8xf32>		# CHECK-NEXT: -> tensor<4x8xf32>
@builtin.FuncOp.from_py_func(RankedTensorType.get((4, 16), f64),		@builtin.FuncOp.from_py_func(
RankedTensorType.get((16, 8), f64),		RankedTensorType.get((4, 16), f64), RankedTensorType.get((16, 8), f64),
RankedTensorType.get((4, 8), f32))		RankedTensorType.get((4, 8), f32))
def test_f64f64f32_matmul(lhs, rhs, init_result):		def test_f64f64f32_matmul(lhs, rhs, init_result):
return matmul_poly(lhs, rhs, outs=[init_result])		return matmul_poly(lhs, rhs, outs=[init_result])

# CHECK-LABEL: @test_fill_rng_2d		# CHECK-LABEL: @test_fill_rng
# CHECK-SAME: %{{.*}} tensor<4x16xi32>, %[[MIN:.+]]: f64, %[[MAX:.+]]: f64, %[[SEED:.+]]: i32		# CHECK-SAME: %{{.*}} tensor<4x16xi32>, %[[MIN:.+]]: f64, %[[MAX:.+]]: f64, %[[SEED:.+]]: i32
# CHECK-DAG: %[[IDX0:.+]] = linalg.index 0 : index		# CHECK-DAG: %[[IDX0:.+]] = linalg.index 0 : index
# CHECK-DAG: %[[IDX1:.+]] = linalg.index 1 : index		# CHECK-DAG: %[[IDX1:.+]] = linalg.index 1 : index
# CHECK-DAG: %[[IDX0_CAST:.+]] = index_cast %[[IDX0]] : index to i32		# CHECK-DAG: %[[IDX0_CAST:.+]] = index_cast %[[IDX0]] : index to i32
# CHECK-DAG: %[[IDX1_CAST:.+]] = index_cast %[[IDX1]] : index to i32		# CHECK-DAG: %[[IDX1_CAST:.+]] = index_cast %[[IDX1]] : index to i32
# CHECK-DAG: %[[RND0:.+]] = addi %[[IDX0_CAST]], %[[SEED]] : i32		# CHECK-DAG: %[[RND0:.+]] = addi %[[IDX0_CAST]], %[[SEED]] : i32
# CHECK-DAG: %[[CST0:.+]] = constant 1103515245 : i32		# CHECK-DAG: %[[CST0:.+]] = constant 1103515245 : i64
# CHECK-DAG: %[[CST1:.+]] = constant 12345 : i32		# CHECK-DAG: %[[CST0_CAST:.+]] = trunci %[[CST0]] : i64 to i32
# CHECK-DAG: %[[RND1:.+]] = muli %[[RND0]], %[[CST0]] : i32		# CHECK-DAG: %[[CST1:.+]] = constant 12345 : i64
# CHECK-DAG: %[[RND2:.+]] = addi %[[RND1]], %[[CST1]] : i32		# CHECK-DAG: %[[CST1_CAST:.+]] = trunci %[[CST1]] : i64 to i32
# CHECK: %[[RND3:.+]] = sitofp %{{.*}} : i32 to f64		# CHECK-DAG: %[[RND1:.+]] = muli %[[RND0]], %[[CST0_CAST]] : i32
		# CHECK-DAG: %[[RND2:.+]] = addi %[[RND1]], %[[CST1_CAST]] : i32
		# Skip random number computation for the second index.
# CHECK-DAG: %[[DIFF:.+]] = subf %[[MAX]], %[[MIN]] : f64		# CHECK-DAG: %[[DIFF:.+]] = subf %[[MAX]], %[[MIN]] : f64
# CHECK-DAG: %[[CST2:.+]] = constant 2.3283063999999999E-10 : f64		# CHECK-DAG: %[[CST3:.+]] = constant 2.3283063999999999E-10 : f64
# CHECK-DAG: %[[FACT:.+]] = mulf %[[DIFF]], %[[CST2]] : f64		# CHECK-DAG: %[[FACT:.+]] = mulf %[[DIFF]], %[[CST3]] : f64
# CHECK-DAG: %[[RND4:.+]] = mulf %[[RND3]], %[[FACT]] : f64		# CHECK-DAG: %[[RND4:.+]] = mulf %{{.+}}, %[[FACT]] : f64
# CHECK-DAG: %[[RND5:.+]] = addf %[[RND4]], %[[MIN]] : f64		# CHECK-DAG: %[[RND5:.+]] = addf %[[RND4]], %[[MIN]] : f64
# CHECK-DAG: %{{.*}} = fptosi %[[RND5]] : f64 to i32		# CHECK-DAG: %{{.*}} = fptosi %[[RND5]] : f64 to i32
@builtin.FuncOp.from_py_func(RankedTensorType.get((4, 16), i32),		@builtin.FuncOp.from_py_func(
f64, f64, i32)		RankedTensorType.get((4, 16), i32), f64, f64, i32)
def test_fill_rng_2d(init_result, min, max, seed):		def test_fill_rng(init_result, min, max, seed):
return fill_rng_2d(outs=[init_result], captures=[min, max, seed])		return fill_rng(outs=[init_result], captures=[min, max, seed])


print(module)		print(module)

mlir/test/python/dialects/linalg/opsrun.py

	# RUN: %PYTHON %s 2>&1 \| FileCheck %s			# RUN: %PYTHON %s 2>&1 \| FileCheck %s

	import sys			import sys
				stellaraccidentUnsubmitted Not Done Reply Inline Actions Can be done in a follow-up, but I feel like this test is firmly in integration test territory and should probably be moved to the integration test directory. stellaraccident: Can be done in a follow-up, but I feel like this test is firmly in integration test territory…
				gysitAuthorUnsubmitted Done Reply Inline Actions yes lets move this to a follow-up iteration gysit: yes lets move this to a follow-up iteration
	from mlir.ir import *			from mlir.ir import *
	from mlir.dialects import builtin			from mlir.dialects import builtin
	from mlir.dialects import linalg			from mlir.dialects import linalg
	from mlir.dialects import std			from mlir.dialects import std
	from mlir.passmanager import *			from mlir.passmanager import *
	from mlir.execution_engine import *			from mlir.execution_engine import *


	# Log everything to stderr and flush so that we have a unified stream to match			# Log everything to stderr and flush so that we have a unified stream to match
	# errors/info emitted by MLIR to stderr.			# errors/info emitted by MLIR to stderr.
	def log(*args):			def log(*args):
	print(*args, file=sys.stderr)			print(*args, file=sys.stderr)
	sys.stderr.flush()			sys.stderr.flush()

	boilerplate = """
				matmul_boiler = """
	func @main() -> f32 attributes {llvm.emit_c_interface} {			func @main() -> f32 attributes {llvm.emit_c_interface} {
	%v0 = constant 0.0 : f32			%v0 = constant 0.0 : f32
	%v1 = constant 1.0 : f32			%v1 = constant 1.0 : f32
	%v2 = constant 2.0 : f32			%v2 = constant 2.0 : f32

	%A = memref.alloc() : memref<4x16xf32>			%A = memref.alloc() : memref<4x16xf32>
	%B = memref.alloc() : memref<16x8xf32>			%B = memref.alloc() : memref<16x8xf32>
	%C = memref.alloc() : memref<4x8xf32>			%C = memref.alloc() : memref<4x8xf32>
	linalg.fill(%A, %v1) : memref<4x16xf32>, f32			linalg.fill(%A, %v1) : memref<4x16xf32>, f32
	linalg.fill(%B, %v2) : memref<16x8xf32>, f32			linalg.fill(%B, %v2) : memref<16x8xf32>, f32
	linalg.fill(%C, %v0) : memref<4x8xf32>, f32			linalg.fill(%C, %v0) : memref<4x8xf32>, f32

	call @matmul_on_buffers(%A, %B, %C) :			call @matmul_on_buffers(%A, %B, %C) :
	(memref<4x16xf32>, memref<16x8xf32>, memref<4x8xf32>) -> ()			(memref<4x16xf32>, memref<16x8xf32>, memref<4x8xf32>) -> ()

	%c0 = constant 0 : index			%c0 = constant 0 : index
	%0 = memref.load %C[%c0, %c0] : memref<4x8xf32>			%0 = memref.load %C[%c0, %c0] : memref<4x8xf32>

	// TODO: FFI-based solution to allow testing and printing with python code.			// TODO: FFI-based solution to allow testing and printing with python code.
	return %0 : f32			return %0 : f32
	}			}
	"""			"""

	def transform(module):			fill_boiler = """
				func @main() -> i32 attributes {llvm.emit_c_interface} {
				%O = memref.alloc() : memref<4x16xi32>

				call @fill_on_buffers(%O) :
				(memref<4x16xi32>) -> ()

				%c0 = constant 0 : index
				%0 = memref.load %O[%c0, %c0] : memref<4x16xi32>

				// TODO: FFI-based solution to allow testing and printing with python code.
				return %0 : i32
				}
				"""


				def transform(module, boilerplate):
	import mlir.conversions			import mlir.conversions
	import mlir.dialects.linalg.passes			import mlir.dialects.linalg.passes
	import mlir.transforms			import mlir.transforms

	# TODO: Allow cloning functions from one module to another.			# TODO: Allow cloning functions from one module to another.
	# Atm we have to resort to string concatenation.			# Atm we have to resort to string concatenation.
	mod = Module.parse(			mod = Module.parse(
	str(module.operation.regions[0].blocks[0].operations[0].operation) +			str(module.operation.regions[0].blocks[0].operations[0].operation) +
	boilerplate)			boilerplate)
	pm = PassManager.parse("func(convert-linalg-to-loops, convert-scf-to-std)," +			pm = PassManager.parse("func(convert-linalg-to-loops, convert-scf-to-std)," +
	"convert-vector-to-llvm," +			"convert-vector-to-llvm," + "convert-std-to-llvm")
	"convert-std-to-llvm")
	pm.run(mod)			pm.run(mod)
	return mod			return mod

	def test_builtin():
				def test_matmul_builtin():
	with Context() as ctx, Location.unknown():			with Context() as ctx, Location.unknown():
	module = Module.create()			module = Module.create()
	f32 = F32Type.get()			f32 = F32Type.get()
	with InsertionPoint(module.body):			with InsertionPoint(module.body):
	@builtin.FuncOp.from_py_func(MemRefType.get((4, 16), f32),
	MemRefType.get((16, 8), f32),			@builtin.FuncOp.from_py_func(
				MemRefType.get((4, 16), f32), MemRefType.get((16, 8), f32),
	MemRefType.get((4, 8), f32))			MemRefType.get((4, 8), f32))
	def matmul_on_buffers(lhs, rhs, out):			def matmul_on_buffers(lhs, rhs, out):
	linalg.matmul(lhs, rhs, outs=[out])			linalg.matmul(lhs, rhs, outs=[out])

	execution_engine = ExecutionEngine(transform(module))			execution_engine = ExecutionEngine(transform(module, matmul_boiler))

	# TODO: FFI-based solution to allow testing and printing with python code.			# TODO: FFI-based solution to allow testing and printing with python code.
	# Prepare arguments: one result f32.			# Prepare arguments: one result f32.
	# Arguments must be passed as pointers.			# Arguments must be passed as pointers.
	c_float_p = ctypes.c_float * 1			c_float_p = ctypes.c_float * 1
	res = c_float_p(-1.)			res = c_float_p(-1.)
	execution_engine.invoke("main", res)			execution_engine.invoke("main", res)

	log('RESULT: ', res[0])			log("RESULT: ", res[0])
	# CHECK: RESULT: 32.0			# CHECK: RESULT: 32.0

	test_builtin()

	def test_generic():			test_matmul_builtin()


				def test_matmul_generic():
	with Context() as ctx, Location.unknown():			with Context() as ctx, Location.unknown():
	module = Module.create()			module = Module.create()
	f32 = F32Type.get()			f32 = F32Type.get()
	with InsertionPoint(module.body):			with InsertionPoint(module.body):
	@builtin.FuncOp.from_py_func(MemRefType.get((4, 16), f32),
	MemRefType.get((16, 8), f32),			@builtin.FuncOp.from_py_func(
				MemRefType.get((4, 16), f32), MemRefType.get((16, 8), f32),
	MemRefType.get((4, 8), f32))			MemRefType.get((4, 8), f32))
	def matmul_on_buffers(lhs, rhs, out):			def matmul_on_buffers(lhs, rhs, out):
	linalg.matmul(lhs, rhs, outs=[out], emit_generic=True)			linalg.matmul(lhs, rhs, outs=[out], emit_generic=True)

	execution_engine = ExecutionEngine(transform(module))			execution_engine = ExecutionEngine(transform(module, matmul_boiler))

	# TODO: FFI-based solution to allow testing and printing with python code.			# TODO: FFI-based solution to allow testing and printing with python code.
	# Prepare arguments: one result f32.			# Prepare arguments: one result f32.
	# Arguments must be passed as pointers.			# Arguments must be passed as pointers.
	c_float_p = ctypes.c_float * 1			c_float_p = ctypes.c_float * 1
	res = c_float_p(-1.)			res = c_float_p(-1.)
	execution_engine.invoke("main", res)			execution_engine.invoke("main", res)

	log('RESULT: ', res[0])			log("RESULT: ", res[0])
	# CHECK: RESULT: 32.0			# CHECK: RESULT: 32.0

	test_generic()
				test_matmul_generic()


				def test_fill_builtin():
				with Context() as ctx, Location.unknown():
				module = Module.create()
				f64 = F64Type.get()
				i32 = IntegerType.get_signless(32)
				with InsertionPoint(module.body):

				@builtin.FuncOp.from_py_func(MemRefType.get((4, 16), i32))
				def fill_on_buffers(out):
				linalg.fill_rng_2d(outs=[out])

				execution_engine = ExecutionEngine(transform(module, fill_boiler))

				# TODO: FFI-based solution to allow testing and printing with python code.
				# Prepare arguments: one result i32.
				# Arguments must be passed as pointers.
				c_int_p = ctypes.c_int * 1
				res = c_int_p(-1)
				execution_engine.invoke("main", res)

				log("RESULT: ", res[0])
				# CHECK: RESULT: -480


				test_fill_builtin()


				def test_fill_generic():
				with Context() as ctx, Location.unknown():
				module = Module.create()
				f64 = F64Type.get()
				i32 = IntegerType.get_signless(32)
				with InsertionPoint(module.body):

				@builtin.FuncOp.from_py_func(MemRefType.get((4, 16), i32))
				def fill_on_buffers(out):
				linalg.fill_rng_2d(outs=[out])

				execution_engine = ExecutionEngine(transform(module, fill_boiler))

				# TODO: FFI-based solution to allow testing and printing with python code.
				# Prepare arguments: one result i32.
				# Arguments must be passed as pointers.
				c_int_p = ctypes.c_int * 1
				res = c_int_p(-1)
				execution_engine.invoke("main", res)

				log("RESULT: ", res[0])
				# CHECK: RESULT: -480


				test_fill_generic()
				nicolasvasilacheUnsubmitted Done Reply Inline Actions nit: NL nicolasvasilache: nit: NL

mlir/tools/mlir-linalg-ods-gen/mlir-linalg-ods-yaml-gen.cpp

Show First 20 Lines • Show All 45 Lines • ▼ Show 20 Lines
namespace {		namespace {

struct LinalgYAMLContext {		struct LinalgYAMLContext {
MLIRContext *mlirContext;		MLIRContext *mlirContext;
};		};

struct LinalgOpMetadata {		struct LinalgOpMetadata {
std::string name;		std::string name;
std::string cppOpName;		std::string cppClassName;
Optional<std::string> doc;		Optional<std::string> doc;
SmallVector<std::string> implements;		SmallVector<std::string> implements;
};		};

struct SerializedAffineMap {		struct SerializedAffineMap {
AffineMapAttr affineMapAttr;		AffineMapAttr affineMapAttr;

AffineMap affineMap() { return affineMapAttr.getValue(); }		AffineMap affineMap() { return affineMapAttr.getValue(); }
Show All 34 Lines	struct ScalarSymbolicCast {
std::string typeVar;		std::string typeVar;
// NOTE: This must be of arity 1, but to break the self-referential cycle,		// NOTE: This must be of arity 1, but to break the self-referential cycle,
// we use a heap allocated vector.		// we use a heap allocated vector.
std::vector<ScalarExpression> operands;		std::vector<ScalarExpression> operands;
};		};

struct ScalarExpression {		struct ScalarExpression {
Optional<std::string> arg;		Optional<std::string> arg;
		Optional<std::string> constant;
		Optional<int64_t> index;
Optional<ScalarApply> apply;		Optional<ScalarApply> apply;
Optional<ScalarSymbolicCast> symbolicCast;		Optional<ScalarSymbolicCast> symbolicCast;
};		};

struct ScalarAssign {		struct ScalarAssign {
std::string arg;		std::string arg;
ScalarExpression value;		ScalarExpression value;
};		};
▲ Show 20 Lines • Show All 90 Lines • ▼ Show 20 Lines	struct ScalarEnumerationTraits<LinalgIteratorTypeDef> {
}		}
};		};

/// Metadata about the op (name, C++ name, and documentation).		/// Metadata about the op (name, C++ name, and documentation).
template <>		template <>
struct MappingTraits<LinalgOpMetadata> {		struct MappingTraits<LinalgOpMetadata> {
static void mapping(IO &io, LinalgOpMetadata &info) {		static void mapping(IO &io, LinalgOpMetadata &info) {
io.mapRequired("name", info.name);		io.mapRequired("name", info.name);
io.mapRequired("cpp_op_name", info.cppOpName);		io.mapRequired("cpp_class_name", info.cppClassName);
io.mapOptional("doc", info.doc);		io.mapOptional("doc", info.doc);
io.mapOptional("implements", info.implements);		io.mapOptional("implements", info.implements);
}		}
};		};

/// How the ops indexing maps are produced. Must be one of:		/// How the ops indexing maps are produced. Must be one of:
/// - static_indexing_maps: A static list of AffineMaps, possibly with		/// - static_indexing_maps: A static list of AffineMaps, possibly with
/// some symbols that bind to attributes of the op. Each indexing map must		/// some symbols that bind to attributes of the op. Each indexing map must
Show All 22 Lines
/// - `scalar_arg`: Name of an argument to the op.		/// - `scalar_arg`: Name of an argument to the op.
/// - `scalar_apply`: Result of evaluating a named function (see		/// - `scalar_apply`: Result of evaluating a named function (see
/// `ScalarApply`).		/// `ScalarApply`).
/// - `symbolic_cast`: Cast to a symbolic TypeVar bound elsewhere.		/// - `symbolic_cast`: Cast to a symbolic TypeVar bound elsewhere.
template <>		template <>
struct MappingTraits<ScalarExpression> {		struct MappingTraits<ScalarExpression> {
static void mapping(IO &io, ScalarExpression &info) {		static void mapping(IO &io, ScalarExpression &info) {
io.mapOptional("scalar_arg", info.arg);		io.mapOptional("scalar_arg", info.arg);
		io.mapOptional("scalar_const", info.constant);
		io.mapOptional("scalar_index", info.index);
io.mapOptional("scalar_apply", info.apply);		io.mapOptional("scalar_apply", info.apply);
io.mapOptional("symbolic_cast", info.symbolicCast);		io.mapOptional("symbolic_cast", info.symbolicCast);
}		}
};		};

/// A scalar expression that evaluates a named function.		/// A scalar expression that evaluates a named function.
/// Functions are generally "math" level and type polymorphic. Builtin		/// Functions are generally "math" level and type polymorphic. Builtin
/// functions include:		/// functions include:
▲ Show 20 Lines • Show All 107 Lines • ▼ Show 20 Lines
findTensorDefArgIndex(StringRef name, SmallVectorImpl<LinalgTensorDef> &args) {		findTensorDefArgIndex(StringRef name, SmallVectorImpl<LinalgTensorDef> &args) {
for (auto it : llvm::enumerate(args)) {		for (auto it : llvm::enumerate(args)) {
if (it.value().name == name)		if (it.value().name == name)
return it.index();		return it.index();
}		}
return None;		return None;
}		}

static Optional<int>		// Try to map the TypeVar to a predefined or an argument type.
findTypeVarArgIndex(StringRef typeVar, SmallVectorImpl<LinalgTensorDef> &args) {		static Optional<std::string>
		findTypeValue(StringRef typeVar, SmallVectorImpl<LinalgTensorDef> &args) {
		// Handle all predefined types.
		if (typeVar == "I32")
		return std::string("helper.getIntegerType(32)");
		if (typeVar == "I64")
		return std::string("helper.getIntegerType(64)");
		if (typeVar == "F32")
		return std::string("helper.getFloat32Type()");
		if (typeVar == "F64")
		return std::string("helper.getFloat64Type()");

		// Search all argument types.
for (auto it : llvm::enumerate(args)) {		for (auto it : llvm::enumerate(args)) {
if (it.value().elementTypeVar == typeVar)		if (it.value().elementTypeVar == typeVar)
return it.index();		return llvm::formatv("block.getArgument({0}).getType()", it.index())
		.str();
}		}

return None;		return None;
}		}

static ScalarAssign *		static ScalarAssign *
findAssignment(StringRef name, SmallVectorImpl<ScalarAssign> &assignments) {		findAssignment(StringRef name, SmallVectorImpl<ScalarAssign> &assignments) {
for (auto &assign : assignments) {		for (auto &assign : assignments) {
if (assign.arg == name)		if (assign.arg == name)
return &assign;		return &assign;
▲ Show 20 Lines • Show All 171 Lines • ▼ Show 20 Lines	if (opConfig.metadata->doc) {
StringRef summary, description;		StringRef summary, description;
std::tie(summary, description) =		std::tie(summary, description) =
StringRef(*opConfig.metadata->doc).trim().split('\n');		StringRef(*opConfig.metadata->doc).trim().split('\n');
doc = llvm::formatv(docFmt, summary.trim(), description.trim());		doc = llvm::formatv(docFmt, summary.trim(), description.trim());
}		}

interfaceNameList = interleaveToString(opConfig.metadata->implements, ", ");		interfaceNameList = interleaveToString(opConfig.metadata->implements, ", ");

os << llvm::formatv(structuredOpOdsHeaderFormat, opConfig.metadata->cppOpName,		os << llvm::formatv(
		structuredOpOdsHeaderFormat, opConfig.metadata->cppClassName,
opConfig.metadata->name, interfaceNameList, doc, attrList,		opConfig.metadata->name, interfaceNameList, doc, attrList,
opConfig.structuredOp->args.size(), attrBuilder,		opConfig.structuredOp->args.size(), attrBuilder, attrMethods);
attrMethods);

return success();		return success();
}		}

static LogicalResult		static LogicalResult
generateNamedGenericOpDefns(LinalgOpConfig &opConfig,		generateNamedGenericOpDefns(LinalgOpConfig &opConfig,
GenerationContext &genContext) {		GenerationContext &genContext) {
if (!genContext.shouldGenerateDefns())		if (!genContext.shouldGenerateDefns())
return success();		return success();

raw_ostream &os = genContext.defns();		raw_ostream &os = genContext.defns();
StringRef className = opConfig.metadata->cppOpName;		StringRef className = opConfig.metadata->cppClassName;

// Implementation banner.		// Implementation banner.
std::string bannerComment = llvm::formatv("Implementation of {0}", className);		std::string bannerComment = llvm::formatv("Implementation of {0}", className);
os << llvm::formatv(bannerFormat, bannerComment);		os << llvm::formatv(bannerFormat, bannerComment);

// Reference iterators.		// Reference iterators.
{		{
std::string iteratorsStr;		std::string iteratorsStr;
▲ Show 20 Lines • Show All 139 Lines • ▼ Show 20 Lines
)FMT";		)FMT";
os << llvm::formatv(structuredOpGetLibraryCallFormat, className);		os << llvm::formatv(structuredOpGetLibraryCallFormat, className);
}		}

// regionBuilder()		// regionBuilder()
{		{
// Generates a regionBuilder method. Parameters.		// Generates a regionBuilder method. Parameters.
// {0}: Class name		// {0}: Class name
// {1}: Statements		// {1}: Number of args
		// {2}: Statements
static const char structuredOpRegionBuilderFormat[] = R"FMT(		static const char structuredOpRegionBuilderFormat[] = R"FMT(
void {0}::regionBuilder(Block &block, ValueRange captures) {{		void {0}::regionBuilder(Block &block, ValueRange captures) {{
RegionBuilderHelper helper(block);		assert({1} > 0 && block.getNumArguments() == {1} &&
		"{0} regionBuilder expects {1} (>=0) args");
		RegionBuilderHelper helper(block.getArgument(0).getContext(), block);
SmallVector<Value> yields;		SmallVector<Value> yields;
{1}		{2}
helper.yieldOutputs(yields);		helper.yieldOutputs(yields);
}		}
)FMT";		)FMT";
auto &args = opConfig.structuredOp->args;		auto &args = opConfig.structuredOp->args;
auto &assignments = opConfig.structuredOp->assignments;		auto &assignments = opConfig.structuredOp->assignments;
size_t generatedAssignmentCount = 0;		size_t generatedAssignmentCount = 0;
int localCounter = 0;		int localCounter = 0;
SmallVector<std::string> stmts;		SmallVector<std::string> stmts;
Show All 13 Lines	for (LinalgTensorDef &arg : args) {
std::function<Optional<std::string>(ScalarExpression &)>		std::function<Optional<std::string>(ScalarExpression &)>
generateExpression =		generateExpression =
[&](ScalarExpression &expression) -> Optional<std::string> {		[&](ScalarExpression &expression) -> Optional<std::string> {
if (expression.arg) {		if (expression.arg) {
// Argument reference.		// Argument reference.
Optional<int> argIndex = findTensorDefArgIndex(*expression.arg, args);		Optional<int> argIndex = findTensorDefArgIndex(*expression.arg, args);
if (!argIndex) {		if (!argIndex) {
emitError(genContext.getLoc())		emitError(genContext.getLoc())
<< "scalar argument not defined on the op: " << arg.name;		<< "scalar argument not defined on the op: " << *expression.arg;
return None;		return None;
}		}
return std::string(		return std::string(
llvm::formatv("block.getArgument({0})", *argIndex));		llvm::formatv("block.getArgument({0})", *argIndex));
} else if (expression.apply) {		}
		if (expression.constant) {
		std::string cppIdent = llvm::formatv("value{0}", ++localCounter);
		stmts.push_back(
		llvm::formatv(R"FMT(Value {0} = helper.constant("{1}");)FMT",
		cppIdent, expression.constant));
		return cppIdent;
		}
		if (expression.index) {
		// Access an iteration index.
		std::string cppIdent = llvm::formatv("value{0}", ++localCounter);
		stmts.push_back(llvm::formatv("Value {0} = helper.index({1});",
		cppIdent, *expression.index));
		return cppIdent;
		}
		if (expression.apply) {
// Apply function.		// Apply function.
// Recursively generate operands.		// Recursively generate operands.
SmallVector<std::string> operandCppValues;		SmallVector<std::string> operandCppValues;
for (ScalarExpression &operand : expression.apply->operands) {		for (ScalarExpression &operand : expression.apply->operands) {
auto operandCppValue = generateExpression(operand);		auto operandCppValue = generateExpression(operand);
if (!operandCppValue)		if (!operandCppValue)
return None;		return None;
operandCppValues.push_back(*operandCppValue);		operandCppValues.push_back(*operandCppValue);
}		}
std::string cppIdent = llvm::formatv("value{0}", ++localCounter);		std::string cppIdent = llvm::formatv("value{0}", ++localCounter);
stmts.push_back(		stmts.push_back(
llvm::formatv("Value {0} = helper.applyfn__{1}({2});", cppIdent,		llvm::formatv("Value {0} = helper.applyfn__{1}({2});", cppIdent,
expression.apply->fnName,		expression.apply->fnName,
interleaveToString(operandCppValues, ", ")));		interleaveToString(operandCppValues, ", ")));
return cppIdent;		return cppIdent;
} else if (expression.symbolicCast) {		}
		if (expression.symbolicCast) {
// Symbolic cast.		// Symbolic cast.
// Operands must be arity 1.		// Operands must be arity 1.
if (expression.symbolicCast->operands.size() != 1) {		if (expression.symbolicCast->operands.size() != 1) {
emitError(genContext.getLoc())		emitError(genContext.getLoc())
<< "symbolic_cast operand arity must be 1";		<< "symbolic_cast operand arity must be 1";
return None;		return None;
}		}
Optional<std::string> operandCppValue =		Optional<std::string> operandCppValue =
generateExpression(expression.symbolicCast->operands[0]);		generateExpression(expression.symbolicCast->operands[0]);
if (!operandCppValue)		if (!operandCppValue)
return None;		return None;

// Try to map the TypeVar to an arg index (which map to block arg		Optional<std::string> typeCppValue =
// indices), since we can just get that type directly.		findTypeValue(expression.symbolicCast->typeVar, args);
// TODO: Handle free type variables which do not map to an argument.		if (!typeCppValue) {
Optional<int> typeArgIndex =
findTypeVarArgIndex(expression.symbolicCast->typeVar, args);
if (!typeArgIndex) {
emitError(genContext.getLoc())		emitError(genContext.getLoc())
<< "type variable " << expression.symbolicCast->typeVar		<< "type variable " << expression.symbolicCast->typeVar
<< ", used in a symbolic cast must map to an argument but it "		<< ", used in a symbolic cast must map to a predefined or "
<< "does not";		<< "an argument type but it does not";
return None;		return None;
}		}
std::string typeCppValue =
llvm::formatv("block.getArgument({0}).getType()", *typeArgIndex);
std::string cppIdent = llvm::formatv("value{0}", ++localCounter);		std::string cppIdent = llvm::formatv("value{0}", ++localCounter);
stmts.push_back(llvm::formatv("Value {0} = helper.cast({1}, {2});",		stmts.push_back(llvm::formatv("Value {0} = helper.cast({1}, {2});",
cppIdent, typeCppValue,		cppIdent, typeCppValue.getValue(),
*operandCppValue));		*operandCppValue));
return cppIdent;		return cppIdent;
} else {		}
emitError(genContext.getLoc()) << "unknown ScalarExpression type";		emitError(genContext.getLoc()) << "unknown ScalarExpression type";
return None;		return None;
}
};		};
Optional<std::string> cppValue = generateExpression(assignment->value);		Optional<std::string> cppValue = generateExpression(assignment->value);
if (!cppValue)		if (!cppValue)
return failure();		return failure();
stmts.push_back(llvm::formatv("yields.push_back({0});", cppValue));		stmts.push_back(llvm::formatv("yields.push_back({0});", cppValue));
}		}

if (generatedAssignmentCount != assignments.size())		if (generatedAssignmentCount != assignments.size())
return emitError(genContext.getLoc())		return emitError(genContext.getLoc())
<< "mismatched number of assignments vs output arguments";		<< "mismatched number of assignments vs output arguments";

os << llvm::formatv(structuredOpRegionBuilderFormat, className,		int64_t numOfArgs = args.size();
		os << llvm::formatv(structuredOpRegionBuilderFormat, className, numOfArgs,
interleaveToString(stmts, "\n "));		interleaveToString(stmts, "\n "));
}		}

// Canonicalizers and folders.		// Canonicalizers and folders.
os << llvm::formatv(structuredOpCanonicalizersAndFoldersFormat, className);		os << llvm::formatv(structuredOpCanonicalizersAndFoldersFormat, className);

return success();		return success();
}		}
▲ Show 20 Lines • Show All 83 Lines • ▼ Show 20 Lines	int main(int argc, char **argv) {
for (auto &opConfig : opConfigs) {		for (auto &opConfig : opConfigs) {
if (!opConfig.metadata) {		if (!opConfig.metadata) {
emitError(genContext.getLoc())		emitError(genContext.getLoc())
<< "missing operation metadata on subsequent op";		<< "missing operation metadata on subsequent op";
return 1;		return 1;
}		}

genContext.setLoc(NameLoc::get(		genContext.setLoc(NameLoc::get(
Identifier::get(opConfig.metadata->cppOpName, &mlirContext)));		Identifier::get(opConfig.metadata->cppClassName, &mlirContext)));
if (failed(generateOp(opConfig, genContext))) {		if (failed(generateOp(opConfig, genContext))) {
return 1;		return 1;
}		}
}		}

if (outputOdsDecl)		if (outputOdsDecl)
outputOdsDecl->keep();		outputOdsDecl->keep();
if (outputCppImpl)		if (outputCppImpl)
outputCppImpl->keep();		outputCppImpl->keep();

return 0;		return 0;
}		}

This is an archive of the discontinued LLVM Phabricator instance.

[mir][Python][linalg] Support OpDSL extensions in C++.
ClosedPublic

Details

Diff Detail

Event Timeline

Revision Contents

Diff 346153

mlir/include/mlir/Dialect/Linalg/IR/LinalgNamedStructuredOps.yaml

mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp

mlir/python/mlir/dialects/linalg/opdsl/lang/comprehension.py

mlir/python/mlir/dialects/linalg/opdsl/lang/config.py

mlir/python/mlir/dialects/linalg/opdsl/lang/emitter.py

mlir/python/mlir/dialects/linalg/opdsl/lang/scalar_expr.py

mlir/python/mlir/dialects/linalg/opdsl/ops/core_named_ops.py

mlir/test/CMakeLists.txt

mlir/test/Dialect/Linalg/generalize-named-polymorphic-ops.mlir

mlir/test/lit.cfg.py

mlir/test/mlir-linalg-ods-gen/test-linalg-ods-yaml-gen.yaml

mlir/test/python/dialects/linalg/opdsl/arguments.py

mlir/test/python/dialects/linalg/opdsl/assignments.py

mlir/test/python/dialects/linalg/opdsl/emit_structured_generic.py

mlir/test/python/dialects/linalg/opsrun.py

mlir/tools/mlir-linalg-ods-gen/mlir-linalg-ods-yaml-gen.cpp

This is an archive of the discontinued LLVM Phabricator instance.

[mir][Python][linalg] Support OpDSL extensions in C++.ClosedPublic

Details

Diff Detail

Event Timeline

Revision Contents

Diff 346153

mlir/include/mlir/Dialect/Linalg/IR/LinalgNamedStructuredOps.yaml

mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp

mlir/python/mlir/dialects/linalg/opdsl/lang/comprehension.py

mlir/python/mlir/dialects/linalg/opdsl/lang/config.py

mlir/python/mlir/dialects/linalg/opdsl/lang/emitter.py

mlir/python/mlir/dialects/linalg/opdsl/lang/scalar_expr.py

mlir/python/mlir/dialects/linalg/opdsl/ops/core_named_ops.py

mlir/test/CMakeLists.txt

mlir/test/Dialect/Linalg/generalize-named-polymorphic-ops.mlir

mlir/test/lit.cfg.py

mlir/test/mlir-linalg-ods-gen/test-linalg-ods-yaml-gen.yaml

mlir/test/python/dialects/linalg/opdsl/arguments.py

mlir/test/python/dialects/linalg/opdsl/assignments.py

mlir/test/python/dialects/linalg/opdsl/emit_structured_generic.py

mlir/test/python/dialects/linalg/opsrun.py

mlir/tools/mlir-linalg-ods-gen/mlir-linalg-ods-yaml-gen.cpp

[mir][Python][linalg] Support OpDSL extensions in C++.
ClosedPublic