[mlir] Fix wrong variable name in Linalg OpDSL

[ftynse]

Summary

The name seems to have been left over from a renaming effort on an unexercised
codepaths that are difficult to catch in Python. Fix it and add a test that
exercises the codepath.

[gysit]

Thanks for fixing!

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

 
 def linalg_structured_op(dsl_func=None,
                          *,
                          op_name=None,
                          op_class_name=None) -> DefinedOpCallable:
   if dsl_func is None:
     # Curry the keyword args in for delayed application.
     return functools.partial(
-        tc_def_op, op_name=op_name, op_class_name=op_class_name)
+        linalg_structured_op, op_name=op_name, op_class_name=op_class_name)
   # Determine default names by introspecting the function.
   if op_name is None:
     op_name = dsl_func.__name__
   if op_class_name is None:
     # Camel case it.
     op_class_name = f"{''.join(x.title() for x in op_name.split('_'))}Op"
 
   tc_model = LinalgOpDef(
       name=op_name, cpp_class_name=op_class_name, doc=inspect.getdoc(dsl_func))
 
   # Extract arguments and TensorDefs from the signature.
   dsl_func_args = list()
   sig = inspect.signature(dsl_func)
   for param_name, param in sig.parameters.items():
     param_default = param.default
     if isinstance(param_default, (TensorDef, ScalarDef, AttributeDef)):
       tc_model.add_operand(param_name, param_default.operand_def)
     else:
-      raise ValueError(f"@tc_def_op function parameters must be defaulted as "
-                       f"TensorDef(...), ScalarDef(...), or AttributeDef(...): "
-                       f"Found {param_name}: {param_default}")
+      raise ValueError(
+          f"@linalg_structured_op function parameters must be defaulted as "
     dsl_func_args.append(param_default)
 
   # Invoke the DSL func to finish populating the model.
   with bind_op_def(tc_model):
     dsl_func(*dsl_func_args)
 
   # TODO: The returned callable should be an IR emitter but that is not
   # upstreamed yet.

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

 
 @linalg_structured_op
 def soft_plus_poly(
     I=TensorDef(T, S.M, S.N), O=TensorDef(U, S.M, S.N, output=True)):
   O[D.m, D.n] = \
       PrimFn.log(cast(U, const(1.0)) + cast(U, PrimFn.exp(I[D.m, D.n])))
 
 
+@linalg_structured_op(op_name="custom_op_name")
+def non_default_op_name(I=TensorDef(T, S.N), O=TensorDef(T, S.N, output=True)):
+  O[D.n] = I[D.n]
+
+
 with Context() as ctx, Location.unknown():
   module = Module.create()
   f16 = F16Type.get()
   f32 = F32Type.get()
   f64 = F64Type.get()
   i8 = IntegerType.get_signless(8)
   i16 = IntegerType.get_signless(16)
   i32 = IntegerType.get_signless(32)
with InsertionPoint(module.body):
     # CHECK-NEXT:   %[[LOG:.+]] = math.log %[[SUM]] : f32
     # CHECK-NEXT:   linalg.yield %[[LOG]] : f32
     # CHECK-NEXT: -> tensor<4x16xf32>
     @builtin.FuncOp.from_py_func(
         RankedTensorType.get((4, 16), f32), RankedTensorType.get((4, 16), f32))
     def test_f32_soft_plus(input, init_result):
       return soft_plus_poly(input, outs=[init_result])
 
+    # Just check that we don't assert out on name mismatch.
+    # CHECK-LABEL: @test_non_default_op_name
+    @builtin.FuncOp.from_py_func(
+        RankedTensorType.get((42,), f32), RankedTensorType.get((42,), f32))
+    def test_non_default_op_name(input, init_result):
+      return non_default_op_name(input, outs=[init_result])
+
 
 print(module)