Skip to content

Instantly share code, notes, and snippets.

@bjacob
Created March 3, 2021 15:40
Show Gist options
  • Save bjacob/eff68b783ea8ac915ef93ba6e225168f to your computer and use it in GitHub Desktop.
Save bjacob/eff68b783ea8ac915ef93ba6e225168f to your computer and use it in GitHub Desktop.
diff --git a/mlir/include/mlir/Dialect/Linalg/IR/LinalgNamedStructuredOps.yaml b/mlir/include/mlir/Dialect/Linalg/IR/LinalgNamedStructuredOps.yaml
index 5752af9bea9a..b56b6ffa9c86 100644
--- a/mlir/include/mlir/Dialect/Linalg/IR/LinalgNamedStructuredOps.yaml
+++ b/mlir/include/mlir/Dialect/Linalg/IR/LinalgNamedStructuredOps.yaml
@@ -305,4 +305,62 @@ structured_op: !LinalgStructuredOpConfig
operands:
- !ScalarExpression
scalar_arg: B
-
+--- !LinalgOpConfig
+metadata: !LinalgOpMetadata
+ name: mmt_kernel
+ cpp_op_name: MmtKernelOp
+ doc: |-
+ WRITE ME
+ implements:
+ - LinalgContractionOpInterface
+structured_op: !LinalgStructuredOpConfig
+ args:
+ - !<LinalgTensorDef>
+ name: A
+ usage: input
+ shape: affine_map<()[s0, s1, s2] -> (s0, s2)>
+ element_type_var: T1
+ - !<LinalgTensorDef>
+ name: B
+ usage: input
+ shape: affine_map<()[s0, s1, s2] -> (s1, s2)>
+ element_type_var: T2
+ - !<LinalgTensorDef>
+ name: C
+ usage: output
+ shape: affine_map<()[s0, s1, s2] -> (s0, s1)>
+ element_type_var: U
+ indexing_maps: !LinalgIndexingMapsConfig
+ static_indexing_maps:
+ - affine_map<(d0, d1, d2)[s0, s1, s2] -> (d0, d2)>
+ - affine_map<(d0, d1, d2)[s0, s1, s2] -> (d1, d2)>
+ - affine_map<(d0, d1, d2)[s0, s1, s2] -> (d0, d1)>
+ iterator_types:
+ - parallel
+ - parallel
+ - reduction
+ assignments:
+ - !ScalarAssign
+ arg: C
+ value: !ScalarExpression
+ scalar_apply:
+ fn_name: add
+ operands:
+ - !ScalarExpression
+ scalar_arg: C
+ - !ScalarExpression
+ scalar_apply:
+ fn_name: mul
+ operands:
+ - !ScalarExpression
+ symbolic_cast:
+ type_var: U
+ operands:
+ - !ScalarExpression
+ scalar_arg: A
+ - !ScalarExpression
+ symbolic_cast:
+ type_var: U
+ operands:
+ - !ScalarExpression
+ scalar_arg: B
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment