fd226c9b02
* As discussed, fixes the ordering or (operands, results) -> (results, operands) in various `create` like methods. * Fixes a syntax error in an ODS accessor method. * Removes the linalg example in favor of a test case that exercises the same. * Fixes FuncOp visibility to properly use None instead of the empty string and defaults it to None. * Implements what was documented for requiring that trailing __init__ args `loc` and `ip` are keyword only. * Adds a check to `InsertionPoint.insert` so that if attempting to insert past the terminator, an exception is raised telling you what to do instead. Previously, this would crash downstream (i.e. when trying to print the resultant module). * Renames `_ods_build_default` -> `build_generic` and documents it. * Removes `result` from the list of prohibited words and for single-result ops, defaults to naming the result `result`, thereby matching expectations and what is already implemented on the base class. * This was intended to be a relatively small set of changes to be inlined with the broader support for ODS generating the most specific builder, but it spidered out once actually testing various combinations, so rolling up separately. Differential Revision: https://reviews.llvm.org/D95320
58 lines
1.9 KiB
Python
58 lines
1.9 KiB
Python
# RUN: %PYTHON %s | FileCheck %s
|
|
|
|
from mlir.ir import *
|
|
from mlir.dialects import builtin
|
|
from mlir.dialects import linalg
|
|
from mlir.dialects import std
|
|
|
|
|
|
def run(f):
|
|
print("\nTEST:", f.__name__)
|
|
f()
|
|
|
|
|
|
# CHECK-LABEL: TEST: testStructuredOpOnTensors
|
|
def testStructuredOpOnTensors():
|
|
with Context() as ctx, Location.unknown():
|
|
module = Module.create()
|
|
f32 = F32Type.get()
|
|
tensor_type = RankedTensorType.get((2, 3, 4), f32)
|
|
with InsertionPoint.at_block_terminator(module.body):
|
|
func = builtin.FuncOp(name="matmul_test",
|
|
type=FunctionType.get(
|
|
inputs=[tensor_type, tensor_type],
|
|
results=[tensor_type]))
|
|
with InsertionPoint(func.add_entry_block()):
|
|
lhs, rhs = func.entry_block.arguments
|
|
result = linalg.MatmulOp([lhs, rhs], results=[tensor_type]).result
|
|
std.ReturnOp([result])
|
|
|
|
# CHECK: %[[R:.*]] = linalg.matmul ins(%arg0, %arg1 : tensor<2x3x4xf32>, tensor<2x3x4xf32>) -> tensor<2x3x4xf32>
|
|
print(module)
|
|
|
|
|
|
run(testStructuredOpOnTensors)
|
|
|
|
|
|
# CHECK-LABEL: TEST: testStructuredOpOnBuffers
|
|
def testStructuredOpOnBuffers():
|
|
with Context() as ctx, Location.unknown():
|
|
module = Module.create()
|
|
f32 = F32Type.get()
|
|
memref_type = MemRefType.get((2, 3, 4), f32)
|
|
with InsertionPoint.at_block_terminator(module.body):
|
|
func = builtin.FuncOp(name="matmul_test",
|
|
type=FunctionType.get(
|
|
inputs=[memref_type, memref_type, memref_type],
|
|
results=[]))
|
|
with InsertionPoint(func.add_entry_block()):
|
|
lhs, rhs, result = func.entry_block.arguments
|
|
linalg.MatmulOp([lhs, rhs], outputs=[result])
|
|
std.ReturnOp([])
|
|
|
|
# CHECK: linalg.matmul ins(%arg0, %arg1 : memref<2x3x4xf32>, memref<2x3x4xf32>) outs(%arg2 : memref<2x3x4xf32>)
|
|
print(module)
|
|
|
|
|
|
run(testStructuredOpOnBuffers)
|